X-Git-Url: http://git.annexia.org/?p=virt-mem.git;a=blobdiff_plain;f=lib%2Fvirt_mem.ml;h=fcefa3cc4bc8492df585f19417441e21d0525cd1;hp=dd3c18a304dd567190ce3ff890049730dd236803;hb=e3b79ee907537feec9274b1bfab7e450fc97dbcf;hpb=35f68724b76ea85e9b949f05ad82b82a7144989c

diff --git a/lib/virt_mem.ml b/lib/virt_mem.ml
index dd3c18a..fcefa3c 100644
--- a/lib/virt_mem.ml
+++ b/lib/virt_mem.ml
@@ -43,7 +43,14 @@ let max_memory_peek = 65536 (* XXX Use D.max_peek function *)
 type ksym = string
 
 type image =
-    string
+    int option
+    * string
+    * Virt_mem_utils.architecture
+    * ([`Wordsize], [`Endian]) Virt_mem_mmap.t
+
+type image_with_ksyms =
+    int option
+    * string
     * Virt_mem_utils.architecture
     * ([`Wordsize], [`Endian]) Virt_mem_mmap.t
     * (ksym -> MMap.addr)
@@ -52,9 +59,120 @@ type kallsyms_compr =
   | Compressed of (string * MMap.addr) list * MMap.addr
   | Uncompressed of (string * MMap.addr) list
 
-let start usage_msg =
-  (* Debug messages. *)
+(* When tools register themselves, they are added to this list.
+ * Later, we will alphabetize the list.
+ *)
+let tools = ref []
+
+(* Registration function used by the tools. *)
+let register ?(external_cmd = true) ?(extra_args = [])
+    ?argcheck ?beforeksyms ?run
+    name summary description =
+  tools :=
+    (name, (name, summary, description, external_cmd, extra_args,
+	    argcheck, beforeksyms, run))
+  :: !tools
+
+(* Main program, called from mem/virt_mem_main.ml when all the
+ * tools have had a chance to register themselves.
+ *)
+let main () =
+  (* Get the registered tools, alphabetically. *)
+  let tools = !tools in
+  let tools = List.sort ~cmp:(fun (a,_) (b,_) -> compare a b) tools in
+
+  (* Which tool did the user want to run?  Look at the executable
+   * name (eg. 'virt-dmesg' => tool == dmesg).  If we don't recognise
+   * the executable name then we must look for the first parameter
+   * which doesn't begin with a '-' character.
+   *
+   * Note that we must do all of this before using the OCaml Arg
+   * module to properly parse the command line (below), so that
+   * we can have a usage message ready.
+   *)
+  let tool, ignore_first_anon_arg =
+    let prog = Sys.executable_name in	(* eg. "/usr/bin/virt-dmesg.opt" *)
+    let prog = Filename.basename prog in(* eg. "virt-dmesg.opt" *)
+    let prog =				(* eg. "virt-dmesg" *)
+      try Filename.chop_extension prog with Invalid_argument _ -> prog in
+    let prog =				(* eg. "dmesg" *)
+      if String.starts_with prog "virt-" then
+	String.sub prog 5 (String.length prog - 5)
+      else prog in
+    try Some (List.assoc prog tools), false
+    with Not_found ->
+      let arg1 =			(* First non-option argument. *)
+	match Array.to_list Sys.argv with
+	| [] -> None
+	| _::args ->
+	    let rec loop = function
+	      | [] -> None
+	      | a::args when String.length a > 0 && a.[0] = '-' -> loop args
+	      | a::_ -> Some a
+	    in
+	    loop args in
+      match arg1 with
+      | None -> None, false
+      | Some prog ->			(* Recognisable first argument? *)
+	  let prog =
+	    try Filename.chop_extension prog with Invalid_argument _ -> prog in
+	  let prog =
+	    if String.starts_with prog "virt-" then
+	      String.sub prog 5 (String.length prog - 5)
+	    else prog in
+	  (try Some (List.assoc prog tools), true
+	   with Not_found -> None, false) in
+
+  (* Make a usage message. *)
+  let usage_msg =
+    match tool with
+    | None ->				(* Generic usage message. *)
+	let tools = List.map (
+	  fun (name, (_, summary, _, external_cmd, _, _, _, _)) ->
+	    if external_cmd then "virt-"^name, summary
+	    else                 "virt-mem "^name, summary
+	) tools in
+	(* Maximum width of field in the left hand column. *)
+	let max_width =
+	  List.fold_left max 0 (List.map String.length (List.map fst tools)) in
+	let tools = List.map (fun (l,r) -> pad max_width l, r) tools in
+	let tools = List.map (fun (l,r) -> "  " ^ l ^ " - " ^ r) tools in
+	let tools = String.concat "\n" tools in
+
+	sprintf (f_"\
+
+virt-mem: Tools for providing information about virtual machines
+
+Currently available tools include:
+%s
+
+General usage is:
+  <tool> [-options] [domains...]
+
+To display extra help for a single tool, do:
+  virt-mem help <tool>
+
+Options:") tools
+
+                                        (* Tool-specific usage message. *)
+    | Some (name, summary, description, external_cmd, _, _, _, _) ->
+	let cmd =
+	  if external_cmd then "virt-" ^ name else "virt-mem " ^ name in
+
+	sprintf (f_"\
+
+%s: %s
+
+Description:
+%s
+
+Options:") cmd summary description in
+
+  (* Now begin proper parsing of the command line arguments. *)
   let debug = ref false in
+  let images = ref [] in
+  let uri = ref "" in
+  let anon_args = ref [] in
 
   (* Default wordsize. *)
   let def_wordsize = ref None in
@@ -70,9 +188,9 @@ let start usage_msg =
   let set_endian = function
     | "auto" -> def_endian := None
     | "le" | "little" | "littleendian" | "intel" ->
-	def_endian := Some Bitmatch.LittleEndian
+	def_endian := Some Bitstring.LittleEndian
     | "be" | "big" | "bigendian" | "motorola" ->
-	def_endian := Some Bitmatch.BigEndian
+	def_endian := Some Bitstring.BigEndian
     | str -> failwith (sprintf (f_"set_endian: %s: unknown endianness") str)
   in
 
@@ -96,17 +214,14 @@ let start usage_msg =
     | str -> def_text_addr := Int64.of_string str
   in
 
-  (* List of kernel images. *)
-  let images = ref [] in
-  let uri = ref "" in
-  let anon_args = ref [] in
-
+  (* Handle -t option. *)
   let memory_image filename =
     images :=
       (!def_wordsize, !def_endian, !def_architecture, !def_text_addr, filename)
     :: !images
   in
 
+  (* Handle --version option. *)
   let version () =
     printf "virt-mem %s\n" Virt_mem_version.version;
 
@@ -117,35 +232,82 @@ let start usage_msg =
     exit 0
   in
 
-  let argspec = Arg.align [
-    "-A", Arg.String set_architecture,
-      "arch " ^ s_"Set kernel architecture, endianness and word size";
-    "-E", Arg.String set_endian,
-      "endian " ^ s_"Set kernel endianness";
-    "-T", Arg.String set_text_addr,
-      "addr " ^ s_"Set kernel text address";
-    "-W", Arg.String set_wordsize,
-      "addr " ^ s_"Set kernel word size";
-    "-c", Arg.Set_string uri,
-      "uri " ^ s_ "Connect to URI";
-    "--connect", Arg.Set_string uri,
-      "uri " ^ s_ "Connect to URI";
-    "--debug", Arg.Set debug,
-      " " ^ s_"Debug mode (default: false)";
-    "-t", Arg.String memory_image,
-      "image " ^ s_"Use saved kernel memory image";
-    "--version", Arg.Unit version,
-      " " ^ s_"Display version and exit";
-  ] in
-
+  (* Function to collect up any anonymous args (domain names/IDs). *)
   let anon_arg str = anon_args := str :: !anon_args in
-  let usage_msg = usage_msg ^ s_"\n\nOPTIONS" in
+
+  (* Construct the argspec.
+   * May include extra arguments specified by the tool.
+   *)
+  let argspec =
+    let extra_args = match tool with
+      | None -> []
+      | Some (_, _, _, _, extra_args, _, _, _) -> extra_args in
+    let argspec = [
+      "-A", Arg.String set_architecture,
+        "arch " ^ s_"Set kernel architecture, endianness and word size";
+      "-E", Arg.String set_endian,
+        "endian " ^ s_"Set kernel endianness";
+      "-T", Arg.String set_text_addr,
+        "addr " ^ s_"Set kernel text address";
+      "-W", Arg.String set_wordsize,
+        "addr " ^ s_"Set kernel word size";
+      "-c", Arg.Set_string uri,
+        "uri " ^ s_ "Connect to URI";
+      "--connect", Arg.Set_string uri,
+        "uri " ^ s_ "Connect to URI";
+      "--debug", Arg.Set debug,
+        " " ^ s_"Debug mode (default: false)";
+      "-t", Arg.String memory_image,
+        "image " ^ s_"Use saved kernel memory image";
+      "--version", Arg.Unit version,
+        " " ^ s_"Display version and exit";
+    ] @ extra_args in
+
+    (* Sort options alphabetically on first alpha character. *)
+    let cmp (a,_,_) (b,_,_) =
+      let chars = "-" in
+      let a = String.strip ~chars a and b = String.strip ~chars b in
+      compare a b
+    in
+    let argspec = List.sort ~cmp argspec in
+    (* Make the options line up nicely. *)
+    Arg.align argspec in
+
+  (* Parse the command line.  This will exit if --version or --help found. *)
   Arg.parse argspec anon_arg usage_msg;
 
   let images = !images in
   let debug = !debug in
   let uri = if !uri = "" then None else Some !uri in
+
+  (* Discard the first anonymous argument if, above, we previously
+   * found it contained the tool name.
+   *)
   let anon_args = List.rev !anon_args in
+  let anon_args =
+    if ignore_first_anon_arg then List.tl anon_args else anon_args in
+
+  (* At this point, either --help was specified on the command line
+   * (and so the program has exited) or we must have determined tool,
+   * or the user didn't give us a valid tool (eg. "virt-mem foobar").
+   * Detect that final case now and give an error.
+   *)
+  let name, _, _, _, _, argcheck, beforeksyms, run =
+    match tool with
+    | Some t -> t
+    | None ->
+	prerr_endline (s_"\
+virt-mem: I could not work out which tool you are trying to run.
+Use 'virt-mem --help' for more help or read the manual page virt-mem(1)");
+	exit 1
+  in
+  if debug then eprintf "tool = %s\n%!" name;
+
+  (* Optional argument checking in the tool. *)
+  (match argcheck with
+   | None -> ()
+   | Some argcheck -> argcheck debug
+  );
 
   (* Get the kernel images. *)
   let images =
@@ -210,6 +372,7 @@ let start usage_msg =
 
       List.map (
 	fun (dom, _) ->
+	  let id = D.get_id dom in
 	  let name = D.get_name dom in
 
 	  let wordsize =
@@ -268,7 +431,7 @@ let start usage_msg =
 	  let mem = MMap.set_wordsize mem wordsize in
 	  let mem = MMap.set_endian mem endian in
 
-	  (name, arch, mem)
+	  ((Some id, name, arch, mem) : image)
       ) xmls
     ) else (
       (* One or more -t options passed. *)
@@ -316,13 +479,27 @@ let start usage_msg =
 	  let mem = MMap.set_wordsize mem wordsize in
 	  let mem = MMap.set_endian mem endian in
 
-	  (filename, arch, mem)
+	  ((None, filename, arch, mem) : image)
       ) images
     ) in
 
+  (* Optional callback into the tool before we start looking for
+   * kernel symbols.
+   *)
+  (match beforeksyms with
+   | None -> ()
+   | Some beforeksyms -> beforeksyms debug images
+  );
+
+  (* If there is no run function, then there is no point continuing
+   * with the rest of the program (kernel symbol analysis) ...
+   *)
+  if run = None then exit 0;
+
+  (* Now kernel symbol analysis starts ... *)
   let images =
     List.map (
-      fun (name, arch, mem) ->
+      fun (domid, name, arch, mem) ->
 	(* Look for some common entries in the exported symbol table and
 	 * from that find the symbol table itself.  These are just
 	 * supposed to be symbols which are very likely to be present
@@ -450,7 +627,7 @@ let start usage_msg =
 
 	      (* Load the whole symbol table as a bitstring. *)
 	      let ksymtab =
-		Bitmatch.bitstring_of_string
+		Bitstring.bitstring_of_string
 		  (MMap.get_bytes mem ksymtab_addr
 		     (Int64.to_int ksymtab_size)) in
 
@@ -727,7 +904,12 @@ let start usage_msg =
 	    lookup_ksym
 	in
 
-	((name, arch, mem, lookup_ksym) : image)
+	((domid, name, arch, mem, lookup_ksym) : image_with_ksyms)
     ) images in
 
-  debug, images
+  (* Run the tool's main function. *)
+  (match run with
+   | None -> ()
+   | Some run ->
+       run debug images
+  )