orb-slot-ctrl: Prettify handling of arguments from user

orb-slot-ctrl/src/lib.rs

@@ -6,6 +6,7 @@
 use std::{
     fmt, io,
     path::{Path, PathBuf},
+    str::FromStr,
 };
 
 mod efivar;
@@ -46,8 +47,15 @@ pub enum Error {
     RemoveEfiVar { path: PathBuf, source: io::Error },
     #[error("failed reading efivar, invalid data length")]
     InvalidEfiVarLen,
+    #[error("invalid slot provided {slot}. Use one of the available slot aliases: \n{help_message}")]
+    InvalidSlotProvided { slot: String, help_message: String },
     #[error("invalid slot configuration")]
     InvalidSlotData,
+    #[error("invalid status provided {status}. Use one of the available status variant aliases: \n{help_message}")]
+    InvalidRootFsStatusProvided {
+        status: String,
+        help_message: String,
+    },
     #[error("invalid rootfs status")]
     InvalidRootFsStatusData,
     #[error("invalid retry counter({counter}), exceeding the maximum ({max})")]
@@ -101,7 +109,7 @@ impl Error {
 }
 
 /// Representation of the slot.
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 #[repr(u8)]
 pub enum Slot {
     /// The Slot A is represented as 0.
@@ -110,6 +118,49 @@ pub enum Slot {
     B = SLOT_B,
 }
 
+impl Slot {
+    fn variants() -> Vec<(Slot, &'static str, Vec<&'static str>)> {
+        vec![
+            (Self::A, "A", vec!["a", "0"]),
+            (Self::B, "B", vec!["b", "1"]),
+        ]
+    }
+
+    /// Retrieves a help message listing each slot variant along with its corresponding aliases.
+    #[must_use]
+    pub fn help_message() -> String {
+        let variants = Self::variants();
+        let message_parts: Vec<String> = variants
+            .iter()
+            .map(|(_, desc, aliases)| format!("{}({})", desc, aliases.join(", ")))
+            .collect();
+
+        let message = message_parts.join(";\n");
+
+        message.to_string()
+    }
+}
+
+impl FromStr for Slot {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Self::variants()
+            .iter()
+            .find_map(|(variant, _, aliases)| {
+                if aliases.contains(&s.to_lowercase().as_str()) {
+                    Some(*variant)
+                } else {
+                    None
+                }
+            })
+            .ok_or_else(|| Error::InvalidSlotProvided {
+                slot: s.to_string(),
+                help_message: Self::help_message(),
+            })
+    }
+}
+
 /// Format slot as lowercase to match Nvidia standard in file system.
 impl fmt::Display for Slot {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -121,7 +172,7 @@ impl fmt::Display for Slot {
 }
 
 /// Representation of the rootfs status.
-#[derive(Debug)]
+#[derive(Clone, Copy, Debug)]
 #[repr(u8)]
 pub enum RootFsStatus {
     /// Default status of the rootfs.
@@ -135,6 +186,37 @@ pub enum RootFsStatus {
 }
 
 impl RootFsStatus {
+    fn variants() -> Vec<(RootFsStatus, &'static str, Vec<&'static str>)> {
+        vec![
+            (Self::Normal, "Normal", vec!["normal", "0"]),
+            (
+                Self::UpdateInProcess,
+                "Update in Process",
+                vec!["updateinprocess", "updinprocess", "1"],
+            ),
+            (
+                Self::UpdateDone,
+                "Update Done",
+                vec!["updatedone", "upddone", "2"],
+            ),
+            (Self::Unbootable, "Unbootable", vec!["unbootable", "3"]),
+        ]
+    }
+
+    /// Retrieves a help message listing each status variant along with its corresponding aliases.
+    #[must_use]
+    pub fn help_message() -> String {
+        let variants = Self::variants();
+        let message_parts: Vec<String> = variants
+            .iter()
+            .map(|(_, desc, aliases)| format!("{}({})", desc, aliases.join(", ")))
+            .collect();
+
+        let message = message_parts.join(";\n");
+
+        message.to_string()
+    }
+
     /// Checks if current status is `RootFsStats::Normal`.
     #[must_use]
     pub fn is_normal(self) -> bool {
@@ -174,6 +256,26 @@ impl TryFrom<u8> for RootFsStatus {
     }
 }
 
+impl FromStr for RootFsStatus {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Self::variants()
+            .iter()
+            .find_map(|(variant, _, aliases)| {
+                if aliases.contains(&s.to_lowercase().as_str()) {
+                    Some(*variant)
+                } else {
+                    None
+                }
+            })
+            .ok_or_else(|| Error::InvalidRootFsStatusProvided {
+                status: s.to_string(),
+                help_message: Self::help_message(),
+            })
+    }
+}
+
 /// Get the current active slot.
 pub fn get_current_slot() -> Result<Slot, Error> {
     match efivar::bootchain::get_current_boot_slot()? {

orb-slot-ctrl/src/main.rs

@@ -1,4 +1,5 @@
 use clap::{Parser, Subcommand};
+use orb_slot_ctrl::{RootFsStatus, Slot};
 use std::{env, process::exit};
 
 use orb_build_info::{make_build_info, BuildInfo};
@@ -26,7 +27,13 @@ enum Commands {
     GetNextSlot,
     /// Set slot for the next boot.
     #[command(name = "set", short_flag = 's')]
-    SetNextSlot { slot: String },
+    SetNextSlot {
+        #[arg(
+            value_parser = clap::value_parser!(Slot),
+            help = Slot::help_message()
+        )]
+        slot: Slot,
+    },
     /// Rootfs status controls.
     Status {
         /// Control the inactive slot instead of the active.
@@ -47,7 +54,13 @@ enum StatusCommands {
     GetRootfsStatus,
     /// Set the rootfs status.
     #[command(name = "set", short_flag = 's')]
-    SetRootfsStatus { status: String },
+    SetRootfsStatus {
+        #[arg(
+            value_parser = clap::value_parser!(RootFsStatus),
+            help = RootFsStatus::help_message()
+        )]
+        status: RootFsStatus,
+    },
     /// Get the retry counter.
     #[command(name = "retries", short_flag = 'c')]
     GetRetryCounter,
@@ -82,120 +95,68 @@ fn main() -> eyre::Result<()> {
             println!("{}", orb_slot_ctrl::get_next_boot_slot()?);
         }
         Commands::SetNextSlot { slot } => {
-            let slot = match slot.as_str() {
-                // Slot A alias.
-                "A" => orb_slot_ctrl::Slot::A,
-                "a" => orb_slot_ctrl::Slot::A,
-                "0" => orb_slot_ctrl::Slot::A,
-                // Slot B alias.
-                "B" => orb_slot_ctrl::Slot::B,
-                "b" => orb_slot_ctrl::Slot::B,
-                "1" => orb_slot_ctrl::Slot::B,
-                _ => {
-                    println!(
-                        "Invalid slot provided, please use either A/a/0 or B/b/1."
-                    );
-                    exit(1)
-                }
-            };
             if let Err(e) = orb_slot_ctrl::set_next_boot_slot(slot) {
                 check_running_as_root(e);
             };
         }
-        Commands::Status { inactive, subcmd } => {
-            match subcmd {
-                StatusCommands::GetRootfsStatus => {
-                    if inactive {
-                        println!(
-                            "{:?}",
-                            orb_slot_ctrl::get_rootfs_status(
-                                orb_slot_ctrl::get_inactive_slot()?
-                            )?
-                        );
-                    } else {
-                        println!("{:?}", orb_slot_ctrl::get_current_rootfs_status()?);
-                    }
+        Commands::Status { inactive, subcmd } => match subcmd {
+            StatusCommands::GetRootfsStatus => {
+                if inactive {
+                    println!(
+                        "{:?}",
+                        orb_slot_ctrl::get_rootfs_status(
+                            orb_slot_ctrl::get_inactive_slot()?
+                        )?
+                    );
+                } else {
+                    println!("{:?}", orb_slot_ctrl::get_current_rootfs_status()?);
                 }
-                StatusCommands::SetRootfsStatus { status } => {
-                    let status = match status.as_str() {
-                        // Status Normal alias.
-                        "Normal" => orb_slot_ctrl::RootFsStatus::Normal,
-                        "normal" => orb_slot_ctrl::RootFsStatus::Normal,
-                        "0" => orb_slot_ctrl::RootFsStatus::Normal,
-                        // Status UpdateInProcess alias.
-                        "UpdateInProcess" => {
-                            orb_slot_ctrl::RootFsStatus::UpdateInProcess
-                        }
-                        "updateinprocess" => {
-                            orb_slot_ctrl::RootFsStatus::UpdateInProcess
-                        }
-                        "updinprocess" => orb_slot_ctrl::RootFsStatus::UpdateInProcess,
-                        "1" => orb_slot_ctrl::RootFsStatus::UpdateInProcess,
-                        // Status UpdateDone alias.
-                        "UpdateDone" => orb_slot_ctrl::RootFsStatus::UpdateDone,
-                        "updatedone" => orb_slot_ctrl::RootFsStatus::UpdateDone,
-                        "upddone" => orb_slot_ctrl::RootFsStatus::UpdateDone,
-                        "2" => orb_slot_ctrl::RootFsStatus::UpdateDone,
-                        // Status Unbootable alias.
-                        "Unbootable" => orb_slot_ctrl::RootFsStatus::Unbootable,
-                        "unbootable" => orb_slot_ctrl::RootFsStatus::Unbootable,
-                        "3" => orb_slot_ctrl::RootFsStatus::Unbootable,
-                        _ => {
-                            println!("Invalid status provided. For a full list of available rootfs status run:");
-                            println!("slot-ctrl status --list");
-                            exit(1)
-                        }
-                    };
-                    if inactive {
-                        if let Err(e) = orb_slot_ctrl::set_rootfs_status(
-                            status,
-                            orb_slot_ctrl::get_inactive_slot()?,
-                        ) {
-                            check_running_as_root(e);
-                        }
-                    } else if let Err(e) =
-                        orb_slot_ctrl::set_current_rootfs_status(status)
-                    {
+            }
+            StatusCommands::SetRootfsStatus { status } => {
+                if inactive {
+                    if let Err(e) = orb_slot_ctrl::set_rootfs_status(
+                        status,
+                        orb_slot_ctrl::get_inactive_slot()?,
+                    ) {
                         check_running_as_root(e);
                     }
+                } else if let Err(e) = orb_slot_ctrl::set_current_rootfs_status(status)
+                {
+                    check_running_as_root(e);
                 }
-                StatusCommands::GetRetryCounter => {
-                    if inactive {
-                        println!(
-                            "{}",
-                            orb_slot_ctrl::get_retry_count(
-                                orb_slot_ctrl::get_inactive_slot()?
-                            )?
-                        );
-                    } else {
-                        println!("{}", orb_slot_ctrl::get_current_retry_count()?);
-                    }
-                }
-                StatusCommands::GetMaxRetryCounter => {
-                    println!("{}", orb_slot_ctrl::get_max_retry_count()?);
+            }
+            StatusCommands::GetRetryCounter => {
+                if inactive {
+                    println!(
+                        "{}",
+                        orb_slot_ctrl::get_retry_count(
+                            orb_slot_ctrl::get_inactive_slot()?
+                        )?
+                    );
+                } else {
+                    println!("{}", orb_slot_ctrl::get_current_retry_count()?);
                 }
-                StatusCommands::ResetRetryCounter => {
-                    if inactive {
-                        if let Err(e) = orb_slot_ctrl::reset_retry_count_to_max(
-                            orb_slot_ctrl::get_inactive_slot()?,
-                        ) {
-                            check_running_as_root(e)
-                        }
-                    } else if let Err(e) =
-                        orb_slot_ctrl::reset_current_retry_count_to_max()
-                    {
+            }
+            StatusCommands::GetMaxRetryCounter => {
+                println!("{}", orb_slot_ctrl::get_max_retry_count()?);
+            }
+            StatusCommands::ResetRetryCounter => {
+                if inactive {
+                    if let Err(e) = orb_slot_ctrl::reset_retry_count_to_max(
+                        orb_slot_ctrl::get_inactive_slot()?,
+                    ) {
                         check_running_as_root(e)
                     }
-                }
-                StatusCommands::ListStatusVariants => {
-                    println!("Available Rootfs status variants with their aliases):");
-                    println!("  Normal (normal, 0)");
-                    println!("  UpdateInProcess (updateinprocess, updinprocess, 1)");
-                    println!("  UpdateDone (updatedone, upddone, 2)");
-                    println!("  Unbootable (unbootable, 3)");
+                } else if let Err(e) = orb_slot_ctrl::reset_current_retry_count_to_max()
+                {
+                    check_running_as_root(e)
                 }
             }
-        }
+            StatusCommands::ListStatusVariants => {
+                println!("Available Rootfs status variants with their aliases:");
+                println!("{}", RootFsStatus::help_message());
+            }
+        },
         Commands::GitDescribe => {
             println!("{}", BUILD_INFO.git.describe);
         }