feat: optimize code

This commit is contained in:
2023-11-02 01:10:07 +08:00
parent c43c7011e7
commit 66fdde9b72
4 changed files with 88 additions and 105 deletions

View File

@@ -1,10 +1,8 @@
use std::collections::BTreeMap;
use clap::{App, Arg, ArgMatches, SubCommand};
use rust_util::{util_msg, XResult};
use rust_util::util_clap::{Command, CommandError};
use rust_util::util_msg;
use yubico_manager::hmacmode::HmacKey;
use yubico_manager::sec::hmac_sha1;
use crate::hmacutil;
pub struct CommandImpl;
@@ -29,65 +27,24 @@ impl Command for CommandImpl {
if json_output { util_msg::set_logger_std_out(false); }
let variable = sub_arg_matches.is_present("variable");
let sha1_output = sub_arg_matches.is_present("sha1");
let sha256_output = sub_arg_matches.is_present("sha256");
let sha384_output = sub_arg_matches.is_present("sha384");
let sha512_output = sub_arg_matches.is_present("sha512");
let secret_bytes: Vec<u8> = if let Some(secret) = sub_arg_matches.value_of("secret") {
opt_result!(hex::decode(secret), "Decode secret hex: {}, failed: {}", secret)
} else {
return simple_error!("Secret must assigned");
};
let challenge_bytes: Vec<u8> = if let Some(challenge) = sub_arg_matches.value_of("challenge") {
challenge.as_bytes().to_vec()
} else if let Some(challenge_hex) = sub_arg_matches.value_of("challenge-hex") {
opt_result!(hex::decode(challenge_hex), "Decode challenge hex: {}, failed: {}", challenge_hex)
} else {
return simple_error!("Challenge must assigned");
};
let secret_bytes = get_secret_bytes(sub_arg_matches)?;
let challenge_bytes = hmacutil::get_challenge_bytes(sub_arg_matches)?;
if secret_bytes.len() != 20 {
return simple_error!("Secret length must be 20, actual is: {}", secret_bytes.len());
}
// Challenge can not be greater than 64 bytes
if challenge_bytes.len() > 64 {
return simple_error!("Challenge bytes is: {}, more than 64", challenge_bytes.len());
}
let hmac_key = HmacKey::from_slice(&secret_bytes);
let mut challenge = [0; 64];
if variable && challenge_bytes.last() == Some(&0) {
challenge = [0xff; 64];
}
(&mut challenge[..challenge_bytes.len()]).copy_from_slice(&challenge_bytes);
let hmac_result = hmac_sha1(&hmac_key, &challenge);
let v: &[u8] = &hmac_result;
let hex_string = hex::encode(v);
let hex_sha1 = iff!(sha1_output, Some(crate::digest::sha1_bytes(v)), None);
let hex_sha256 = iff!(sha256_output, Some(crate::digest::sha256_bytes(v)), None);
let hex_sha384 = iff!(sha384_output, Some(crate::digest::sha384_bytes(v)), None);
let hex_sha512 = iff!(sha512_output, Some(crate::digest::sha512_bytes(v)), None);
if json_output {
let mut json = BTreeMap::<&'_ str, String>::new();
json.insert("challenge_hex", hex::encode(challenge_bytes));
json.insert("response_hex", hex_string);
hex_sha1.map(|hex_sha1| json.insert("response_sha1_hex", hex::encode(hex_sha1)));
hex_sha256.map(|hex_sha256| json.insert("response_sha256_hex", hex::encode(hex_sha256)));
hex_sha384.map(|hex_sha384| json.insert("response_sha384_hex", hex::encode(hex_sha384)));
hex_sha512.map(|hex_sha512| json.insert("response_sha512_hex", hex::encode(hex_sha512)));
println!("{}", serde_json::to_string_pretty(&json).expect("Convert to JSON failed!"));
} else {
success!("Challenge HEX: {}", hex::encode(challenge_bytes));
success!("Response HEX: {}", hex_string);
if let Some(hex_sha256) = hex_sha256 { success!("Response SHA256 HEX: {}", hex::encode(hex_sha256)); }
if let Some(hex_sha384) = hex_sha384 { success!("Response SHA384 HEX: {}", hex::encode(hex_sha384)); }
if let Some(hex_sha512) = hex_sha512 { success!("Response SHA512 HEX: {}", hex::encode(hex_sha512)); }
}
let hmac_result = hmacutil::calculate_hmac_sha1_result(&secret_bytes, &challenge_bytes, variable);
hmacutil::output_hmac_result(sub_arg_matches, json_output, challenge_bytes, &hmac_result);
Ok(None)
}
}
fn get_secret_bytes(sub_arg_matches: &ArgMatches) -> XResult<Vec<u8>> {
let secret_bytes: Vec<u8> = if let Some(secret) = sub_arg_matches.value_of("secret") {
opt_result!(hex::decode(secret), "Decode secret hex: {}, failed: {}", secret)
} else {
return simple_error!("Secret must assigned");
};
if secret_bytes.len() != 20 {
return simple_error!("Secret length must be 20, actual is: {}", secret_bytes.len());
}
Ok(secret_bytes)
}