Source

screentime-monitor / client.ml

Full commit
open Core.Std
open Async.Std
module Shell = Core_extended.Std.Shell

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 setup_conn username k =
  load_config ()
  >>= fun config ->
  (match username with
   | Some x -> return x
   | None ->
     match config.Config.user with
     | Some x -> return x
     | None -> get_username ())
  >>= fun username ->
  Common.with_rpc_conn
    ~host:config.Config.server
    ~port:Common.port
    (fun conn -> k username conn)

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

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 start =
  Command.async_basic
    ~summary:"Start the timer"
    Command.Spec.(
      shared_flags ()
      +> anon ("category" %: category_spec)
    )
    (fun username category () ->
       setup_conn username (fun username conn ->
         Rpc.Rpc.dispatch_exn Protocol.start conn (username,category)))

let stop =
  Command.async_basic
    ~summary:"Stop the timer"
    (shared_flags ())
    (fun username () -> setup_conn username (fun username conn ->
       Rpc.Rpc.dispatch_exn Protocol.stop conn username))

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 "Failed to parse rule:\n%s\n"
        (e |> Error.sexp_of_t |> Sexp.to_string_hum);
      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 spinner =
  let states = [| "|" ; "/"; "-"; "\\" |] in
  let pos = ref 0 in
  (fun () ->
     pos := (!pos + 1) % Array.length states;
     states.(!pos)
  )

let span_string span =
  let { Time.Span.Parts.
        sign = _
      ; hr; min; sec
      ; ms = _ ; us = _
      } =
    Time.Span.to_parts span
  in
  sprintf "%02d:%02d:%02d" hr min sec
;;

let print_violations violations =
  let module Ascii_table = Textutils.Std.Ascii_table in
  force clear_string
  >>= fun clear_string ->
  printf "%s" clear_string;
  let cols = 
    Ascii_table.(
      [ Column.create "rule"  (fun (r,_,_) -> r.Rule.name |> Rule.Name.to_string)
      ; Column.create "time"  (fun (_,s,_) -> Time.Span.to_string s)
      ; Column.create "status" (fun (_,_,(status : Rule_store.Status.t)) ->
          match status with Acked -> "X" | Unacked -> " ")
      ]
    )
  in
  printf "%s\n" (Ascii_table.to_string ~display:Ascii_table.Display.line cols violations);
  return ()
;;

let monitor_violations user conn ~stop =
  Clock.every' (sec 10.) ~stop (fun () ->
    Rpc.Rpc.dispatch_exn Protocol.todays_violations conn user
    >>= print_violations
  )
;;

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


let monitor_violations =
  Command.async_basic
    ~summary:"Monitor for violations"
    (shared_flags ())
    (fun user () -> setup_conn user (fun user conn ->
         let stop = stop_on_term () in
         monitor_violations user conn ~stop;
         stop
       ))

let monitor =
  Command.async_basic
    ~summary:"Monitor your time spent"
    (shared_flags ())
    (fun user () -> setup_conn user (fun user conn ->
       force clear_string
       >>= fun cstring ->
       let stop = stop_on_term  () in
       Clock.every' (sec 1.) ~stop (fun () ->
         Rpc.Rpc.dispatch_exn Protocol.status conn user
         >>| fun { Protocol.Status. state; elapsed } ->
         let total_elapsed =
           Map.data elapsed
           |> List.fold ~init:Time.Span.zero ~f:Time.Span.(+)
         in
         printf "%s%s %s %s\n"
           cstring
           (spinner ())
           (span_string total_elapsed)
           (match state with
            | `Active cat -> sprintf "{%s}" (Category.to_string cat)
            | `Not_running -> "(paused)"
           )
       );
       stop
     ))

module Ascii_table = Textutils.Ascii_table

let report =
  Command.async_basic
    ~summary:"Print a report of time spent"
    Command.Spec.(
      shared_flags ()
      +> days_flag ()
    )
    (fun user days () -> setup_conn user (fun user conn ->
       Rpc.Rpc.dispatch_exn Protocol.report conn user
       >>= fun report ->
       let report =
         Map.to_alist report
         |> (fun l -> List.drop l (List.length l - days))
       in
       let categories =
         List.map ~f:snd report
         |> List.map ~f:Map.keys
         |> List.concat
         |> List.dedup
       in
       let open Ascii_table in
       let date_column =
         Column.create "date"
           (fun (date,_) -> Date.to_string date)
       in
       let cat_columns =
         List.map categories ~f:(fun cat ->
           Column.create
             (Category.to_string cat)
             (fun (_date,by_cat) ->
                match Map.find by_cat cat with
                | None -> "-"
                | Some span -> span_string span
             ))
       in
       let total_column =
         Column.create "total"
           (fun (_date,by_cat) ->
              Map.data by_cat
              |> List.fold ~init:Time.Span.zero ~f:Time.Span.(+)
              |> span_string
           )
       in
       printf "%s\n"
         (to_string
            ~display:Display.column_titles
            ([date_column]
             @ (if List.length categories > 1 then cat_columns else [])
             @ [total_column]
            )
            report);
       return ()
     ))

let full_report =
  Command.async_basic
   ~summary:"Print a detailed report of all the screentime sessions"
   Command.Spec.(
     shared_flags ()
     +> days_flag ()
   )
   (fun user days () -> setup_conn user (fun user conn ->
      Rpc.Rpc.dispatch_exn Protocol.full_report conn user
      >>= fun report ->
      let open Ascii_table in
      let ofday_string t =
        Time.to_local_ofday t |> Time.Ofday.to_sec_string
      in
      let columns =
        [ Column.create "category" (fun (c,_,_) -> Category.to_string c)
        ; Column.create "start"    (fun (_,t,_) -> ofday_string t)
        ; Column.create "stop"     (fun (_,_,t) -> ofday_string t)
        ; Column.create "min"     (fun (_,start,stop) ->
          Time.diff stop start
          |> Time.Span.to_min
          |> Float.to_int
          |> Int.to_string)
        ]
      in
      Map.to_alist report
      |> (fun l -> List.drop l (List.length l - days))
      |> List.iter ~f:(fun (date,sessions) ->
        printf "*** Date: %s ***\n\n" (Date.to_string date);
        printf "%s\n" (to_string
                         ~display:Display.column_titles
                         columns
                         (List.rev sessions)));
      Deferred.unit
    ))

let rules =
  let show_rules =
    Command.async_basic
      ~summary:"Retrieve the screentime rules"
      (shared_flags ())
      (fun username () -> setup_conn username (fun username conn ->
         Rpc.Rpc.dispatch_exn Protocol.get_rules conn username
         >>| fun rules ->
         printf "%s\n" (With_format.format rules)
       ))
  in
  let edit_rules =
    Command.async_basic
      ~summary:"Set the screentime rule"
      (shared_flags ())
      (fun username () -> setup_conn username (fun username conn ->
         Rpc.Rpc.dispatch_exn Protocol.get_rules conn username
         >>= fun rule ->
         let tempfile = Filename.temp_file "rule" ".scm" in
         Writer.save tempfile ~contents:(With_format.format rule)
         >>= fun () ->
         let editor =
           match Sys.getenv "EDITOR" with None -> "emacs" | Some x -> x
         in
         edit_file (module Rule.List) ~editor ~tempfile
         >>= function
         | None -> return ()
         | Some rules ->
           printf "Set the rules for %s? (y/N):" (Username.to_string username);
           Reader.read_line (Lazy.force Reader.stdin)
           >>= fun response ->
           let upload =
             match response with
             | `Eof -> false
             | `Ok s ->
               match s |> String.lowercase |> String.strip with
               | "y" | "yes" -> true
               | _ -> false
           in
           if not upload then
             (printf "Not setting rules.\n"; return ())
           else
             Rpc.Rpc.dispatch_exn Protocol.set_rules conn (username,rules)
             >>| fun () ->
             printf "Rules set\n"
       ))
  in
  let rule_violations =
    Command.async_basic
      ~summary:"Get any rule violations"
      (shared_flags ())
      (fun username () -> setup_conn username (fun username conn ->
         Rpc.Rpc.dispatch_exn Protocol.todays_violations conn username
         >>| fun violations ->
         violations
         |> <:sexp_of<(Rule.t * Time.Span.t * Rule_store.Status.t) list>>
         |> Sexp.to_string_hum
         |> printf "%s\n"
       ))
  in
  Command.group
    ~summary:"Tools for interacting with screentime rules"
    [ "show", show_rules
    ; "edit", edit_rules
    ; "violations", rule_violations
    ; "monitor",   monitor_violations
    ]

let commands =
  [ "start"       , start
  ; "stop"        , stop
  ; "monitor"     , monitor
  ; "report"      , report
  ; "full-report" , full_report
  ; "rules"       , rules
  ]