feat: works2

src/cmd_signjwt.rs

@@ -1,7 +1,9 @@
+use std::borrow::Cow;
 use std::collections::BTreeMap;
 
 use clap::{App, Arg, ArgMatches, SubCommand};
 use jwt::{AlgorithmType, Header, ToBase64};
+use jwt::header::HeaderType;
 use rust_util::{util_msg, XResult};
 use rust_util::util_clap::{Command, CommandError};
 use yubikey::piv::{AlgorithmId, sign_data};
@@ -20,6 +22,9 @@ impl Command for CommandImpl {
         SubCommand::with_name(self.name()).about("Sign JWT subcommand")
             .arg(Arg::with_name("pin").short("p").long("pin").takes_value(true).help("PIV card user pin"))
             .arg(Arg::with_name("slot").short("s").long("slot").takes_value(true).help("PIV slot, e.g. 82, 83 ... 95, 9a, 9c, 9d, 9e"))
+            .arg(Arg::with_name("key-id").short("K").long("key-id").takes_value(true).help("Header key ID"))
+            .arg(Arg::with_name("claims").short("C").long("claims").takes_value(true).multiple(true).help("Claims, key:value"))
+            .arg(Arg::with_name("payload").short("P").long("payload").takes_value(true).help("Claims in JSON"))
             .arg(Arg::with_name("json").long("json").help("JSON output"))
     }
 
@@ -33,15 +38,29 @@ impl Command for CommandImpl {
         let slot = opt_value_result!(
             sub_arg_matches.value_of("slot"), "--slot must assigned, e.g. 82, 83 ... 95, 9a, 9c, 9d, 9e");
 
-        // TODO custom header
+        let key_id = sub_arg_matches.value_of("key-id");
+        let claims = opt_value_result!(sub_arg_matches.values_of("claims"), "Claims is required.");
+        let payload = sub_arg_matches.value_of("payload");
+
         let header = Header {
+            key_id: key_id.map(ToString::to_string),
+            type_: Some(HeaderType::JsonWebToken),
             ..Default::default()
         };
-        let mut claims = BTreeMap::new();
-        claims.insert("sub".to_string(), "someone".to_string());
-
-        let token_string = sign_jwt(slot, &pin_opt, header, &claims)?;
+        let mut jwt_claims = BTreeMap::new();
+        if payload.is_none() {
+            for claim in claims {
+                match split_claim(claim) {
+                    None => { warning!("Claim '{}' do not contains ':'", claim); }
+                    Some((k, v)) => { jwt_claims.insert(k, v); }
+                }
+            }
+            if !jwt_claims.contains_key("sub") {
+                return simple_error!("Claim sub is not assigned.");
+            }
+        }
 
+        let token_string = sign_jwt(slot, &pin_opt, header, &payload, &jwt_claims)?;
         success!("Singed JWT: {}", token_string);
         if json_output {
             json.insert("token", token_string.clone());
@@ -55,7 +74,7 @@ impl Command for CommandImpl {
 }
 
 
-fn sign_jwt(slot: &str, pin_opt: &Option<&str>, mut header: Header, claims: &BTreeMap<String, String>) -> XResult<String> {
+fn sign_jwt(slot: &str, pin_opt: &Option<&str>, mut header: Header, payload: &Option<&str>, claims: &BTreeMap<String, String>) -> XResult<String> {
     let mut yk = opt_result!(YubiKey::open(), "Find YubiKey failed: {}");
     let slot_id = opt_result!(pivutil::get_slot_id(slot), "Get slot id failed: {}");
 
@@ -71,7 +90,10 @@ fn sign_jwt(slot: &str, pin_opt: &Option<&str>, mut header: Header, claims: &BTr
     debugging!("Claims: {:?}", claims);
 
     let header = opt_result!(header.to_base64(), "Header to base64 failed: {}");
-    let claims = opt_result!(claims.to_base64(), "Claims to base64 failed: {}");
+    let claims = match payload {
+        Some(payload) => Cow::Owned(payload.to_string()),
+        None => opt_result!(claims.to_base64(), "Claims to base64 failed: {}"),
+    };
 
     let mut tobe_signed = vec![];
     tobe_signed.extend_from_slice(header.as_bytes());
@@ -85,4 +107,24 @@ fn sign_jwt(slot: &str, pin_opt: &Option<&str>, mut header: Header, claims: &BTr
     let signature = util::base64_encode_url_safe_no_pad(signed_data);
 
     Ok([&*header, &*claims, &signature].join(SEPARATOR))
-}
+}
+
+fn split_claim(claim: &str) -> Option<(String, String)> {
+    let mut k = String::new();
+    let mut v = String::new();
+
+    let mut is_k = true;
+    for c in claim.chars() {
+        if is_k {
+            if c == ':' {
+                is_k = false;
+            } else {
+                k.push(c);
+            }
+        } else {
+            v.push(c);
+        }
+    }
+
+    iff!(is_k, None, Some((k, v)))
+}