Use clap to parse command arguments

.gitignore

@@ -0,0 +1,10 @@
+/*
+!/.cargo
+!/.github
+!/pics
+!/src
+!/.gitignore
+!/Cargo.lock
+!/Cargo.toml
+!/LICENSE
+!/README.md

Cargo.toml

@@ -11,6 +11,7 @@ crate-type = ["cdylib"]
 [dependencies]
 anyhow = { version = "1.0.79" }
 chrono = "0.4.33"
+clap = { version = "4.5.0", features = ["derive"] }
 serde = { version = "1.0.195", features = ["derive"] }
 serde_json = { version = "1.0.111"}
 windows = { version = "0.52.0", features = ["Win32_Foundation", "Win32_System", "Win32_System_Diagnostics", "Win32_System_Diagnostics_Debug", "Win32_System_Diagnostics_Debug_Extensions", "Win32_System_SystemInformation" ] }

src/lib.rs

@@ -12,6 +12,7 @@ use std::{env, mem};
 
 use anyhow::{bail, Result};
 use chrono::Local;
+use clap::{Parser, ValueEnum};
 use debug_client::DebugClient;
 use serde_json::Value;
 use state::{Float80, GlobalSeg, State, Zmm};
@@ -40,11 +41,6 @@ mod msr {
     pub const TSC_AUX: u32 = 0xc000_0103;
 }
 
-enum SnapshotKind {
-    ActiveKernel,
-    Full,
-}
-
 /// Check if an address lives in user-mode.
 /// https://learn.microsoft.com/en-us/windows-hardware/drivers/gettingstarted/virtual-address-spaces
 fn is_usermode_addr(addr: u64) -> bool {
@@ -291,9 +287,26 @@ fn state(dbg: &DebugClient) -> Result<State> {
     })
 }
 
+#[derive(Clone, Default, ValueEnum)]
+enum SnapshotKind {
+    #[default]
+    ActiveKernel,
+    Full,
+}
+
+#[derive(Parser)]
+#[clap(no_binary_name = true)]
+struct SnapshotArgs {
+    /// The kind of snapshot to take.
+    #[clap(default_value = "full", long, short)]
+    kind: SnapshotKind,
+    /// The path to save the snapshot to.
+    state_path: Option<PathBuf>,
+}
+
 /// This is where the meat is - this function generates the `state` folder made
 /// of the CPU register as well as the memory dump.
-fn snapshot_with_kind_inner(kind: SnapshotKind, dbg: &DebugClient, args: String) -> Result<()> {
+fn snapshot_inner(dbg: &DebugClient, args: SnapshotArgs) -> Result<()> {
     // Let's make sure this is a live kernel, not a dump, etc..
     let is_live_kernel = matches!(
         dbg.debuggee_type()?,
@@ -319,12 +332,7 @@ fn snapshot_with_kind_inner(kind: SnapshotKind, dbg: &DebugClient, args: String)
     }
 
     // Build the state path.
-    let state_path = if args.is_empty() {
-        env::temp_dir()
-    } else {
-        PathBuf::from(args)
-    };
-
+    let state_path = args.state_path.unwrap_or(env::temp_dir());
     if !state_path.exists() {
         bail!("the directory {:?} doesn't exist", state_path);
     }
@@ -383,7 +391,7 @@ fn snapshot_with_kind_inner(kind: SnapshotKind, dbg: &DebugClient, args: String)
     // Generate the `mem.dmp`.
     dbg.exec(format!(
         ".dump /{} {:?}",
-        match kind {
+        match args.kind {
             // Create a dump with active kernel and user mode memory.
             SnapshotKind::ActiveKernel => "ka",
             // A Complete Memory Dump is the largest kernel-mode dump file. This file includes all
@@ -403,7 +411,11 @@ pub type RawIUnknown = *mut std::ffi::c_void;
 
 /// This is a wrapper function made to be able to display the error in case the
 /// inner function fails.
-fn snapshot_with_kind(raw_client: RawIUnknown, kind: SnapshotKind, args: PCSTR) -> HRESULT {
+fn wrap<P: Parser>(
+    raw_client: RawIUnknown,
+    args: PCSTR,
+    callback: impl FnOnce(&DebugClient, P) -> Result<()>,
+) -> HRESULT {
     // We do not own the `raw_client` interface  so we want to created a borrow. If
     // we don't, the object will get Release()'d when it gets dropped which will
     // lead to a use-after-free.
@@ -419,9 +431,19 @@ fn snapshot_with_kind(raw_client: RawIUnknown, kind: SnapshotKind, args: PCSTR)
         return E_ABORT;
     };
 
-    match snapshot_with_kind_inner(kind, &dbg, args) {
+    // Parse the arguments using `clap`. Currently splitting arguments by
+    // whitespaces. whitespace.
+    let args = match P::try_parse_from(args.split_whitespace()) {
+        Ok(a) => a,
+        Err(e) => {
+            let _ = dbg.logln(format!("{e}"));
+            return E_ABORT;
+        }
+    };
+
+    match callback(&dbg, args) {
         Err(e) => {
-            dbg.logln(format!("Ran into an error: {e:?}")).unwrap();
+            let _ = dbg.logln(format!("Ran into an error: {e:?}"));
 
             E_ABORT
         }
@@ -431,12 +453,7 @@ fn snapshot_with_kind(raw_client: RawIUnknown, kind: SnapshotKind, args: PCSTR)
 
 #[no_mangle]
 extern "C" fn snapshot(raw_client: RawIUnknown, args: PCSTR) -> HRESULT {
-    snapshot_with_kind(raw_client, SnapshotKind::Full, args)
-}
-
-#[no_mangle]
-extern "C" fn snapshot_active_kernel(raw_client: RawIUnknown, args: PCSTR) -> HRESULT {
-    snapshot_with_kind(raw_client, SnapshotKind::ActiveKernel, args)
+    wrap(raw_client, args, snapshot_inner)
 }
 
 /// The DebugExtensionInitialize callback function is called by the engine after