feat: v1.0.9, add init via ssh

This commit is contained in:
Hatter Jiang
2025-06-28 12:12:16 +08:00
parent a8d3f6dadb
commit 88154c9397
10 changed files with 167 additions and 73 deletions

View File

@@ -6,7 +6,7 @@ use clap::{App, Arg, ArgMatches, SubCommand};
use hyper::body::Buf;
use hyper::{Body, Client, Method, Request, Response, StatusCode};
use rust_util::util_clap::{Command, CommandError};
use rust_util::{debugging, opt_result, opt_value_result, simple_error, success, XResult};
use rust_util::{debugging, iff, opt_result, opt_value_result, simple_error, success, XResult};
use serde_json::{json, Map, Value};
use crate::jose;
@@ -34,6 +34,8 @@ impl Command for CommandImpl {
.arg(Arg::with_name("yubikey-challenge").long("yubikey-challenge").short("c").takes_value(true).help("Yubikey challenge"))
.arg(Arg::with_name("comment").long("comment").takes_value(true).help("Comment"))
.arg(Arg::with_name("force-write").long("force-write").short("F").help("Force write value"))
.arg(Arg::with_name("read-from-pinentry").long("read-from-pinentry").help("Read from pin-entry"))
.arg(Arg::with_name("ssh-remote").long("ssh-remote").takes_value(true).help("SSH remote, root@example or localhost"))
}
fn run(&self, arg_matches: &ArgMatches, sub_arg_matches: &ArgMatches) -> CommandError {
@@ -104,17 +106,16 @@ async fn do_direct_init(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatc
}
async fn do_init(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let ssh_remote = sub_arg_matches.value_of("ssh-remote").map(|s| s.to_string());
let connect = sub_arg_matches.value_of("connect").expect("Get argument listen error");
let read_from_pinentry = sub_arg_matches.is_present("read-from-pinentry");
let client = Client::new();
// let client = Client::new();
let uri = format!("http://{}/status", connect);
debugging!("Request uri: {}", &uri);
let req = Request::builder().method(Method::GET).uri(uri).body(Body::empty())?;
let req_response = client.request(req).await?;
if req_response.status() != StatusCode::OK {
return simple_error!("Server status is not success: {}", req_response.status().as_u16());
}
let data = response_to_value(req_response).await?;
let data = send_kms_request_with_ssh_enabled(&ssh_remote, true, &uri, &None).await?;
debugging!("Get status: {}", &data);
let status = &data["status"];
if let Some(status) = status.as_str() {
@@ -129,25 +130,81 @@ async fn do_init(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>
let instance_public_key_jwk = &data["instance_public_key_jwk"];
println!("Instance server public key JWK: {}", instance_public_key_jwk);
let line = read_line("Input encrypted master key: ")?;
let line = {
let line = read_line("Input clear(starts with hex: or base64:) or encrypted master key: ", read_from_pinentry)?;
if line.starts_with("hex:") || line.starts_with("base64:") {
let jwk = opt_result!(serde_json::to_string(&instance_public_key_jwk), "Serialize instance server public key JWK: {} failed");
master_key_encrypt(&line, &jwk)?
} else {
line
}
};
let uri = format!("http://{}/init", connect);
debugging!("Request uri: {}", &uri);
let body = json!({
"encrypted_master_key": line,
});
let body = serde_json::to_string(&body)?;
let req = Request::builder().method(Method::POST).uri(uri).body(Body::from(body))?;
let req_response = client.request(req).await?;
if req_response.status() != StatusCode::OK {
let status = req_response.status().as_u16();
let data = response_to_value(req_response).await?;
return simple_error!("Server status is not success: {}, response: {}", status, data);
}
let _ = send_kms_request_with_ssh_enabled(&ssh_remote, false, &uri, &Some(body)).await?;
success!("Init finished");
Ok(Some(0))
}
async fn send_kms_request_with_ssh_enabled(ssh_remote: &Option<String>, get_request: bool, uri: &str, body: &Option<String>) -> XResult<Value> {
match ssh_remote {
None => {
let client = Client::new();
let method = iff!(get_request, Method::GET, Method::POST);
let request_body = match body {
None => Body::empty(),
Some(body) => Body::from(body.clone()),
};
let req = Request::builder().method(method).uri(uri).body(request_body)?;
let req_response = client.request(req).await?;
if req_response.status() != StatusCode::OK {
return simple_error!("Server status is not success: {}", req_response.status().as_u16());
}
let data = response_to_value(req_response).await?;
Ok(data)
}
Some(ssh_remote) => {
let mut c;
if ssh_remote == "localhost" {
c = std::process::Command::new("curl");
} else {
c = std::process::Command::new("ssh");
c.args([ssh_remote, "curl"]);
}
c.arg("-s");
if !get_request {
c.args(["-X", "POST"]);
}
if let Some(body) = body {
c.args(["-H", "x-body-based64-encoded:1"]);
c.args(["--data-raw", &format!("{}", STANDARD.encode(&body))]);
}
c.arg(uri);
debugging!("Run command: {:?}", c);
let output = opt_result!(c.output(), "Call: {:?} failed: {}", c);
if !output.status.success() {
return simple_error!("Call: {:?} exit with error", output);
}
debugging!("Output: {:?}", output);
let data: Value = serde_json::from_slice(&output.stdout)?;
if let Value::Object(data_map) = &data {
if let Some(error) = data_map.get("error") {
if let Value::String(error) = error {
return simple_error!("Get error: {}, details: {}", error, data);
}
}
}
Ok(data)
}
}
}
async fn do_read(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let body = if let Some(name) = sub_arg_matches.value_of("name") {
json!({ "name": name })
@@ -214,15 +271,26 @@ async fn do_decrypt(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<
Ok(Some(0))
}
fn do_offline_init(_arg_matches: &ArgMatches<'_>, _sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let line = read_line("Input master key: ")?;
let master_key = if line.starts_with("hex:") {
let hex: String = line.chars().skip(4).collect();
fn do_offline_init(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let read_from_pinentry = sub_arg_matches.is_present("read-from-pinentry");
let line = read_line("Input master key: ", read_from_pinentry)?;
let jwk = read_line("Input JWK: ", read_from_pinentry)?;
let encrypted_master_key = master_key_encrypt(&line, &jwk)?;
success!("Encrypted master key: {}", encrypted_master_key);
Ok(Some(0))
}
fn master_key_encrypt(master_key: &str, jwk: &str) -> XResult<String> {
let master_key = if master_key.starts_with("hex:") {
let hex: String = master_key.chars().skip(4).collect();
hex::decode(&hex)?
} else if line.starts_with("base64:") {
let base64: String = line.chars().skip(7).collect();
} else if master_key.starts_with("base64:") {
let base64: String = master_key.chars().skip(7).collect();
STANDARD.decode(&base64)?
} else if line.starts_with("LKMS:") {
} else if master_key.starts_with("LKMS:") {
#[cfg(feature = "yubikey")]
{
use crate::yubikey_hmac;
@@ -230,23 +298,28 @@ fn do_offline_init(_arg_matches: &ArgMatches<'_>, _sub_arg_matches: &ArgMatches<
let challenge = opt_result!(
pinentry_util::read_pin(Some("Input yubikey challenge"), Some("Challenge: ")), "Read challenge failed: {}");
let derived_key = yubikey_hmac::yubikey_challenge_as_32_bytes(challenge.get_pin().as_bytes())?;
let (key, _) = jose::deserialize_jwe_aes(&line, &derived_key)?;
let (key, _) = jose::deserialize_jwe_aes(&master_key, &derived_key)?;
key
}
#[cfg(not(feature = "yubikey"))]
return simple_error!("Yubikey feature is not enabled.");
} else {
line.as_bytes().to_vec()
master_key.as_bytes().to_vec()
};
let jwk = read_line("Input JWK: ")?;
let rsa_public_key = jwk_to_rsa_pubic_key(&jwk)?;
let encrypted_master_key = jose::serialize_jwe_rsa(&master_key, &rsa_public_key)?;
success!("Encrypted master key: {}", encrypted_master_key);
Ok(Some(0))
Ok(encrypted_master_key)
}
fn read_line(prompt: &str) -> XResult<String> {
fn read_line(prompt: &str, pinentry: bool) -> XResult<String> {
if pinentry {
read_line_from_pinentry_util(prompt)
} else {
read_line_from_terminal(prompt)
}
}
fn read_line_from_terminal(prompt: &str) -> XResult<String> {
std::io::stdout().write_all(prompt.as_bytes()).ok();
std::io::stdout().flush().ok();
let mut line = String::new();
@@ -256,6 +329,11 @@ fn read_line(prompt: &str) -> XResult<String> {
Ok(line.trim().to_string())
}
fn read_line_from_pinentry_util(prompt: &str) -> XResult<String> {
let pin = opt_result!(pinentry_util::read_pin(Some(prompt.to_string()), Some("PIN:".to_string())), "Read from pin-entry failed: {}");
Ok(pin.get_pin().to_string())
}
async fn response_to_value(response: Response<Body>) -> XResult<Value> {
let req_body = response.into_body();
let whole_body = hyper::body::aggregate(req_body).await?;