hatter/local-mini-kms
1
1
Fork 0
Code Issues 2 Pull Requests Projects Releases Wiki Activity

feat: cli read/write

Browse Source
This commit is contained in:
Hatter Jiang
2023-08-13 16:06:48 +08:00
parent 167e37f9f8
commit 05446372cf
2 changed files with 82 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
src/cli.rs
View File

@@ -1,7 +1,7 @@
use std::io::Write;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use clap::{App, Arg, ArgMatches, SubCommand};
use hyper::{Body, Client, Method, Request, Response, StatusCode};
use hyper::body::Buf;
@@ -25,9 +25,13 @@ impl Command for CommandImpl {
.arg(Arg::with_name("offline-init").long("offline-init").help("Offline init server"))
.arg(Arg::with_name("encrypt").long("encrypt").help("Encrypt text"))
.arg(Arg::with_name("decrypt").long("decrypt").help("Decrypt text"))
.arg(Arg::with_name("read").long("read").help("Read value"))
.arg(Arg::with_name("write").long("write").help("Write value"))
.arg(Arg::with_name("value").long("value").takes_value(true).help("Value, for encrypt or decrypt"))
.arg(Arg::with_name("key").long("key").takes_value(true).help("Read/Write key name"))
.arg(Arg::with_name("value-hex").long("value-hex").takes_value(true).help("Value(hex), for encrypt"))
.arg(Arg::with_name("value-base64").long("value-base64").takes_value(true).help("Value(base64), for encrypt"))
.arg(Arg::with_name("force-write").long("force-write").help("Force write value"))
}
fn run(&self, arg_matches: &ArgMatches, sub_arg_matches: &ArgMatches) -> CommandError {
@@ -36,6 +40,8 @@ impl Command for CommandImpl {
let offline_init = sub_arg_matches.is_present("offline-init");
let encrypt = sub_arg_matches.is_present("encrypt");
let decrypt = sub_arg_matches.is_present("decrypt");
let read = sub_arg_matches.is_present("read");
let write = sub_arg_matches.is_present("write");
let rt = tokio::runtime::Runtime::new().expect("Create tokio runtime error");
if init {
rt.block_on(async { do_init(arg_matches, sub_arg_matches).await })
@@ -47,6 +53,10 @@ impl Command for CommandImpl {
rt.block_on(async { do_encrypt(arg_matches, sub_arg_matches).await })
} else if decrypt {
rt.block_on(async { do_decrypt(arg_matches, sub_arg_matches).await })
} else if read {
rt.block_on(async { do_read(arg_matches, sub_arg_matches).await })
} else if write {
rt.block_on(async { do_write(arg_matches, sub_arg_matches).await })
} else {
simple_error!("Need a flag")
}
@@ -130,15 +140,46 @@ async fn do_init(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>
Ok(Some(0))
}
async fn do_read(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let body = if let Some(key) = sub_arg_matches.value_of("key") {
json!({ "name": key })
} else {
return simple_error!("Require key");
};
let data = do_inner_request(sub_arg_matches, "read", &body).await?;
success!("Value: {}", serde_json::to_string_pretty(&data)?);
Ok(Some(0))
}
async fn do_write(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let key = if let Some(key) = sub_arg_matches.value_of("key") {
key
} else {
return simple_error!("Require key");
};
let value = sub_arg_matches.value_of("value");
let value_hex = sub_arg_matches.value_of("value-hex");
let value_base64 = sub_arg_matches.value_of("value-base64");
let force_write = sub_arg_matches.is_present("force-write");
let body = if let Some(value) = value {
json!({ "name": key, "force_write": force_write, "value": json!({"value": value}) })
} else if let Some(value_hex) = value_hex {
json!({ "name": key, "force_write": force_write, "value": json!({"value_hex": value_hex}) })
} else if let Some(value_base64) = value_base64 {
json!({ "name": key, "force_write": force_write, "value": json!({"value_base64": value_base64}) })
} else {
return simple_error!("Require one of value, value-hex, value-base64");
};
let data = do_inner_request(sub_arg_matches, "write", &body).await?;
success!("Value: {}", serde_json::to_string_pretty(&data)?);
Ok(Some(0))
}
async fn do_encrypt(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let connect = sub_arg_matches.value_of("connect").expect("Get argument listen error");
let value = sub_arg_matches.value_of("value");
let value_hex = sub_arg_matches.value_of("value-hex");
let value_base64 = sub_arg_matches.value_of("value-base64");
let client = Client::new();
let uri = format!("http://{}/encrypt", connect);
debugging!("Request uri: {}", &uri);
let body = if let Some(value) = value {
json!({ "value": value })
} else if let Some(value_hex) = value_hex {
@@ -148,41 +189,18 @@ async fn do_encrypt(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<
} else {
return simple_error!("Require one of value, value-hex, value-base64");
};
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 data = response_to_value(req_response).await?;
success!("Encrypted value: {}", data["encrypted_value"].as_str().expect("Get encrypted_value error"));
let data = do_inner_request(sub_arg_matches, "encrypt", &body).await?;
success!("Value: {}", serde_json::to_string_pretty(&data)?);
Ok(Some(0))
}
async fn do_decrypt(_arg_matches: &ArgMatches<'_>, sub_arg_matches: &ArgMatches<'_>) -> CommandError {
let connect = sub_arg_matches.value_of("connect").expect("Get argument listen error");
let value = opt_value_result!(sub_arg_matches.value_of("value"), "Argument value required");
let client = Client::new();
let uri = format!("http://{}/decrypt", connect);
debugging!("Request uri: {}", &uri);
let body = json!({ "encrypted_value": value });
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 data = response_to_value(req_response).await?;
success!("Encrypted value(hex): {}", data["value_hex"].as_str().expect("Get value_hex error"));
success!("Encrypted value(base64): {}", data["value_base64"].as_str().expect("Get value_base64 error"));
success!("Encrypted value: {}", String::from_utf8_lossy(&hex::decode(data["value_hex"].as_str().expect("Get value_hex error"))?).to_string());
let data = do_inner_request(sub_arg_matches, "decrypt", &body).await?;
success!("Value: {}", serde_json::to_string_pretty(&data)?);
Ok(Some(0))
}
@@ -220,4 +238,20 @@ async fn response_to_value(response: Response<Body>) -> XResult<Value> {
let whole_body = hyper::body::aggregate(req_body).await?;
let data: Value = serde_json::from_reader(whole_body.reader())?;
Ok(data)
}
}
async fn do_inner_request(sub_arg_matches: &ArgMatches<'_>, action: &str, body: &Value) -> XResult<Value> {
let connect = sub_arg_matches.value_of("connect").expect("Get argument listen error");
let body = serde_json::to_string(&body)?;
let client = Client::new();
let uri = format!("http://{}/{}", connect, action);
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);
}
Ok(response_to_value(req_response).await?)
}

20
src/serve_read_write.rs
View File

@@ -94,9 +94,19 @@ async fn inner_write(req: Request<Body>) -> XResult<(StatusCode, Value)> {
db::insert_key(&conn, &new_db_key)?;
}
Ok((StatusCode::OK, json!({
"name": name.to_string(),
"override": db_key.is_some(),
"encrypted_value": encrypt_value,
})))
let response_body = if let Some(db_key) = db_key {
json!({
"name": name.to_string(),
"override": true,
"encrypted_value": encrypt_value,
"previous_encrypted_value": db_key.encrypted_key,
})
} else {
json!({
"name": name.to_string(),
"override": false,
"encrypted_value": encrypt_value,
})
};
Ok((StatusCode::OK, response_body))
}