feat: supports yubikey init

src/serve.rs

@@ -5,17 +5,20 @@ use hyper::{Body, Client, Method, Request, Response, Server, StatusCode};
 use hyper::client::HttpConnector;
 use hyper::server::conn::AddrStream;
 use hyper::service::{make_service_fn, service_fn};
-use rust_util::{failure_and_exit, information, success, XResult};
+use rust_util::{failure_and_exit, information, success, warning, XResult};
 use rust_util::util_clap::{Command, CommandError};
 use serde_json::{json, Value};
+use tokio::runtime::Runtime;
 
 use crate::{db, jose, proc};
 use crate::do_response;
 use crate::serve_common::{self, GenericError, MemoryKey, Result};
 use crate::serve_encrypt_decrypt;
 use crate::serve_init;
+use crate::serve_init::InitRequest;
 use crate::serve_read_write;
 use crate::serve_status;
+use crate::yubikey_hmac;
 
 pub struct CommandImpl;
 
@@ -24,8 +27,10 @@ impl Command for CommandImpl {
 
     fn subcommand<'a>(&self) -> App<'a, 'a> {
         SubCommand::with_name(self.name()).about("Local mini KMS serve")
-            .arg(Arg::with_name("listen").long("listen").takes_value(true).default_value("127.0.0.1:5567").help("Listen"))
-            .arg(Arg::with_name("local-db").long("local-db").takes_value(true).default_value("local-mini-kms.db").help("Local db file"))
+            .arg(Arg::with_name("listen").long("listen").short("L").takes_value(true).default_value("127.0.0.1:5567").help("Listen"))
+            .arg(Arg::with_name("local-db").long("local-db").short("d").takes_value(true).default_value("local-mini-kms.db").help("Local db file"))
+            .arg(Arg::with_name("yubikey-challenge").long("yubikey-challenge").short("c").takes_value(true).help("Yubikey challenge"))
+            .arg(Arg::with_name("init-encrypted-master-key").long("init-encrypted-master-key").short("k").takes_value(true).help("Init encrypted mater key"))
     }
 
     fn run(&self, _arg_matches: &ArgMatches, sub_arg_matches: &ArgMatches) -> CommandError {
@@ -36,8 +41,10 @@ impl Command for CommandImpl {
             Err(e) => failure_and_exit!("Init server failed: {}", e),
         }
 
+        let rt = Runtime::new().expect("Create tokio runtime error");
+        init_with_yubikey_challenge(&rt, sub_arg_matches);
+
         let listen = sub_arg_matches.value_of("listen").expect("Get argument listen error");
-        let rt = tokio::runtime::Runtime::new().expect("Create tokio runtime error");
         rt.block_on(async {
             let addr = listen.parse().expect(&format!("Parse listen error: {}", listen));
             let client = Client::new();
@@ -156,3 +163,49 @@ Supports commands:
         Ok(Response::builder().body("Root Not Found\n".into())?)
     }
 }
+
+fn init_with_yubikey_challenge(rt: &Runtime, sub_arg_matches: &ArgMatches) {
+    let mut yubikey_challenge = sub_arg_matches.value_of("yubikey-challenge").map(ToString::to_string);
+    let init_encrypted_master_key = sub_arg_matches.value_of("init-encrypted-master-key");
+
+    if yubikey_challenge.is_none() {
+        yubikey_challenge = rpassword::prompt_password("Yubikey challenge: ").ok();
+    }
+
+    let (challenge_key, init_encrypted_master_key) = match (yubikey_challenge, init_encrypted_master_key) {
+        (Some(yubikey_challenge), Some(init_encrypted_master_key)) => {
+            match yubikey_hmac::yubikey_challenge_as_32_bytes(yubikey_challenge.as_bytes()) {
+                Err(e) => {
+                    warning!("Yubikey challenge failed: {}", e);
+                    return;
+                }
+                Ok(challenge_key) => (challenge_key, init_encrypted_master_key),
+            }
+        }
+        (Some(_), None) | (None, Some(_)) => {
+            warning!("Arguments yubikey-challenge and init-encrypted-master-key should both assigned.");
+            return;
+        }
+        _ => return,
+    };
+
+    match jose::deserialize_jwe_aes(&init_encrypted_master_key, &challenge_key) {
+        Err(e) => warning!("Yubikey seal master key failed: {}", e),
+        Ok((key, _)) => {
+            success!("Yubikey un-seal master key success");
+            let init_master_key_result = rt.block_on(async {
+                serve_init::inner_init_request(InitRequest {
+                    yubikey_challenge: None,
+                    clear_master_key_hex: Some(hex::encode(&key)),
+                    clear_master_key_base64: None,
+                    encrypted_master_key: None,
+                }).await
+            });
+            match init_master_key_result {
+                Err(e) => warning!("Init master key failed: {}", e),
+                Ok((StatusCode::OK, _)) => success!("Init master key success"),
+                Ok((_, response)) => warning!("Init master failed: {}", response),
+            }
+        }
+    }
+}