Source

screentime-monitor / client_common.ml

Full commit
open Core.Std
open Async.Std

module Config = struct
  type t = { server: string
           ; user: Username.t sexp_option
           } 
  with sexp
end

let load_config () =
  match Unix.getenv "HOME" with
  | None -> failwith "no $HOME defined.  Can't open config"
  | Some home ->
    Reader.load_sexp (home ^/ ".screentime-monitor") Config.t_of_sexp
    >>| function
    | Error err ->
      failwiths "Failed to load config file" err <:sexp_of<Error.t>>
    | Ok config ->
      config
;;

let get_username () =
  let uid = Unix.getuid () in
  Unix.Passwd.getbyuid uid
  >>| function
  | None -> failwith "Could not compute username"
  | Some pwd -> Username.of_string pwd.Unix.Passwd.name
;;

let username_spec = Command.Spec.Arg_type.create Username.of_string
let category_spec = Command.Spec.Arg_type.create Category.of_string

let on_term_signal =
  lazy (
    let stop_ivar = Ivar.create () in
    Signal.handle [Signal.term] ~f:(fun (_:Signal.t) ->
        Ivar.fill stop_ivar ());
    Ivar.read stop_ivar
  )

let shared_flags () =
  Command.Spec.(
    empty
    +> flag "-username" (optional username_spec)
      ~doc:"Username to act as"
  )

let days_flag () =
  Command.Spec.(
    flag "-days" (optional_with_default 4 int)
      ~doc:"number of days the report should go back"
  )

let edit_file sexpable ~editor ~tempfile =
  let rec loop () =
    Unix.system_exn (String.concat [editor;" ";tempfile])
    >>= fun () ->
    With_format.load tempfile sexpable
    >>= function
    | Ok resp -> return (Some resp)
    | Error e ->
      printf "Unable to read data:\n%s\n" (Error.to_string_hum e);
      printf "Try again? (Y/n): ";
      Reader.read_line (Lazy.force Reader.stdin)
      >>= fun response ->
      let reread =
        match response with
        | `Eof -> true
        | `Ok s ->
          match s |> String.lowercase |> String.strip with
          | "n" | "no" -> false
          | _ -> true
      in
      if not reread then (printf "Abandoning edit\n"; return None)
      else loop ()
  in
  loop ()
;;

let shell_run cmd args =
  In_thread.run(fun () ->
      Core_extended.Shell.run_full cmd args)

let clear_string =
  lazy (shell_run "clear" [])

let clear_screen () =
  force clear_string
  >>| fun clear_string ->
  print_endline clear_string


let with_rpc_conn f ~host ~port =
  Tcp.with_connection
    (Tcp.to_host_and_port host port)
    ~timeout:(sec 1.)
    (fun _ r w ->
       Rpc.Connection.create r w ~connection_state:()
       >>= function
         | Error exn -> raise exn
         | Ok conn -> f conn
    )

let resolve_username username config =
  match username with
  | Some x -> return x
  | None ->
    match config.Config.user with
    | Some x -> return x
    | None -> get_username ()


let setup_conn username k =
  load_config ()
  >>= fun config ->
  resolve_username username config
  >>= fun username ->
  with_rpc_conn
    ~host:config.Config.server
    ~port:Common.port
    (fun conn -> k username conn)

module Robust_connection : sig
  type t
  val get : t -> Rpc.Connection.t option

  val with_connect
    :  reconnect_after:Time.Span.t
    -> 'addr Tcp.where_to_connect
    -> (t -> 'a Deferred.t)
    -> 'a Deferred.t
    
end = struct
  type t = Rpc.Connection.t option ref

  let get t = !t

  let create ~reconnect_after ~stop where_to_connect =
    let t = ref None in
    let rec connect_loop () =
      if Deferred.is_determined stop then Deferred.unit
      else (
        try_with (fun () -> 
            Tcp.connect ~interrupt:stop where_to_connect
            >>= fun (_,r,w) ->
            Rpc.Connection.create r w ~connection_state:()
            >>= function
            | Error exn -> raise exn
            | Ok conn -> return conn            
          )
        >>= function
        | Error _ ->
          after reconnect_after
          >>= fun () ->
          connect_loop ()
        | Ok conn ->
          t := Some conn;
          Rpc.Connection.close_finished conn
          >>= fun () ->
          t := None;
          connect_loop ()
      )
    in
    don't_wait_for (connect_loop ());
    t

  let with_connect ~reconnect_after where_to_connect k =
    let stop = Ivar.create () in
    let t = create where_to_connect
        ~reconnect_after ~stop:(Ivar.read stop)
    in
    Monitor.protect (fun () -> k t)
      ~finally:(fun () -> Ivar.fill stop (); Ivar.read stop)
end

let setup_robust_conn ~reconnect_after username k =
  load_config ()
  >>= fun config ->
  resolve_username username config
  >>= fun username ->
  let host = config.Config.server in
  let port = Common.port in
  Robust_connection.with_connect
    ~reconnect_after
    (Tcp.to_host_and_port host port)
    (fun conn -> k username conn)