Improve cli abstractions

Use `Cmdliner.Arg.file` conv to parse valid files. Cmdliners' file
conv already does a `Sys.file_exists` so we avoid code duplication
here.

Also define `Mode.t` as an algebraic datatype to reduce string
allocations and allow for faster instant `is_` mode checks.

bin/main.ml

@@ -10,24 +10,24 @@ open Auxiliary.Functions
 
 let env_path = Filename.dirname (Filename.dirname Sys.executable_name) ^ "/lib/ast_gen/";;
 
-let setup_output (output_path : string) : (string * string * string) = 
-  let code_dir = output_path  ^ "/code/" in 
-  let graph_dir = output_path ^ "/graph/" in 
+let setup_output (output_path : string) : (string * string * string) =
+  let code_dir = output_path  ^ "/code/" in
+  let graph_dir = output_path ^ "/graph/" in
   let run_dir   = output_path ^ "/run/" in
 
   File_system.clean_dir output_path;
   File_system.create_dir code_dir;
   File_system.create_dir graph_dir;
   File_system.create_dir run_dir;
-  
+
   code_dir, graph_dir, run_dir
 
-let setup_node (mode : string) : string = 
+let setup_node (mode : Mode.t) : string =
   (* !REFACTOR : inbed this functionality into the dune build and install process *)
-  let js_folder    = env_path ^ "js/" in 
-  let node_modules = js_folder ^ "node_modules" in 
-  let package_info = js_folder ^ "package.json" in 
-  let script       = js_folder ^ "generate_cg.js" in 
+  let js_folder    = env_path ^ "js/" in
+  let node_modules = js_folder ^ "node_modules" in
+  let package_info = js_folder ^ "package.json" in
+  let script       = js_folder ^ "generate_cg.js" in
 
   if not (Sys.file_exists script)
     then failwith "[ERROR] Dependency tree genetarion script not found";
@@ -38,63 +38,63 @@ let setup_node (mode : string) : string =
   if Mode.is_multi_file mode && not (Sys.file_exists node_modules)
     then (
       print_endline "installing js dependencies";
-      let result = Sys.command ("npm install --prefix " ^ js_folder) in 
+      let result = Sys.command ("npm install --prefix " ^ js_folder) in
       if result != 0 then failwith "[ERROR] Unable to install js depedencies";
       print_endline "DONE!");
 
   script
 
 
-let main (filename : string) (output_path : string) (config_path : string) (mode : string) (generate_mdg : bool) (no_dot : bool) (verbose : bool) : int =
-  
+let main (filename : string) (output_path : string) (config_path : string) (mode : Mode.t) (generate_mdg : bool) (no_dot : bool) (verbose : bool) : int =
+
   (* SETUP *)
   let script = setup_node mode in
-  let dep_tree = DependencyTree.generate script filename mode in  
-  let code_dir, graph_dir, _ = setup_output output_path in 
+  let dep_tree = DependencyTree.generate script filename mode in
+  let code_dir, graph_dir, _ = setup_output output_path in
 
   (* process dependencies first with the aid of the depedency tree *)
-  let summaries = Summaries.empty () in 
-  let module_graphs = ModuleGraphs.empty () in 
-  List.iter (fun file_path -> 
-    let dir = Filename.dirname file_path ^ "/" in 
-    let filename = File_system.file_name file_path in 
+  let summaries = Summaries.empty () in
+  let module_graphs = ModuleGraphs.empty () in
+  List.iter (fun file_path ->
+    let dir = Filename.dirname file_path ^ "/" in
+    let filename = File_system.file_name file_path in
 
     (* STEP 0 : Generate AST using Flow library *)
-    let ast = Js_parser.from_file file_path in 
+    let ast = Js_parser.from_file file_path in
 
     (* STEP 1 : Normalize AST *)
     let norm_program = Ast.Normalize.program ast file_path in
-    let norm_program = if file_path = dep_tree.main then Program.set_main norm_program else norm_program in 
+    let norm_program = if file_path = dep_tree.main then Program.set_main norm_program else norm_program in
     let js_program = Ast.Pp.Js.print norm_program in
     File_system.write_to_file (code_dir ^ filename) js_program;
 
     (* STEP 2 : Generate MDG for the normalized code *)
     if generate_mdg then (
       let graph, exportedObject, external_calls = Mdg.Analyse.program mode verbose config_path norm_program in
-    
+
       ExternalReferences.iter (fun locs info ->
         let l_call = LocationSet.min_elt locs in
-        
-        (* module information *) 
-        let module_name = dir ^ info._module in 
-        let moduleEO = Summaries.get_opt summaries module_name in 
+
+        (* module information *)
+        let module_name = dir ^ info._module in
+        let moduleEO = Summaries.get_opt summaries module_name in
         option_may (fun moduleEO ->
-          let moduleGraph = ModuleGraphs.get module_graphs module_name in 
-          
+          let moduleGraph = ModuleGraphs.get module_graphs module_name in
+
           (* exported function information *)
-          let func_loc = ExportedObject.get_value_location moduleEO info.properties in 
+          let func_loc = ExportedObject.get_value_location moduleEO info.properties in
           if not (Graph.has_external_function graph func_loc) then (
-            let func_graph = Graph.get_function moduleGraph func_loc in 
+            let func_graph = Graph.get_function moduleGraph func_loc in
             Graph.add_external_func graph func_graph l_call func_loc
           );
-          
+
           Graph.add_call_edge graph l_call func_loc;
         ) moduleEO;
 
       ) external_calls;
-      
+
       (* save current module info*)
-      let alter_name = String.sub file_path 0 (String.length file_path - 3) in 
+      let alter_name = String.sub file_path 0 (String.length file_path - 3) in
       ModuleGraphs.add module_graphs file_path graph;
       ModuleGraphs.add module_graphs alter_name graph;
       Summaries.add summaries file_path exportedObject;
@@ -105,9 +105,9 @@ let main (filename : string) (output_path : string) (config_path : string) (mode
 
   (* output *)
   if generate_mdg then (
-    let main = DependencyTree.get_main dep_tree in 
-    let graph = ModuleGraphs.get module_graphs main in 
-    if not no_dot then 
+    let main = DependencyTree.get_main dep_tree in
+    let graph = ModuleGraphs.get module_graphs main in
+    if not no_dot then
       Mdg.Pp.Dot.output graph_dir graph;
     Mdg.Pp.CSV.output graph_dir graph;
   );
@@ -118,18 +118,18 @@ let main (filename : string) (output_path : string) (config_path : string) (mode
 let input_file : string Term.t =
   let doc = "Path to JavaScript file (.js) or directory containing JavaScript files for analysis." in
   let docv = "FILE_OR_DIR" in
-  let is_dir_or_file (param : string) : string =  
-    if Sys.file_exists param 
-      then param
-      else failwith ("[ERROR] Invalid input file : " ^ param) 
-  in 
-
-  Term.(const is_dir_or_file $  Arg.(required & pos 0 (some string) None & info [] ~doc ~docv))
-
-
-let mode : string Term.t =
+  Arg.(required & pos 0 (some file) None & info [] ~doc ~docv)
+
+let mode : Mode.t Term.t =
+  let mode_enum =
+    Arg.enum
+      [ ("basic", Mode.Basic)
+      ; ("single_file", Mode.Single_file)
+      ; ("multi_file", Mode.Multi_file);
+      ]
+  in
   let doc = "Analysis mode.\n\t 1) basic: attacker controlls all parameters from all functions \n\t 2) single_file: the attacker controlls the functions that were exported by the input file \n\t 3) multi_file: the attacker controlls the functions that were exported in the \"main\" file" in
-  Term.(const Mode.is_valid $ Arg.(value & opt string Mode.default & info ["m"; "mode"] ~doc))
+  Arg.(value & opt mode_enum Mode.single_file & info ["m"; "mode"] ~doc)
 
 let mdg : bool Term.t =
   let doc = "Generates Multiversion Dependency Graph." in
@@ -141,12 +141,12 @@ let no_dot : bool Term.t =
 
 let output_path : string Term.t =
   let doc = "Path to store all output files." in
-  let default_path = "graphjs-results" in 
+  let default_path = "graphjs-results" in
   Arg.(value & opt string default_path & info ["o"; "output"] ~doc)
 
 let config_path : string Term.t =
   let doc = "Path to configuration file." in
-  let default_path = env_path ^ "config.json" in 
+  let default_path = env_path ^ "config.json" in
   Arg.(value & opt non_dir_file default_path & info ["c"; "config"] ~doc)
 
 let verbose : bool Term.t =

lib/auxiliary/mode.ml

@@ -1,14 +1,22 @@
-let basic : string = "basic"
-let single_file : string = "single_file"
-let multi_file : string = "multi_file"
-let default : string = single_file
+type t = Basic | Single_file | Multi_file
 
-let is_valid (mode : string) : string =  
-  if mode = basic || mode = single_file || mode = multi_file 
-    then mode
-    else failwith "[ERROR] Invalid mode. Try using: basic, single_file or multi_file"
+let basic = Basic
+let single_file = Single_file
+let multi_file = Multi_file
+let is_basic = function Basic -> true | _ -> false
+let is_single_file = function Single_file -> true | _ -> false
+let is_multi_file = function Multi_file -> true | _ -> false
 
+let of_string = function
+  | "basic" -> Ok Basic
+  | "single_file" -> Ok Single_file
+  | "multi_file" -> Ok Multi_file
+  | _ ->
+      Error "[ERROR] Invalid mode. Try using: basic, single_file or multi_file"
 
-let is_basic (mode : string) : bool = mode = basic
-let is_single_file (mode : string) : bool = mode = single_file
-let is_multi_file (mode : string) : bool = mode = multi_file
+let pp fmt = function
+  | Basic -> Format.pp_print_string fmt "basic"
+  | Single_file -> Format.pp_print_string fmt "single_file"
+  | Multi_file -> Format.pp_print_string fmt "multi_file"
+
+let to_string mode = Format.asprintf "%a" pp mode

lib/auxiliary/mode.mli

@@ -0,0 +1,19 @@
+type t = Basic | Single_file | Multi_file
+
+val basic : t
+
+val single_file : t
+
+val multi_file : t
+
+val is_basic : t -> bool
+
+val is_single_file : t -> bool
+
+val is_multi_file : t -> bool
+
+val of_string : string -> (t, string) result
+
+val to_string : t -> string
+
+val pp : Format.formatter -> t -> unit