From d369e50d341774ff5511190207eaec8054a19c7a Mon Sep 17 00:00:00 2001 From: Vivian Roest Date: Mon, 1 May 2023 13:38:37 +0200 Subject: [PATCH] cleaned up the code --- Cargo.lock | 121 +++++++++++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 4 +- flake.nix | 15 ++++--- src/main.rs | 82 +++++++++++++++++++++++++++-------- 4 files changed, 196 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 07a4b26..d62f8a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -91,6 +91,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + [[package]] name = "libc" version = "0.2.142" @@ -115,18 +121,40 @@ dependencies = [ "adler", ] +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + [[package]] name = "once_cell" version = "1.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7e5500299e16ebb147ae15a00a942af264cf3688f47923b8fc2cd5858f23ad3" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + [[package]] name = "percent-encoding" version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e" +[[package]] +name = "pin-project-lite" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" + [[package]] name = "proc-macro2" version = "1.0.56" @@ -219,6 +247,21 @@ dependencies = [ "serde", ] +[[package]] +name = "sharded-slab" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "900fba806f70c630b0a382d0d825e17a0f19fcd059a2ade1ff237bcddf446b31" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "smallvec" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a507befe795404456341dfab10cef66ead4c041f62b8b11bbb92bffe5d0953e0" + [[package]] name = "spin" version = "0.5.2" @@ -247,6 +290,16 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "thread_local" +version = "1.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdd6f064ccff2d6567adcb3873ca630700f00b5ad3f060c25b5dcfd9a4ce152" +dependencies = [ + "cfg-if", + "once_cell", +] + [[package]] name = "tinyvec" version = "1.6.0" @@ -262,6 +315,64 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tracing" +version = "0.1.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ce8c33a8d48bd45d624a6e523445fd21ec13d3653cd51f681abf67418f54eb8" +dependencies = [ + "cfg-if", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f57e3ca2a01450b1a921183a9c9cbfda207fd822cef4ccb00a65402cbba7a74" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.15", +] + +[[package]] +name = "tracing-core" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24eb03ba0eab1fd845050058ce5e616558e8f8d8fca633e6b163fe25c797213a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ddad33d2d10b1ed7eb9d1f518a5674713876e97e5bb9b7345a7984fbb4f922" +dependencies = [ + "lazy_static", + "log", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30a651bc37f915e81f087d86e62a18eec5f79550c7faff886f7090b4ea757c77" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", +] + [[package]] name = "unicode-bidi" version = "0.3.13" @@ -319,12 +430,20 @@ dependencies = [ ] [[package]] -name = "vault-unseal" +name = "valuable" version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + +[[package]] +name = "vault-unseal" +version = "0.2.0" dependencies = [ "dotenv", "serde", "serde_json", + "tracing", + "tracing-subscriber", "ureq", ] diff --git a/Cargo.toml b/Cargo.toml index 353c368..96bbe99 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "vault-unseal" -version = "0.1.0" +version = "0.2.0" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html @@ -10,3 +10,5 @@ ureq = { version = "*", features = ["json"] } serde_json = "*" dotenv = "*" serde = { version = "*", features = ["derive"] } +tracing = "0.1" +tracing-subscriber = "0.3" diff --git a/flake.nix b/flake.nix index 970755f..956bc89 100644 --- a/flake.nix +++ b/flake.nix @@ -10,12 +10,15 @@ let pkgs = nixpkgs.legacyPackages.${system}; in { packages = { - default = pkgs.rustPlatform.buildRustPackage { - pname = "vault-unseal"; - version = "0.1.0"; - src = self; - cargoSha256 = "sha256-nCOHQU62fzJ9uwUK8n5JsVkKmqQwhG/5GI6rvtejZjY="; - }; + default = + let toml = (builtins.fromTOML (builtins.readFile ./Cargo.toml)); + in pkgs.rustPlatform.buildRustPackage { + pname = toml.package.name; + version = toml.package.version; + src = self; + cargoSha256 = + "sha256-eOvTR7TpFpi83J3G8HPXgOBryTzkq4XWp6CER6UDCbo="; + }; }; }); } diff --git a/src/main.rs b/src/main.rs index 683f7be..b9e0e51 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,10 @@ use std::io::Read; use std::thread; use std::time::Duration; use std::{env, fs::File}; +use tracing::{info, subscriber, warn}; +use tracing_subscriber::FmtSubscriber; +use ureq::Error::Status; +use ureq::Response; #[derive(Debug, Deserialize)] struct KeyFile { @@ -17,8 +21,65 @@ struct HealthCheck { sealed: bool, } +fn is_sealed(health_url: &str) -> bool { + fn parse_hc(x: Response) -> bool { + match x.into_json() { + Ok(HealthCheck { sealed }) => sealed, + Err(_) => false, + } + } + + let resp = ureq::get(health_url).call(); + match resp { + Ok(x) => parse_hc(x), + Err(Status(503, resp)) => parse_hc(resp), + Err(Status(429, _)) => { + info!("got code 429: too many requests, waiting"); + // too many requests + thread::sleep(Duration::from_secs(15)); + false + } + Err(Status(code, resp)) => { + info!( + "error checking health, got code: '{code}', with message: {}", + resp.status_text() + ); + false + } + Err(e) => { + warn!("Got error: {e}"); + false + } + } +} + +fn unseal(keyfile: &KeyFile, unseal_url: &str) { + let len = keyfile.keys.len(); + for (i, key) in keyfile.keys.iter().enumerate() { + let i = i + 1; + match ureq::post(unseal_url).send_json(json!({ "key": key })) { + Ok(resp) if resp.status() == 200 => { + if i < len { + info!("unsealed vault partially {i}/{len}"); + } else { + info!("fully unsealed vault {i}/{len}"); + } + } + Ok(resp) => warn!( + "error unsealing vault, got code '{}', with message: {}", + resp.status(), + resp.status_text() + ), + Err(err) => warn!("error unsealing vault: {err}"), + } + } +} + fn main() -> Result<(), Box> { dotenv().ok(); + let subscriber = FmtSubscriber::new(); + subscriber::set_global_default(subscriber)?; + let vault_addr = env::var("VAULT_ADDR")?; let file = env::var("VAULT_KEY_FILE")?; let interval = env::var("UNSEAL_INTERVAL").unwrap_or(String::from("15")); @@ -33,25 +94,10 @@ fn main() -> Result<(), Box> { let unseal_url = format!("{vault_addr}/v1/sys/unseal"); let health_url = format!("{vault_addr}/v1/sys/health"); - println!("Starting vault unsealer ..."); + info!("Starting vault unsealer at {vault_addr}"); loop { - match ureq::get(&health_url).call() { - Err(ureq::Error::Status(code, resp)) if code == 503 => { - if let Ok(HealthCheck { sealed: true }) = resp.into_json() { - for key in &keyfile.keys { - match ureq::post(&unseal_url).send_json(json!({ "key": key })) { - Ok(resp) if resp.status() != 200 => eprintln!("error unsealing vault"), - Ok(_) => println!("unsealed vault partially"), - Err(err) => eprintln!("error unsealing vault: {err}"), - } - } - } else { - eprintln!("Can't unseal"); - } - } - Err(ureq::Error::Status(_, _)) => (), - Err(e) => eprintln!("{e}"), - _ => (), + if is_sealed(&health_url) { + unseal(&keyfile, &unseal_url); } thread::sleep(interval);