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,4 +1,3 @@
use std::collections::BTreeMap;
use std::ops::Deref;
use clap::{App, Arg, ArgMatches, SubCommand};
@@ -7,6 +6,8 @@ use rust_util::util_msg;
use yubico_manager::config::{Config, Mode, Slot};
use yubico_manager::Yubico;
use crate::hmacutil;
pub struct CommandImpl;
impl Command for CommandImpl {
@@ -27,25 +28,9 @@ impl Command for CommandImpl {
let json_output = sub_arg_matches.is_present("json");
if json_output { util_msg::set_logger_std_out(false); }
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 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");
};
// 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 challenge_bytes = hmacutil::get_challenge_bytes(sub_arg_matches)?;
let mut yubi = Yubico::new();
if let Ok(device) = yubi.find_yubikey() {
success!("Found key, Vendor ID: {:?}, Product ID: {:?}", device.vendor_id, device.product_id);
@@ -59,31 +44,7 @@ impl Command for CommandImpl {
// In HMAC Mode, the result will always be the SAME for the SAME provided challenge
let hmac_result = opt_result!(yubi.challenge_response_hmac(&challenge_bytes, config), "Challenge HMAC failed: {}");
// Just for debug, lets check the hex
let v: &[u8] = hmac_result.deref();
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)); }
}
hmacutil::output_hmac_result(sub_arg_matches, json_output, challenge_bytes, hmac_result.deref());
} else {
warning!("YubiKey not found");
return Ok(Some(1));

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)
}

64
src/hmacutil.rs Normal file
View File

@@ -0,0 +1,64 @@
use std::collections::BTreeMap;
use clap::ArgMatches;
use rust_util::XResult;
use yubico_manager::hmacmode::HmacKey;
use yubico_manager::sec::hmac_sha1;
pub fn get_challenge_bytes(sub_arg_matches: &ArgMatches) -> XResult<Vec<u8>> {
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");
};
if challenge_bytes.len() > 64 {
return simple_error!("Challenge bytes is: {}, more than 64", challenge_bytes.len());
}
Ok(challenge_bytes)
}
pub fn calculate_hmac_sha1_result(secret_bytes: &Vec<u8>, challenge_bytes: &Vec<u8>, variable: bool) -> [u8; 20] {
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);
hmac_result
}
pub fn output_hmac_result(sub_arg_matches: &ArgMatches, json_output: bool, challenge_bytes: Vec<u8>, result: &[u8]) {
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 hex_string = hex::encode(result);
let hex_sha1 = iff!(sha1_output, Some(crate::digest::sha1_bytes(result)), None);
let hex_sha256 = iff!(sha256_output, Some(crate::digest::sha256_bytes(result)), None);
let hex_sha384 = iff!(sha384_output, Some(crate::digest::sha384_bytes(result)), None);
let hex_sha512 = iff!(sha512_output, Some(crate::digest::sha512_bytes(result)), 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)); }
}
}

View File

@@ -11,6 +11,7 @@ mod digest;
mod pivutil;
mod rsautil;
mod pkiutil;
mod hmacutil;
mod pgpcardutil;
mod cmd_list;
mod cmd_u2fregister;