open Unix

exception No_home_variable
exception No_home_to_chdir
exception Fails_daemonize
exception Already_running
exception No_pidfile
exception Cant_find_logger
exception Problem_with_logger
exception Invalid_pidfile
exception No_proc_entry of string

type daemon_state = Terminate | Respawn | MlnetDied | Nop

type action_type = Start | Stop | Restart 

let (prio_min, prio_max) = -20, 20

type mldonkey_server_state = {
  real_home : string option;
  daemon : bool;
  quiet : bool;
  logfile : string option;
  pidfile : string option;
  program : string;
  logger : string;
  prg_args : string list;
  chdir : string option;
  chuid : int;
  chgid : int;
  umask : int;
  nice  : int;
  action : action_type;

  (* Variable to determine the way we manage mlnet *)
  end_duration : float;
  max_alive : float;
  min_alive : float;
  respawn_after : float;

  (* Filename that should be inspected. *)
  (* I put them here, because this should 
   become a command line option -- if needed,
   *)
  (* All path should be relative to chdir, where
   all the .ini files should be.
   *)

  inifiles : string list;
  passwdfiles : string list;
}

let user_home st =
  match st.real_home with 
    | None ->
        raise No_home_variable
    | Some x ->
        x

let get_chdir st =
  match st.chdir with
    | None ->
        let real_filename = Filename.concat (user_home st) ".mldonkey"
        in
          real_filename
    | Some x ->
        x

let get_pidfile st =
  match st.pidfile with
    | None ->
        let real_filename = Filename.concat (get_chdir st) "mldonkey.pid"
        in
          real_filename
    | Some x ->
        x

let get_logfile st =
  match st.logfile with
    | None ->
        let real_filename = Filename.concat (get_chdir st) "mldonkey.log"
        in
          real_filename
    | Some x ->
        x

let get_inifiles st =
  List.map (Filename.concat (get_chdir st)) st.inifiles

let get_passwdfiles st =
  List.map (Filename.concat (get_chdir st)) st.passwdfiles

let default_mldonkey_server_state = 
  {
    real_home = 
      (
        try 
          Some (getenv "HOME") 
        with Not_found -> 
          None
      );
    daemon = false;
    quiet = false;
    logfile = None;
    pidfile = None;
    program = "/usr/bin/mlnet";
    logger = "/usr/bin/logger";
    prg_args = [];
    chdir = None;
    chuid = Unix.getuid ();
    chgid = Unix.getgid ();
    umask = 0o0022;
    nice  = 0;
    action = Start;

    (* 5 min *)
    end_duration = 300.;

    (* 15 min *)
    min_alive = 900.;

    (* 1 day *)
    max_alive = 88400.;

    (* 10 min *)
    respawn_after = 600.;

    inifiles = 
      [
        "bittorrent.ini";
        "donkey.ini";
        "files.ini";
        "file_sources.ini";
        "fileTP.ini";
        "friends.ini";
        "gnutella2.ini";
        "gnutella.ini";
        "searches.ini";
        "servers.ini";
        "shared_files_new.ini";
        "stats.ini";
        "stats_mod.ini";
        "downloads.ini";
      ];

    passwdfiles =
      [
        "users.ini"
      ];
  }

let get_mldonkey_server_state () =
  let state = 
    ref default_mldonkey_server_state
  in
  let add_args x = 
    state := { !state with prg_args = x :: !state.prg_args }
  in
  let _ = Arg.parse [
    "--daemon", 
    Arg.Unit (fun x -> state := {!state with daemon = true}), 
    "Run in daemon mode";
    "--quiet", 
    Arg.Unit (fun x -> state := {!state with quiet = true}), 
    "Keep quiet";
    "--logfile", 
    Arg.String (fun x -> state := {!state with logfile = Some x}), 
    "Where to put the log";
    "--pidfile", 
    Arg.String (fun x -> state := {!state with pidfile = Some x}), 
    "Where to put the pid";
    "--program", 
    Arg.String (fun x -> state := {!state with program = x}), 
    "Which program to start";
    "--logger", 
    Arg.String (fun x -> state := {!state with logger = x}),
    "Logger program";
    "--chdir", 
    Arg.String (fun x -> state := {!state with chdir = Some x}), 
    "Where to chdir";
    "--chuid", 
    Arg.String (fun x -> 
                  let passwd_ent =
                    try
                      Unix.getpwnam x
                    with
                      | Not_found ->
                          raise (Arg.Bad ("wrong argument `"^x^"'; option `--chuid' expects a user name"))
                  in
                    state := {!state with chuid = passwd_ent.pw_uid}
    ), 
    "Which user own the process";
    "--chgid",
    Arg.String (fun x ->
                  let group_ent =
                    try
                      Unix.getgrnam x
                    with
                      | Not_found ->
                          raise (Arg.Bad ("wrong argument `"^x^"'; option `--chgid' expects a group name"))
                  in
                    state := {!state with chgid = group_ent.gr_gid}
    ),
    "Which group own the process";
    "--umask", 
    Arg.String (fun x ->
                  let mask =
                    try
                      int_of_string ("0o"^x)
                    with
                      | Failure "int_of_string" -> raise (Arg.Bad ("wrong argument `"^x^"'; option `--umask' expects an octal umask"))
                  in
                    state := {!state with umask = mask}),
    "What umask to use";
    "--nice",
    Arg.Int (fun x -> state := {!state with nice = x}),
    "Niceness of the process";
    "--end-duration", 
    Arg.Int (fun x -> state := {!state with end_duration = float_of_int x }),
    "How much time does it take to end mlnet";
    "--max-alive", 
    Arg.Int (fun x -> state := { !state with max_alive = (float_of_int x) *. 3600.}),
    "For how long an instance of mlnet should run";
    "--min-alive", 
    Arg.Int (fun x -> state := {!state with min_alive = float_of_int x}),
    "Minimun time between respawning";
    "--respawn-after", 
    Arg.Int (fun x -> state := {!state with respawn_after = float_of_int x}),
    "When mlnet fails, how long to wait before restarting";
    "--start", 
    Arg.Unit (fun x -> state := {!state with action = Start}),
    "Start mldonkey_server";
    "--stop", 
    Arg.Unit (fun x -> state := {!state with action = Stop}),
    "Stop a running mldonkey_server (use the pidfile)";
    "--restart", 
    Arg.Unit (fun x -> state := {!state with action = Restart}),
    "Restart a running mldonkey_server (use the pifile, only respawn mlnet)";
    "--", 
    Arg.Rest (fun x -> add_args x),
    "MLnet arguments"
  ]
            add_args
            "Usage mldonkey_server [options] -- [mlnet options] where options are:"
  in
    !state

let print_log st str =
  if st.quiet then
    true
  else if st.daemon then
    (
      (
        match Unix.system (st.logger ^ " -t mldonkey_server \"" ^ str ^ "\"") with
          | WEXITED(0) ->
              ()
          | WEXITED(127) ->
              raise Cant_find_logger
          | _ ->
              raise Problem_with_logger
      );
      true
    )
  else
    false

let debug st str =
  if print_log st str then
    ()
  else
    (    
      print_string str;
      print_newline ()
    )

let warning st str = 
  if print_log st str then
    ()
  else
    (
      prerr_string str;
      prerr_newline ()
    )

let fatal st str =
  prerr_string (" " ^ str);
  prerr_newline ();
  ignore (print_log st str)

let go_home st =
  debug st ("Chdir to chdir dir: " ^ (get_chdir st));
  Unix.putenv "MLDONKEY_STRINGS" (Filename.concat (get_chdir st) "mlnet_strings");
  Sys.chdir (get_chdir st)

let create_home st =
  if Sys.file_exists (get_chdir st) then
    ()
  else
    if not st.daemon then
      let answer =    
        prerr_string ((get_chdir st)^" doesn't exists." ^ 
                      " Do you want to create it? (y/N)");
        flush(Pervasives.stderr);
        read_line ()
      in
        match answer with
          |  "y" ->
              Unix.mkdir (get_chdir st) 0o0755;
              debug st ("Creating home dir: " ^ (get_chdir st))
          | _ ->
              raise No_home_to_chdir
    else
      raise No_home_to_chdir

let set_uid_gid st =
  debug st (
    "Set uid/gid of the process ("^(string_of_int st.chuid)
    ^", "^(string_of_int st.chgid)^")");
  setgid st.chgid;
  setuid st.chuid

let set_umask st =
  debug st ("Set umask of the process: "^(string_of_int st.umask));
  ignore(umask st.umask)

let set_nice st =
  debug st ("Set niceness of the process: "^(string_of_int st.nice));
  (* We are in the bound of the normal niceness *)
  if prio_min <= st.nice && st.nice <= prio_max then
    let current_nice = Unix.nice 0
    in
      (* Only root can lower the niceness of a process *)
      if current_nice > st.nice && (Unix.getuid ()) <> 0 then
        warning st ("Only root can lower the niceness of a process ("
                    ^(string_of_int current_nice)^" > "^(string_of_int st.nice)^")")
      else
        ignore (Unix.nice (st.nice - current_nice))
      else
        warning st ("Niceness out of bound ("^(string_of_int st.nice)
                    ^"not in ["^(string_of_int prio_min)^"; "^(string_of_int prio_max)^"])")

(** Create a pidfile, holding the PID value of the process *)
let create_pidfile st =
  let pidfile = open_out (get_pidfile st)
  in
    debug st ("Writing PID ("^(string_of_int (Unix.getpid ()))^") to pidfile: "^(get_pidfile st));
    output_string pidfile (string_of_int (Unix.getpid ()));
    output_string pidfile "\n";
    close_out pidfile

(** Read a pidfile, return the PID value stored in it *)
let read_pidfile st =
  try
    let pidfile = open_in (get_pidfile st) in
    let pid_server = int_of_string (input_line pidfile) in
      debug st ("Reading PID ("^(string_of_int pid_server)^") from pidfile: "^(get_pidfile st));
      close_in pidfile;
      pid_server
  with 
    | Sys_error(_) ->
        raise No_pidfile
    | End_of_file 
    | Failure "int_of_string" ->
        raise Invalid_pidfile

(** Remove a pidfile *)
let close_pidfile st =
  try
    debug st ("Removing pidfile: "^(get_pidfile st));
    Sys.remove (get_pidfile st)
  with Sys_error(_) ->
    raise No_pidfile

(** Check that the given PID is a running instance of the program which we are
  in *)
let daemon_is_running st = 
  let prog_inode_of_pid pid =
    let proc_filename = 
      List.fold_left Filename.concat "/proc" [(string_of_int pid); "exe"]
    in
      if Sys.file_exists proc_filename then
        (* This condition is too hard: when upgrading you loose the inode number
         because the script is reinstalled
         *)
        (*(Unix.stat proc_filename).Unix.st_ino*)
        Unix.readlink proc_filename
      else
          raise (No_proc_entry proc_filename)
  in
    if Sys.file_exists (get_pidfile st) then
      (
        try 
          let prev_pid = read_pidfile st
          in
          let real_prog_inode = prog_inode_of_pid (Unix.getpid ())
          in
          let prev_prog_inode = prog_inode_of_pid prev_pid
          in
            real_prog_inode = prev_prog_inode 
        with 
          | Invalid_pidfile ->
              (
                warning st ("Invalid pidfile: "^(get_pidfile st));
                false
              )
          | No_proc_entry proc ->
              warning st ("Cannot open "^proc^" entry for the given pidfile: "
                          ^(get_pidfile st));
              false
      )
    else
      false

(** Remove stale pidfile 
  *)
let remove_stale_pidfile st =
  if not (daemon_is_running st) && Sys.file_exists (get_pidfile st) then
    (
      debug st ("Removing stale pidfile: "^(get_pidfile st));
      Sys.remove (get_pidfile st)
    )
  else
    ()

(** Get problematic INI file for mldonkey. Returns a list of all problematic
  files 
  *)
let check_tmp_ini_files lst_fl =
  let check_one_tmp_ini_file lst fl =
    let tmp_fl = 
      fl ^ ".tmp"
    in
      if Filename.check_suffix fl ".ini" && Sys.file_exists tmp_fl then
        tmp_fl :: lst
      else
        lst
  in
    List.fold_left check_one_tmp_ini_file [] lst_fl

(** Check that user/group and perms are correctly sets for any 
  directory/file that could be used by mldonkey. Returns a list of all
  problematic files/directories.
  *)
let check_file_owner_perms (uid,gid) lst_fl =
  let check_one_file_owner_perms lst fl =
    let match_perm perm perm_to_match =
      (* Check that the perm given are enough to match the perm_to_match,
       in other word, that you have all the bit of the perm_to_match
       in the perm 
       *)
      ( perm land perm_to_match ) = perm_to_match
    in
      (* First of all, does the file exist ? *)
      if Sys.file_exists fl then
        (* Get property of the file *)
        let stat = Unix.stat fl
        in
        let enough_right =
          (* We try to be sure that any program running with (uid,gid) has enough
           right to read/write the file considered *)
          (* The owner is the user *)
          ( stat.Unix.st_uid = uid && match_perm stat.Unix.st_perm 0o600 )
            (* The group owner match the group of the user *)
            || ( stat.Unix.st_gid = gid && match_perm stat.Unix.st_perm 0o060 )
            (* Rights given to "other" *)
            || ( match_perm stat.Unix.st_perm 0o006 ) 
        in
          if enough_right then
            lst
          else
            fl :: lst
          else
            (* The file doesn't exist, it will be created by mldonkey, if needed *)
            lst
  in
    List.fold_left check_one_file_owner_perms [] lst_fl

(** Check that the given list of file are only readable by the owner/group. It 
  is used to check that the file containing password are not readable by the
  group "other". Returns a list of problematic files
  *)
let check_file_security_perms lst_fl =
  let check_one_file_security_perms lst fl =
    if Sys.file_exists fl then
      let stat = Unix.stat fl
      in
        (* Does other have read access to this file ? *)
        if ( stat.Unix.st_perm land 0o004 ) <> 0 then
          fl :: lst
        else
          lst
        else
          (* The file doesn't exist : we don't have problem *)
          lst
  in
    List.fold_left check_one_file_security_perms [] lst_fl

let sanity_check st = 
  (* Checker function: apply the given function on the given list. If the
   result is not an empty message, display the given message and solution 
   to solve the problem and exit with the given exit code. 
   *)
  let check_fun check lst_fl message proposed_solution exit_code =
    let result = check lst_fl
    in
      if result <> [] then
        (
          fatal st (message^": "^(String.concat ", " result)^" -- "^proposed_solution);
          exit exit_code
        )
      else
        ()
  in
  let pidfile = get_pidfile st 
  in
  let passwdfiles = get_passwdfiles st
  in
  let inifiles = get_inifiles st
  in
  let basedirs =
    [
      get_chdir st;
      Filename.dirname pidfile
    ]
  in
  let old_inifiles =
    List.map ( fun x -> x ^ ".old" )
      inifiles
  in
    (* Test existence of a few dirs *)
    check_fun
      (List.filter (fun fl -> not (Sys.file_exists fl)))
      basedirs
      "directory[ies] doesn't exist"
      "create it first"
      1;
    check_fun 
      check_tmp_ini_files 
      (inifiles @ passwdfiles)
      "temporary file[s] left" 
      "delete it first" 
      1;
    check_fun 
      (check_file_owner_perms (st.chuid,st.chgid))
      (basedirs @ inifiles @ old_inifiles @ passwdfiles)
      ("file[s] not owned by user "^(string_of_int st.chuid)
       ^" or group "^(string_of_int st.chgid))
      "reown it first"
      1;
    check_fun
      check_file_security_perms 
      passwdfiles
      ("file[s] should be only readable by owner "^(string_of_int st.chuid)
       ^" or group "^(string_of_int st.chgid))
      "change permission"
      1;
    check_fun
      (List.filter Sys.file_exists)
      [Filename.concat (get_chdir st) "mlnet.pid"]
      "file[s] should no exist"
      "delete it first"
      1

let stop_or_die st pid =
  let timeout = ref false
  in
  let _ =
    Sys.set_signal Sys.sigalrm (Sys.Signal_handle
                                  ( fun x -> timeout := true ));
    ignore (Unix.alarm (int_of_float st.end_duration));
    debug st ("Waiting termination of process "^(string_of_int pid));
    try 
      Unix.kill pid Sys.sigterm;
      ignore (waitpid [] pid )
    with Unix.Unix_error(_, _, _) ->
      ()
  in
    if !timeout then
      begin
        debug st ("Process "^(string_of_int pid)^" not responding, taking measure: SIGKILL");
        try
          Unix.kill pid Sys.sigkill
        with Unix.Unix_error(_, _, _) ->
          ()
      end
    else
      debug st ("Process "^(string_of_int pid)^" terminated")

let daemonize st = 
  if st.daemon then
    (
      debug st ("Fork the process");
      if Unix.fork () = 0 then
        (
          debug st ("Reset the group leader");
          if Unix.setsid () = Unix.getpid () then
            (
              debug st ("Fork a second time the process");
              if Unix.fork () = 0 then
                (
                  debug st ("Close standard IO");
                  let fd = Unix.openfile "/dev/null" [ Unix.O_RDWR ] 0o0644 
                  in
                    List.iter (
                      fun fd_std -> 
                        Unix.close fd_std; 
                        Unix.dup2 fd fd_std
                    )
                      [Unix.stdin; Unix.stdout; Unix.stderr];
                    Unix.close fd;
                    debug st ("Process is running in the background");
                    ()
                )
              else
                exit 0
            )
          else
            raise Fails_daemonize 
        )
      else
        exit 0
    )
  else
    ()  

let start_mldonkey_server st =
  sanity_check st;
  if daemon_is_running st then
    raise Already_running
  else
    ();
  remove_stale_pidfile st;
  set_nice st;
  set_uid_gid st;
  set_umask st;
  create_home st;
  go_home st;
  daemonize st;
  create_pidfile st;
  let launch_mlnet st =
    let (logger_stderr, mlnet_stderr) =
      if st.daemon then
        Unix.pipe ()
      else
        (stdin, stderr)
    in
    let (logger_stdout, mlnet_stdout) =
      if st.daemon then
        Unix.pipe ()
      else
        (stdin, stdout)
    in
    let args = Array.of_list 
                 (st.program :: (List.rev st.prg_args))
    in
    let pid_mlnet = 
      debug st ("Launching MLnet process");
      create_process st.program args
        stdin
        mlnet_stdout
        mlnet_stderr
    in
    let pid_logger_stderr =
      if logger_stderr != stdin then
        begin
          debug st ("Launching MLnet stderr logger");
          create_process st.logger [| st.logger ; "-t"; "mlnet_error" |]
            logger_stderr
            stdout
            stderr
        end
      else
        0
    in
    let pid_logger_stdout =
      if logger_stdout != stdin then
        begin
          debug st ("Launching MLnet stdout logger");
          create_process st.logger [| st.logger ; "-t"; "mlnet" |]
            logger_stdout
            stdout
            stderr
        end
      else
        0
    in
      (
        [pid_mlnet; pid_logger_stderr; pid_logger_stdout],
        [logger_stderr; mlnet_stderr; logger_stdout; mlnet_stdout]
      )
  in
  let stop_mlnet st (pids, fds) =
    let str_pids = List.fold_left 
                     ( fun str x -> str^" "^(string_of_int x) ) 
                     "" pids
    in
      debug st ("Stopping processes PID ("^str_pids^" )") ;
      begin
        try
          let close_fds x =
            if x != stdout && x != stdin && x != stderr then
              Unix.close x
            else
              ()
          in
          let stop_pids x =
            if x != 0 then
              stop_or_die st x
            else
              ()
          in
            List.iter stop_pids pids;
            List.iter close_fds fds
        with Unix.Unix_error(_,_,_) ->
          ()
      end;
      debug st ("Process stopped PID ("^(str_pids)^" )")
  in
  let state = ref Nop
  in
  let terminate = ref false
  in
  let reload = ref false
  in
  let _ =
    Sys.set_signal Sys.sigint ( Sys.Signal_handle 
                                  ( fun x -> state := Terminate ));
    Sys.set_signal Sys.sigterm ( Sys.Signal_handle
                                   ( fun x -> state := Terminate ));
    Sys.set_signal Sys.sighup ( Sys.Signal_handle
                                  ( fun x -> state := Respawn ));
    Sys.set_signal Sys.sigchld ( Sys.Signal_handle
                                   ( fun x -> state := MlnetDied ));
    Sys.set_signal Sys.sigpipe ( Sys.Signal_handle
                                   ( fun x -> state := Terminate ))
  in
    while not !terminate do
      let mlnet = launch_mlnet st 
      in
      let last_respawn = Unix.time ()
      in
        Sys.set_signal Sys.sigalrm ( Sys.Signal_handle
                                       ( fun x -> state := Respawn ));
        ignore (Unix.alarm ( int_of_float st.max_alive ));
        reload := false;
        while not !reload && not !terminate do
          let _ = 
            try
              state := Nop;
              begin
                match Unix.wait () with
                  | x,WEXITED(y) ->
                      debug st ("Process PID ("^(string_of_int x)^
                                ") exit with return code "^(string_of_int y))
                  | x,WSIGNALED(y) ->
                      debug st ("Process PID ("^(string_of_int x)^
                                ") was killed by signal "^(string_of_int y))
                  | x,WSTOPPED(y) ->
                      debug st ("Process PID ("^(string_of_int x)^
                                ") was stopped by signal "^(string_of_int y))
              end;
              state := MlnetDied
            (* On peut etre interrompu par un signal extrieur *)
            with Unix.Unix_error(EINTR,_,_) ->
              ()
          in
            match !state with
              | Terminate ->
                  debug st ("Terminate process");
                  stop_mlnet st mlnet;
                  terminate := true
              | Respawn ->
                  debug st ("Respawn process");
                  stop_mlnet st mlnet;
                  reload := true
              | MlnetDied ->
                  if Unix.time () -. last_respawn < st.min_alive then
                    begin
                      debug st ("Process respawning too fast: only live "^
                                (string_of_float(Unix.time () -. last_respawn)));
                      stop_mlnet st mlnet;
                      terminate := true
                    end
                  else
                    begin
                      debug st ("Process died, respawning: live for "^
                                (string_of_float(Unix.time () -. last_respawn)));
                      stop_mlnet st mlnet;
                      reload := true
                    end
              | Nop ->
                  ()
        done;
    done;
    debug st "MLDonkey server end";
    close_pidfile st

let kill_mldonkey_server st signal signal_name=
  if daemon_is_running st then
    let pid_server = read_pidfile st
    in
      debug st ("Sending signal "^signal_name^" to process PID ("^(string_of_int pid_server)^")");
      if signal = Sys.sigterm then
        stop_or_die st pid_server
      else
        begin
          try
            Unix.kill pid_server signal
          with Unix.Unix_error (_, _, _) ->
            ()
        end
      else
        debug st ("Daemon is not running")

let stop_mldonkey_server st =
  kill_mldonkey_server st Sys.sigterm "SIGTERM"

let restart_mldonkey_server st =
  kill_mldonkey_server st Sys.sighup "SIGHUP"

let () = 
  try
    let state = get_mldonkey_server_state () in
      begin
        match state.action with
          | Start ->
              start_mldonkey_server state
          | Stop ->
              stop_mldonkey_server state
          | Restart ->
              restart_mldonkey_server state
      end;
      exit 0
  with
    | No_home_variable ->  
        prerr_string (" Could not guess $HOME environment variable: provide a --chdir or $HOME");
        prerr_newline ()
    | No_home_to_chdir ->
        prerr_string (" Home dir doesn't exist");
        prerr_newline ()
    | Fails_daemonize ->
        prerr_string (" Cannot daemonize process");
        prerr_newline ()
    | Already_running ->
        prerr_string (" Some others mldonkey_server are running (a pidfile exists)");
        prerr_newline ()
    | No_pidfile ->
        prerr_string (" No pidfile, maybe no mldonkey_server are running");
        prerr_newline ()
    | Invalid_pidfile ->
        prerr_string (" Invalid pidfile, maybe the pidfile is corrupted");
        prerr_newline ();
    | No_proc_entry str ->
        prerr_string (" Cannot find /proc entry for " ^ str);
        prerr_newline ();
    | Unix.Unix_error (error,_,_) ->
        prerr_string (" " ^ error_message error);
        prerr_newline ()

let () = exit 1