feat: Support multi-format integer parsing for CLI arguments

README.md

@@ -131,7 +131,7 @@ snpguest fetch <SUBCOMMAND>
     | `-r, --report $ATT_REPORT_PATH` | Specifies the attestation report to detect the host processor model (conflict with `$PROCESSOR_MODEL`) | — |
     | `-e, --endorser $ENDORSER` | Specifies the endorser type, possible values: "vcek", "vlek". | "vcek" |
 
-    Example
+    **Example**
     ```bash
     snpguest fetch ca der ./certs-kds milan -e vlek
     ```
@@ -163,7 +163,7 @@ snpguest fetch <SUBCOMMAND>
 
 ### 5. `key` 
 
-Creates the derived key based on input parameters and stores it. `$KEY_PATH` is the path to store the derived key. `$ROOT_KEY_SELECT` is the root key from which to derive the key (either "vcek" or "vmrk"). The `--guest_field_select` option specifies which Guest Field Select bits to enable as a binary string of length 6 or 7. Each of the bits from *right to left* correspond to Guest Policy, Image ID, Family ID, Measurement, SVN and TCB Version, Launch Mitigation Vector, respectively. For each bit, 0 denotes off, and 1 denotes on. The `--guest_svn` option specifies the guest SVN to mix into the key, the `--tcb_version` option specifies the TCB version to mix into the derived key, and the `--launch_mit_vector` option specifies the launch mitigation vector value to mix into the derived key. The `--vmpl` option specifies the VMPL level the Guest is running on and defaults to 1.
+Creates the derived key based on input parameters and stores it. `$KEY_PATH` is the path to store the derived key. `$ROOT_KEY_SELECT` is the root key from which to derive the key (either "vcek" or "vmrk"). The `--guest_field_select` option specifies which Guest Field Select bits to enable as a 64-bit integer. Only the least-significant 6 bits (Message Version 1) or 7 bits (Message Version 2) are currently defined. Each of the bits from *right to left* correspond to Guest Policy, Image ID, Family ID, Measurement, SVN and TCB Version, Launch Mitigation Vector, respectively. For each bit, 0 denotes off, and 1 denotes on. The `--guest_svn` option specifies the guest SVN to mix into the key, the `--tcb_version` option specifies the TCB version to mix into the derived key, and the `--launch_mit_vector` option specifies the launch mitigation vector value to mix into the derived key. The `--vmpl` option specifies the VMPL level the Guest is running on and defaults to 1.
 
 **Usage**
 ```bash
@@ -176,25 +176,30 @@ snpguest key $KEY_PATH $ROOT_KEY_SELECT [-v, --vmpl] [-g, --guest_field_select]
 | `$KEY_PATH` | The path to store the derived key. | required |
 | `$ROOT_KEY_SELECT` | is the root key from which to derive the key (either "vcek" or "vmrk"). | required |
 | `-v, --vmpl $VMPL` | option specifies the VMPL level the Guest is running on. | 1 |
-| `-g, --guest_field_select $GFS` | option specifies which Guest Field Select bits to enable as a binary string (length 6 or 7). For each bit, 0 denotes off, and 1 denotes on. | all 0s |
-| `-s, --guest_svn $GSVN` | option specifies the guest SVN to mix into the key. | — |
-| `-t, --tcb_version $TCBV` | option specifies the TCB version to mix into the derived key. | — |
-| `-l, --launch_mit_vector $LMV` | option specifies the launch mitigation vector value to mix into the derived key (only available for `MSG_KEY_REQ` version ≥ 2 ). | — |
+| `-g, --guest_field_select $GFS` | option specifies which Guest Field Select bits to enable as a 64-bit integer (decimal, prefixed hex or prefixed bin). | 0 |
+| `-s, --guest_svn $GSVN` | option specifies the guest SVN to mix into the key (decimal, prefixed hex or prefixed bin). | 0 |
+| `-t, --tcb_version $TCBV` | option specifies the TCB version to mix into the derived key (decimal, prefixed hex or prefixed bin). | 0 |
+| `-l, --launch_mit_vector $LMV` | option specifies the launch mitigation vector value to mix into the derived key (decimal, prefixed hex or prefixed bin). Only available for `MSG_KEY_REQ` message version ≥ 2. | — |
 
 **Guest Field Select**
 
 | Bit      | Field             | Note |
 | :--      | :--               | :--  |
-| 0 (LSB)  | Guest Policy      |      |
-| 1        | Image ID          |      |
-| 2        | Family ID         |      |
-| 3        | Measurement       |      |
-| 4        | SVN               |      |
+| 63:7     | (Reserved)        | Currently not supported. |
+| 6 (MSB)  | Launch MIT Vector | Set to 0 if not specified; supported for `MSG_KEY_REQ` message version ≥ 2. |
 | 5        | TCB Version       |      |
-| 6 (MSB)  | Launch MIT Vector | Set to 0 if not specified; supported for `MSG_KEY_REQ` version ≥ 2. |
-| (7–63)   | (Reserved)        | Currently not supported. |
+| 4        | SVN               |      |
+| 3        | Measurement       |      |
+| 2        | Family ID         |      |
+| 1        | Image ID          |      |
+| 0 (LSB)  | Guest Policy      |      |
+
+For example, all of
+- `--guest_field_select 49`
+- `--guest_field_select 0x31`
+- `--guest_field_select 0b110001`
 
-For example, `--guest_field_select 110001` denotes 
+denote the following specification:
 ```
 Guest Policy:On (1), 
 Image ID:Off (0), 
@@ -208,7 +213,7 @@ Launch MIT Vector: Off (none).
 **Example**
 ```bash
 # Creating and storing a derived key
-snpguest key derived-key.bin vcek --guest_field_select 110001 --guest_svn 2 --tcb_version 1 --vmpl 3
+snpguest key derived-key.bin vcek --guest_field_select 0b110001 --guest_svn 2 --tcb_version 1 --vmpl 3
 ```
 
 ### 6. `report` 

src/clparser/mod.rs

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// This file contains custom u32/u64 parsers for supporting decimal, hexadecimal and binary formats.
+
+use std::num::ParseIntError;
+
+pub fn parse_int_auto_radix<T>(s: &str) -> Result<T, ParseIntError>
+where
+    T: FromStrRadix,
+{
+    if let Some(hex) = s.strip_prefix("0x") {
+        T::from_str_radix(hex, 16)
+    } else if let Some(bin) = s.strip_prefix("0b") {
+        T::from_str_radix(bin, 2)
+    } else {
+        T::from_str_radix(s, 10)
+    }
+}
+
+pub trait FromStrRadix: Sized {
+    fn from_str_radix(src: &str, radix: u32) -> Result<Self, ParseIntError>;
+}
+
+impl FromStrRadix for u32 {
+    fn from_str_radix(src: &str, radix: u32) -> Result<Self, ParseIntError> {
+        u32::from_str_radix(src, radix)
+    }
+}
+
+impl FromStrRadix for u64 {
+    fn from_str_radix(src: &str, radix: u32) -> Result<Self, ParseIntError> {
+        u64::from_str_radix(src, radix)
+    }
+}

src/key.rs

@@ -18,77 +18,55 @@ pub struct KeyArgs {
 
     /// Specify an integer VMPL level between 0 and 3 that the Guest is running on.
     #[arg(short, long, value_name = "vmpl", default_value = "1")]
-    pub vmpl: Option<u32>,
+    pub vmpl: u32,
 
-    /// Specify which Guest Field Select bits to enable. It is a 6 or 7 digit binary string. For each bit, 0 denotes off and 1 denotes on.
-    /// The least significant (rightmost) bit is Guest Policy followed by Image ID, Family ID, Measurement, SVN, TCB Version, and Launch Mitigation Vector which is the most significant (leftmost) bit.
-    #[arg(short, long = "guest_field_select", value_name = "#######")]
-    pub gfs: Option<String>,
+    /// Specify which Guest Field Select bits to enable. It is 64-bit wide but only the least-significant 6 bits (Message Version 1) or 7 bits (Message Version 2) are currently defined; all higher bits must be zero. The bits (LSB->MSB) are: 0 = Guest Policy, 1 = Image ID, 2 = Family ID, 3 = Measurement, 4 = SVN, 5 = TCB Version, 6 = Launch Mitigation Vector (only available for Message Version 2). Accepts an integer in decimal (e.g. `63`), prefixed hex (e.g. `0x3f`) or prefixed bin (e.g. `0b111111`).
+    #[arg(short, long = "guest_field_select", value_name = "gfs", value_parser = clparser::parse_int_auto_radix::<u64>, default_value = "0")]
+    pub gfs: u64,
 
-    /// Specify the guest SVN to mix into the key. Must not exceed the guest SVN provided at launch in the ID block.
-    #[arg(short = 's', long = "guest_svn")]
-    pub gsvn: Option<u32>,
+    /// Specify the guest SVN to mix into the key. Must not exceed the guest SVN provided at launch in the ID block. Accepts an integer in decimal, prefixed hex or prefixed bin.
+    #[arg(short = 's', long = "guest_svn", value_name = "gsvn", value_parser = clparser::parse_int_auto_radix::<u32>, default_value = "0")]
+    pub gsvn: u32,
 
-    /// Specify the TCB version to mix into the derived key. Must not exceed CommittedTcb.
-    #[arg(short, long = "tcb_version")]
-    pub tcbv: Option<u64>,
+    /// Specify the TCB version to mix into the derived key. Must not exceed CommittedTcb. Accepts an integer in decimal, prefixed hex or prefixed bin.
+    #[arg(short, long = "tcb_version", value_name = "tcbv", value_parser = clparser::parse_int_auto_radix::<u64>, default_value = "0")]
+    pub tcbv: u64,
 
-    /// Specify the launch mitigation vector to mix into the derived key.
-    #[arg(short, long = "launch_mit_vector")]
+    /// Specify the launch mitigation vector to mix into the derived key (only available for Message Version 2). Accepts an integer in decimal, hexadecimal or binary string.
+    #[arg(short, long = "launch_mit_vector", value_name = "lmv", value_parser = clparser::parse_int_auto_radix::<u64>)]
     pub lmv: Option<u64>,
 }
 
 pub fn get_derived_key(args: KeyArgs) -> Result<()> {
+    // Validate arguments
     let root_key_select = match args.root_key_select.as_str() {
         "vcek" => false,
         "vmrk" => true,
         _ => return Err(anyhow::anyhow!("Invalid input. Enter either vcek or vmrk")),
     };
 
-    let vmpl = match args.vmpl {
-        Some(level) => {
-            if level <= 3 {
-                level
-            } else {
-                return Err(anyhow::anyhow!("Invalid Virtual Machine Privilege Level."));
-            }
-        }
-        None => 1,
-    };
-
-    let gfs = match args.gfs {
-        Some(gfs) => {
-            if gfs.len() != 6 && gfs.len() != 7 {
-                return Err(anyhow::anyhow!(
-                    "Invalid Guest Field Select option. Must be 6 or 7 digits."
-                ));
-            }
-            if gfs.chars().any(|c| c != '0' && c != '1') {
-                return Err(anyhow::anyhow!(
-                    "Invalid Guest Field Select option. Must be a binary string."
-                ));
-            }
-            let value: u64 = u64::from_str_radix(gfs.as_str(), 2).unwrap();
-            value
-        }
-        None => 0,
-    };
+    if args.vmpl > 3 {
+        return Err(anyhow::anyhow!(
+            "Invalid Virtual Machine Privilege Level. Must betwee"
+        ));
+    }
 
-    let gsvn: u32 = args.gsvn.unwrap_or(0);
+    if args.gfs > 0b1111111 {
+        return Err(anyhow::anyhow!("Invalid Guest Field Select option."));
+    }
 
-    let tcbv: u64 = args.tcbv.unwrap_or(0);
+    // Switch message version of MSG_KEY_REQ
+    let msg_ver = if args.lmv.is_some() { 2 } else { 1 };
 
+    // Request derived key
     let request = DerivedKey::new(
         root_key_select,
-        GuestFieldSelect(gfs),
-        vmpl,
-        gsvn,
-        tcbv,
+        GuestFieldSelect(args.gfs),
+        args.vmpl,
+        args.gsvn,
+        args.tcbv,
         args.lmv,
     );
-
-    let msg_ver = if args.lmv.is_some() { 2 } else { 1 };
-
     let mut sev_fw = Firmware::open().context("failed to open SEV firmware device.")?;
     let derived_key: [u8; 32] = sev_fw
         .get_derived_key(Some(msg_ver), request)

src/main.rs

@@ -10,6 +10,8 @@ mod preattestation;
 mod report;
 mod verify;
 
+mod clparser;
+
 #[cfg(feature = "hyperv")]
 mod hyperv;