open Unix

(** Excerpts from liblos to avoid a dependency *)

module Stream2 = struct

  (** Return a list of the stream elements *)
  let to_list strm =    
    let result = ref [] in
    Stream.iter (fun x -> result := x :: !result) strm;
    List.rev !result

  let map f s =
    Stream.from (fun n ->
		   try Some (f (Stream.next s))
		   with Stream.Failure -> None)

end

module Linestream = struct

  let to_stringstream strm =
    Stream2.map (fun line -> line ^ "\n") strm

  let to_file path strm =
    let chan = open_out_gen [Open_creat; Open_trunc; Open_wronly] 0o0644 path in
    Stream.iter (output_string chan) (to_stringstream strm);
    close_out chan

  (** Return a stream of the lines of a file *)
  let of_file path =
    let chan = open_in path in
    Stream.from (fun _ ->
		   try Some (input_line chan)
		   with
		     End_of_file -> close_in chan; None
		   | exn -> close_in chan; raise exn)

end

module Textfile = struct

  (** Return a stream of lines *)
  let to_stream path =
    let chan = open_in path in
    Stream.from (fun _ ->
		   try Some (input_line chan)
		   with
		     End_of_file -> close_in chan; None
		   | exn -> close_in chan; raise exn)

  (** Return a list of a lines *)
  let to_list path = Stream2.to_list (to_stream path)

end

module File = struct

  (* This needs to work on filesystems like /proc where st_size is invalid. *)
  let to_string path =
    let chan = open_in path in
    let rec loop lst =
      let buf = String.create 1024 in
      match input chan buf 0 1024 with
	1024 -> loop (buf :: lst)
      | 0 -> close_in chan; String.concat "" (List.rev lst)
      | n when n > 0 -> loop (String.sub buf 0 n :: lst)
      | _ -> close_in chan; raise (Failure "File.to_string") in
    loop []

end

(** concat componants of a path together ensuring there is exactly one 
    / between directories and no trailing /
*)

(** [realpath pathname] return the canonicalized absolute pathname *)
external realpath : string -> string = "realpath_native"

(** [fchdir file_descr] change the current directory to that specified by the open [file_descr] *)
external fchdir : file_descr -> unit = "fchdir_native";;

external pipe_buf : unit -> int = "pipe_buf_native"

external openlog : string -> int -> int -> unit = "openlog_native"
external syslog : int -> string -> unit = "syslog_native"
external closelog : unit -> unit = "closelog_native"

let ( ^/^ ) d1 d2 = 
  begin
    if d1 <> "" then
	(Pcre.qreplace ~pat:"/+$" ~templ:"" d1)  ^ "/" 
    else
      ""
  end
  ^ (Pcre.qreplace ~pat:"(^/+|/+$)" ~templ:"" d2)

(** [regexp_of_glob glob] convert the glob into an equivalent regular expersion pattern
@param glob - shell file glob as a string
@returns - equivalent regular expression as a string
*)
let regexp_of_glob glob =
  let escape str = 
    match str with
	"*" -> ".*"
      | "?" -> "."
      | "[" -> "["
      | "]" -> "]"
      | w -> Pcre.quote w
  in

  let strings_of_string str = Pcre.split ~pat:"" str in
    String.concat "" (List.map escape (strings_of_string glob))

(** [date_str ()] A nicely formatted current date/time string that sorts easily *)
let date_str () =
  let ct = localtime (time ()) in
    (Printf.sprintf "%4d-%02d-%02d %02d:%02d:%02d" (ct.tm_year+1900) ct.tm_mon ct.tm_mday ct.tm_hour ct.tm_min ct.tm_sec)

let verbose = List.length (List.filter (fun x -> x = "-v") (Array.to_list Sys.argv))

(** print a string, prepended with the current time *)
let message str =
  if verbose > 0 then prerr_endline ((date_str ()) ^ " " ^ str)

(** run a system command, raise an error if the command does not return 0 *)
let command str =
  let _ = message ("-> " ^ str) in
  let result = Sys.command str in
    if result != 0 then raise (Failure (str ^ "-> " ^ (string_of_int result)))
      
(** return entire contents of [ichan] as a [string list]. *)
let input_lines ichan =
  let result = ref [] in
    try
      while true do result := (input_line ichan) :: !result done; []
    with
	End_of_file -> List.rev !result

(** helper functions for [get_process_output] *)
let new_get_process_output cmd =
  let _ = message (" -> " ^ cmd) in
  let ichan = Unix.open_process_in cmd in
  let lines = input_lines ichan in
    Unix.close_process_in ichan, lines
      
let rec old_get_process_output cmd =
  let (status, lines) = new_get_process_output cmd in lines

(** [get_process_output] runs [cmd] and returns its output as a [string list] *)
let rec get_process_output cmd =
  match new_get_process_output cmd with
      Unix.WEXITED 0, lines -> lines
    | Unix.WEXITED n, _ ->
	raise (Failure (cmd ^ " -> " ^ (string_of_int n)))
    | Unix.WSIGNALED n, _ ->
	raise (Failure (cmd ^ " -> killed " ^ (string_of_int n)))
    | Unix.WSTOPPED n, _ ->
	raise (Failure (cmd ^ " -> stopped " ^ (string_of_int n)))

(* Check that the file exists and is readable by the current user *)
let exists filename =
  try
    let _ = Unix.access filename [Unix.F_OK ; Unix.R_OK] in
      true
  with
      Unix.Unix_error (Unix.ENOENT, "access", _) -> false
	
let sh_mount ?(opts=[]) device mntpnt : unit =
  command ("/bin/mount " ^ (String.concat " " opts) ^ " '" ^ device ^ "' '" ^ mntpnt ^ "'")
    
let sh_umount ?(opts=[]) device_or_mntpnt : unit =
  command ("/bin/umount " ^ (String.concat " " opts) ^ " '" ^ device_or_mntpnt ^ "'")

type mountflag =
    MS_RDONLY		(* Mount read-only *)
  | MS_NOSUID		(* Ignore suid and sgid bits *)
  | MS_NODEV		(* Disallow access to device special files *)
  | MS_NOEXEC		(* Disallow program execution *)
  | MS_SYNCHRONOUS	(* Writes are synced at once *)
  | MS_REMOUNT		(* Alter flags of a mounted FS *)
  | MS_MANDLOCK		(* Allow mandatory locks on an FS *)
  | MS_DIRSYNC		(* Directory modifications are synchronous *)
  | MS_NOATIME		(* Do not update access times. *)
  | MS_NODIRATIME	(* Do not update directory access times *)
  | MS_BIND
  | MS_MOVE
  | MS_REC
  | MS_VERBOSE
  | MS_POSIXACL		(* VFS does not apply the umask *)
  | MS_ONE_SECOND	(* fs has 1 sec a/m/ctime resolution *)
  | MS_ACTIVE
  | MS_NOUSER

type umount2flag =
    MNT_FORCE
  | MNT_DETACH

external mount:
  string		(* source *)
  -> string		(* target *)
  -> string		(* e.g. reiserfs, proc *)
  -> mountflag list
  -> string		(* usually a comma separated list of options *)
  -> unit		(* 0 for success, -1 for error *)
  = "mount_native"

external umount
  : string		(* target *)
  -> unit		(* 0 for success, -1 for error *)
  = "umount_native"
external umount2
  : string		(* target *)
  -> umount2flag list
  -> unit
  = "umount2_native"

type file_kind =
    S_REG
  | S_DIR
  | S_CHR
  | S_BLK
  | S_LNK
  | S_FIFO
  | S_SOCK

external mknod
  : string		(* pathname *)
  -> file_kind		(* file_kind, e.g., S_CHR = character special *)
  -> int		(* mode *)
  -> int		(* major device number *)
  -> int		(* minor device number *)
  -> unit
  = "mknod_native"

external pivot_root
  : string              (* new root *)
  -> string             (* put_old *)
  -> unit
  = "pivot_root_native"

external init_module
  : string		(* module name *)
  -> string		(* path name of module *)
  -> unit
  = "init_module_native"

external uname_r
  : unit
  -> string		(* OS release *)
  = "uname_r_native"
    
let mkdir ?(opts=[]) directory : unit =
  command ("/bin/mkdir " ^ (String.concat " " opts) ^ " '" ^ directory ^ "'")

let rmdir ?(opts=[]) directory : unit =
  command ("/bin/rmdir " ^ (String.concat " " opts) ^ " '" ^ directory ^ "'")
    
let grep ?(opts=[]) pattern file : string list =
  get_process_output ("/bin/grep " ^ (String.concat " " opts) ^ " " ^ pattern ^ " " ^ file)

let cp ?(opts=["-p"]) src dst : unit = 
  command ("cp " ^ (String.concat " " opts) ^ " '" ^ src ^ "' '" ^ dst ^ "'")

let mv ?(opts=[]) src dst : unit = 
  command ("mv " ^ (String.concat " " opts) ^ " '" ^ src ^ "' '" ^ dst ^ "'")

let rm ?(opts=[""]) dst : unit = 
  command ("rm " ^ (String.concat " " opts) ^ " '" ^ dst ^ "'")

let touch ?(opts=[""]) dst : unit = 
  command ("touch " ^ (String.concat " " opts) ^ " '" ^ dst ^ "'")

(** [umount_if_mounted] if [mntpnt] currently has something mounted on it, umount whateven is mounted *)
let umount_if_mounted mntpnt =
  if (try 
	let _ = grep mntpnt "/etc/mtab" in
	  true
      with
	  Failure _ -> false) 
  then
    umount mntpnt

(** [mount_once] mounts [device] on [mntpnt]. If [mntpnt] does not exist it is created. An umount will be performed is something is mounted on [mntpnt] already. *)      
let mount_once ?(opts=[]) device mntpnt = 
  if (exists mntpnt) then
    umount_if_mounted mntpnt
  else
    mkdir ~opts:["-p"] mntpnt ;
  
  sh_mount ~opts:opts device mntpnt

(** [map_over_file] [func] [filename] applies [func] (which has type string -> string) to every line in [file] and write output back to [file]. *)
let map_over_file func file =
  Linestream.to_file file
    (Stream.of_list (List.map func (Stream2.to_list
				      (Linestream.of_file file))))

(** [kill_by_cwd path] kills all the processes whose cwd is [path] *)
let kill_by_cwd path =
  let rec map_with_exceptions func lyst =
    match lyst with
	(head::tail) -> 
	  begin
	    try
	      (func head)::(map_with_exceptions func tail)
	    with
		_ -> map_with_exceptions func tail
	  end
      | [] -> []
  in
  let path = List.hd (get_process_output ("readlink -f " ^ path)) in
  let pids = get_process_output ("cd /proc ; echo [0-9]*") in
  let cwds = map_with_exceptions (fun pid -> (pid,Unix.readlink ("/proc" ^/^ pid ^/^ "cwd"))) pids in
  let kill_list = List.filter (fun (_,cwd) -> path = cwd) cwds in
    List.iter (fun (pid,_) -> command ("ps uw " ^ pid ^ " ; kill " ^ pid)) kill_list

(** [umount_below path] umount any directories that are mount below [path] *)
let umount_below path =
  let unescape str =
    let stream = Stream.of_string str in
    let escape_char a b c =
      match (a, b, c) with
	  ('0','4','0') -> " "
	| ('0','1','1') -> "\t"
	| ('0','1','2') -> "\n"
	| ('1','3','4') -> "\\"
	| _ -> 
	    let str = (String.create 4) in
	      String.set str 0 '\\' ;
	      String.set str 1 a ;
	      String.set str 2 b ;
	      String.set str 3 c ;
	      str
    in
    let rec unescape' =
      parser
	  [< ''\\' ; 'a ; 'b ; 'c ; rest = unescape' >] -> (escape_char a b c) ^ rest
	| [< 'c ; rest = unescape' >] -> (String.make 1 c) ^ rest
	| [< >] -> ""
    in
      unescape' stream
  in
  let proc_mount = List.map (Pcre.split ~pat:" ") (Textfile.to_list "/proc/mounts") in
  let mount_point mount_line =
    match mount_line with
	(_::mp::_) -> (unescape mp)
      | o -> raise Not_found
  in
  let mount_points = List.map mount_point proc_mount in
  let umount_regexp = Pcre.regexp ("^" ^ (Pcre.quote (realpath path))) in
  let needs_umounting =
    List.filter (fun mount -> Pcre.pmatch ~rex:umount_regexp mount) mount_points
  in
    List.iter (sh_umount ~opts:["-l"]) needs_umounting ;
    needs_umounting
;;

type ('a, 'b) either = Left of 'a
		       | Right of 'b

(** [fchroot path f] runs function [f] in a chroot at [path] *)
let fchroot path f =
  let orig_wd = Unix.getcwd () in
  let root_fd = Unix.openfile "/" [Unix.O_RDONLY] 0o640 in
    Unix.chroot path ;
    Unix.chdir "/" ;
    let ret_val =
      try 
	Right (f ())
      with
	  exn -> Left exn
    in
      fchdir root_fd ;
      Unix.close root_fd ;
      Unix.chroot "." ;
      Unix.chdir orig_wd ;
      match ret_val with
	  Right v ->  v
	| Left exn -> raise exn

(** [fork_and_chroot path f] forks process, chroots to [path] and executes
  function [f]. The parent process will wait for the child process to
  exit before continuing 
  
  The function f should return a status (integer).  Fchroot will return this value.
  If f throws an exception, fchroot will return -1.
*)
let fork_and_chroot path f =
  match Unix.fork () with
    0 -> Unix.chroot path ; let result = try f () with _ -> -1 in exit result
  | child_id -> let (_, status) = Unix.waitpid [] child_id in status

external lchown: string -> int -> int -> unit = "lchown_native"
external seteuid: int -> unit = "seteuid_native"
external setresuid: int -> int -> int -> unit = "setresuid_native"

(* Evaluate a thunk as the effective user.  This will convince things
   like mount(2) that you really are root if you are executing suid. *)
let su ?(newuid=Unix.geteuid ()) f =
  let ruid = Unix.getuid ()
  and euid = Unix.geteuid () in
  setresuid newuid euid (-1);
  let result = f () in
  setresuid ruid euid (-1);
  result

let do_locked ?(wait=0.01) f path =
  let pid = string_of_int (Unix.getpid ()) ^ "\n" in
  let len = String.length pid in
  let rec loop () =
    try
      let fd = Unix.openfile path [Unix.O_CREAT; Unix.O_EXCL] 0o644 in
      if Unix.write fd pid 0 len != len
      then failwith "do_locked: can't write lock file";
      Unix.close fd;
      begin
	try f ()
	with exn ->
	  Unix.unlink path;
	  raise exn
      end;
      Unix.unlink path
    with
      Unix.Unix_error (Unix.EEXIST, _, _) ->
	let pid = File.to_string path in
	let pid = String.sub pid 0 (String.length pid - 1) in
	if Sys.file_exists ("/proc/" ^ pid)
	then ignore (Unix.select [] [] [] wait)
	else Unix.unlink path;
	loop ()
  in loop ()