agent-snapshot.git

index f6d74c7..20f07fe 100644
--- a/dune-project
+++ b/dune-project
@@ -7,6 +7,7 @@
  (depends
   (ocaml (>= 5.4))
   dune
+  cmdliner
   yojson
   ppx_deriving_yojson
   camomile

index 2cb5d54..5639d6b 100644
--- a/src/ocaml/agent_snapshot.ml
+++ b/src/ocaml/agent_snapshot.ml
@@ -718,47 +718,63 @@ let run_snapshot (output : string) (command : string list) : int =
   write_manifest output command 0;
   0
 
+open Cmdliner
+
+let output_arg =
+  Arg.(
+    required
+    & opt (some string) None
+    & info [ "o"; "output" ] ~docv:"SNAPDIR"
+        ~doc:"Write the snapshot under this directory.")
+
+let command_arg =
+  Arg.(
+    value
+    & pos_all string []
+    & info [] ~docv:"ARG" ~doc:"Command to trace and its arguments.")
+
+let restore_snapdir =
+  Arg.(
+    required
+    & pos 0 (some string) None
+    & info [] ~docv:"SNAPDIR"
+        ~doc:"Snapshot directory produced by a prior run.")
+
+let snapshot_term : unit Term.t =
+  Term.(
+    const (fun output cmd ->
+      if cmd = [] then
+        failwith "usage: agent-snapshot --output SNAPDIR -- command args...";
+      ignore (run_snapshot output cmd : int))
+    $ output_arg
+    $ command_arg)
+
+let restore_term : unit Term.t =
+  Term.(const (fun dir -> restore_snapshot dir) $ restore_snapdir)
+
+let cmd_main : unit Cmd.t =
+  let doc = "Filesystem snapshotter for traced commands" in
+  let man =
+    [
+      `S Manpage.s_description;
+      `P "Traces a command with ptrace and records a manifest plus file blobs.";
+      `S Manpage.s_examples;
+      `P "$(tool) $(b,--output) /tmp/snap $(b,--) make all";
+      `P "$(tool) $(b,--output) /tmp/snap make all";
+      `P "$(tool) $(b,restore) /tmp/snap";
+    ]
+  in
+  let main_info = Cmd.info "agent-snapshot" ~doc ~man in
+  let restore_info =
+    Cmd.info "restore" ~docs:Manpage.s_commands
+      ~doc:"Restore blobbed files and tombstones from a snapshot directory."
+  in
+  let restore_cmd = Cmd.v restore_info restore_term in
+  Cmd.group main_info ~default:snapshot_term [ restore_cmd ]
+
 let main () : unit =
   try
-    let argv = Sys.argv in
-    let rc =
-      if Array.length argv >= 2 && argv.(1) = "restore" then (
-        if Array.length argv < 3 then failwith "usage: agent-snapshot restore SNAPDIR";
-        if Array.length argv > 3 then failwith "usage: agent-snapshot restore SNAPDIR";
-        restore_snapshot argv.(2);
-        0)
-      else
-        let output = ref "" in
-        let command_rev = ref [] in
-        let anon s = command_rev := s :: !command_rev in
-        let speclist =
-          [
-            ( "--output",
-              Arg.Set_string output,
-              "SNAPDIR  Write the snapshot under this directory (required)." );
-          ]
-        in
-        let aligned = Arg.align speclist in
-        let usage = "agent-snapshot --output SNAPDIR -- command [args...]\nOptions:" in
-        let command =
-          let n = Array.length argv in
-          let rec ddash i =
-            if i >= n then None else if argv.(i) = "--" then Some i else ddash (i + 1)
-          in
-          match ddash 1 with
-          | Some i ->
-              if i + 1 >= n then failwith "usage: agent-snapshot --output SNAPDIR -- command args...";
-              let opt_argv = Array.sub argv 0 i in
-              Arg.parse_argv ~current:(ref 0) opt_argv aligned anon usage;
-              Array.to_list (Array.sub argv (i + 1) (n - i - 1))
-          | None ->
-              Arg.parse aligned anon usage;
-              List.rev !command_rev
-        in
-        if !output = "" || command = [] then
-          failwith "usage: agent-snapshot --output SNAPDIR -- command args...";
-        run_snapshot !output command
-    in
+    let rc = Cmd.eval ~catch:false cmd_main in
     Ocaml_git.shutdown ();
     exit rc
   with exn ->

index 0ddbecc..7f9c26f 100644
--- a/src/ocaml/dune
+++ b/src/ocaml/dune
@@ -6,4 +6,4 @@
  (foreign_stubs
   (language c)
   (names ptrace_stubs))
- (libraries unix yojson camomile ocaml-git fpath fileutils arrow.c_api arrow.ppx_runtime))
+ (libraries unix cmdliner yojson camomile ocaml-git fpath fileutils arrow.c_api arrow.ppx_runtime))