2 files changed, 48 insertions, 1 deletions

diff --git a/nitrocli/src/args.rs b/nitrocli/src/args.rs
index 1a54c43..9982d0f 100644
--- a/nitrocli/src/args.rs
+++ b/nitrocli/src/args.rs
@@ -327,12 +327,14 @@ impl str::FromStr for PinCommand {
 #[derive(Debug)]
 enum PwsCommand {
   Get,
+  Set,
 }
 
 impl PwsCommand {
   fn execute(&self, args: Vec<String>) -> Result<()> {
     match *self {
       PwsCommand::Get => pws_get(args),
+      PwsCommand::Set => pws_set(args),
     }
   }
 }
@@ -344,6 +346,7 @@ impl fmt::Display for PwsCommand {
       "{}",
       match *self {
         PwsCommand::Get => "get",
+        PwsCommand::Set => "set",
       }
     )
   }
@@ -355,6 +358,7 @@ impl str::FromStr for PwsCommand {
   fn from_str(s: &str) -> result::Result<Self, Self::Err> {
     match s {
       "get" => Ok(PwsCommand::Get),
+      "set" => Ok(PwsCommand::Set),
       _ => Err(()),
     }
   }
@@ -796,7 +800,7 @@ fn pws(args: Vec<String>) -> Result<()> {
   let _ = parser.refer(&mut subcommand).required().add_argument(
     "subcommand",
     argparse::Store,
-    "The subcommand to execute (get)",
+    "The subcommand to execute (get|set)",
   );
   let _ = parser.refer(&mut subargs).add_argument(
     "arguments",
@@ -851,6 +855,40 @@ fn pws_get(args: Vec<String>) -> Result<()> {
   commands::pws_get(slot, name, login, password, quiet)
 }
 
+/// Set a slot of the password safe on the Nitrokey.
+fn pws_set(args: Vec<String>) -> Result<()> {
+  let mut slot: u8 = 0;
+  let mut name = String::new();
+  let mut login = String::new();
+  let mut password = String::new();
+  let mut parser = argparse::ArgumentParser::new();
+  parser.set_description("Writes a password safe slot");
+  let _ = parser.refer(&mut slot).required().add_argument(
+    "slot",
+    argparse::Store,
+    "The PWS slot to read",
+  );
+  let _ = parser.refer(&mut name).required().add_argument(
+    "name",
+    argparse::Store,
+    "The name to store on the slot",
+  );
+  let _ = parser.refer(&mut login).required().add_argument(
+    "login",
+    argparse::Store,
+    "The login to store on the slot",
+  );
+  let _ = parser.refer(&mut password).required().add_argument(
+    "password",
+    argparse::Store,
+    "The password to store on the slot",
+  );
+  parse(&parser, args)?;
+  drop(parser);
+
+  commands::pws_set(slot, &name, &login, &password)
+}
+
 /// Parse the command-line arguments and return the selected command and
 /// the remaining arguments for the command.
 fn parse_arguments(args: Vec<String>) -> Result<(Command, Vec<String>)> {
diff --git a/nitrocli/src/commands.rs b/nitrocli/src/commands.rs
index 125b344..e1c88c8 100644
--- a/nitrocli/src/commands.rs
+++ b/nitrocli/src/commands.rs
@@ -584,6 +584,15 @@ pub fn pws_get(
   Ok(())
 }
 
+/// Write a PWS slot.
+pub fn pws_set(slot: u8, name: &str, login: &str, password: &str) -> Result<()> {
+  let device = get_device()?;
+  let pws = get_password_safe(&device)?;
+  pws
+    .write_slot(slot, name, login, password)
+    .map_err(|err| get_error("Could not write PWS slot", &err))
+}
+
 #[cfg(test)]
 mod tests {
   use super::*;