diff --git a/Cargo.lock b/Cargo.lock index f0adec9..f69d2d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -20,6 +20,35 @@ dependencies = [ "winapi", ] +[[package]] +name = "argh" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbb41d85d92dfab96cb95ab023c265c5e4261bb956c0fb49ca06d90c570f1958" +dependencies = [ + "argh_derive", + "argh_shared", +] + +[[package]] +name = "argh_derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be69f70ef5497dd6ab331a50bd95c6ac6b8f7f17a7967838332743fbd58dc3b5" +dependencies = [ + "argh_shared", + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "argh_shared" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6f8c380fa28aa1b36107cd97f0196474bb7241bb95a453c5c01a15ac74b2eac" + [[package]] name = "atty" version = "0.2.14" @@ -59,6 +88,7 @@ checksum = "a4a45a46ab1f2412e53d3a0ade76ffad2025804294569aae387231a0cd6e0899" name = "burp_exporter" version = "0.1.0" dependencies = [ + "argh", "gethostname", "openssl", "prometheus_exporter_base", @@ -66,6 +96,7 @@ dependencies = [ "serde_json", "tokio", "tokio-native-tls", + "toml", "tracing", "tracing-subscriber", ] @@ -227,6 +258,15 @@ dependencies = [ "winapi", ] +[[package]] +name = "heck" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d621efb26863f0e9924c6ac577e8275e5e6b77455db64ffa6c65c904e9e132c" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "hermit-abi" version = "0.1.19" @@ -941,6 +981,15 @@ dependencies = [ "webpki", ] +[[package]] +name = "toml" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31142970826733df8241ef35dc040ef98c679ab14d7c3e54d827099b3acecaa" +dependencies = [ + "serde", +] + [[package]] name = "tower-service" version = "0.3.1" @@ -1015,6 +1064,12 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59547bce71d9c38b83d9c0e92b6066c4253371f15005def0c30d9657f50c7642" +[[package]] +name = "unicode-segmentation" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e8820f5d777f6224dc4be3632222971ac30164d4a258d595640799554ebfd99" + [[package]] name = "unicode-width" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index e5ad5fa..36b419b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,10 +6,12 @@ license = "MIT OR Apache-2.0" edition = "2021" [dependencies] +argh = "^0.1.7" openssl = "^0.10.38" gethostname = "^0.2.2" serde_json = "1.0.78" tokio-native-tls = "^0.3.0" +toml = "^0.5.8" tracing = "^0.1.30" [dependencies.prometheus_exporter_base] diff --git a/src/main.rs b/src/main.rs index f79c691..1650931 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,2 +1,90 @@ -fn main() { +use std::fs::File; +use std::io::prelude::*; +use std::net::IpAddr; +use std::path::Path; + +use argh::FromArgs; +use prometheus_exporter_base::render_prometheus; +use toml; +use tracing::{debug, error}; + +mod burp; +mod config; +mod metrics; +use burp::client::ClientConnector; +use config::Config; + +/// BURP exporter +#[derive(FromArgs)] +struct Args { + /// configuration file path + #[argh(option, short = 'c')] + config_file: Option, +} + +#[tokio::main(flavor = "current_thread")] +async fn main() { + tracing_subscriber::fmt::init(); + let args: Args = argh::from_env(); + let config = match load_config(args.config_file) { + Ok(c) => c, + Err(e) => { + eprintln!("Failed to load configuration: {}", e); + ::std::process::exit(1); + } + }; + let host: IpAddr = match config.listen_host.parse() { + Ok(a) => a, + Err(e) => { + eprintln!("Invalid listen host ({}): {}", &config.listen_host, e); + ::std::process::exit(1); + } + }; + let addr = (host, config.listen_port).into(); + debug!("Using configuration: {:?}", &config); + + render_prometheus(addr, config, |_request, config| async move { + let mut connector = ClientConnector::new(&config.server); + if let Some(server_name) = &config.server_name { + connector = connector.server_name(server_name); + } + if let Some(ca_cert) = &config.ca_cert { + connector = connector.ca_cert_file(&ca_cert)?; + } + if let Some(identity_file) = &config.identity { + let password = match &config.identity_password { + Some(p) => &p, + None => "", + }; + connector = connector.identity_file(identity_file, password)?; + } + debug!("Connecting to server: {}", &config.server); + let mut client = match connector.connect().await { + Ok(c) => Ok(c), + Err(e) => { + error!("Error connecting to server: {}", e); + Err(e) + } + }?; + if let Some(p) = &config.password { + client.handshake(&config.client_name, Some(p)).await + } else { + client.handshake(&config.client_name, None).await + }?; + Ok(metrics::get_metrics(&mut client).await?) + }) + .await; +} + +fn load_config>( + path: Option

, +) -> Result> { + if let Some(path) = path { + let mut f = File::open(path)?; + let mut contents = String::new(); + f.read_to_string(&mut contents)?; + Ok(toml::from_str(&contents)?) + } else { + Ok(Config::default()) + } }