prometheus_wireguard_exporter/src/main.rs

167 lines
5.1 KiB
Rust
Raw Normal View History

2019-04-23 23:06:35 +02:00
extern crate serde_json;
#[macro_use]
extern crate failure;
use clap;
use clap::Arg;
use futures::future::{done, ok, Either, Future};
use http::StatusCode;
use hyper::service::service_fn;
use hyper::{Body, Request, Response, Server};
use log::{error, info, trace};
use std::env;
mod options;
use options::Options;
mod exporter_error;
use exporter_error::ExporterError;
mod render_to_prometheus;
use render_to_prometheus::RenderToPrometheus;
mod wireguard;
use std::convert::TryFrom;
use std::process::Command;
use std::string::String;
use wireguard::WireGuard;
2019-05-15 16:55:07 +02:00
mod wireguard_config;
2019-05-20 10:09:02 +02:00
use wireguard_config::peer_entry_hashmap_try_from;
2019-04-23 23:06:35 +02:00
fn check_compliance(req: &Request<Body>) -> Result<(), Response<Body>> {
if req.uri() != "/metrics" {
trace!("uri not allowed {}", req.uri());
Err(Response::builder()
.status(StatusCode::NOT_FOUND)
.body(hyper::Body::empty())
.unwrap())
} else if req.method() != "GET" {
trace!("method not allowed {}", req.method());
Err(Response::builder()
.status(StatusCode::METHOD_NOT_ALLOWED)
.body(hyper::Body::empty())
.unwrap())
} else {
Ok(())
}
}
fn handle_request(
req: Request<Body>,
options: Options,
) -> impl Future<Item = Response<Body>, Error = failure::Error> {
trace!("{:?}", req);
done(check_compliance(&req)).then(move |res| match res {
Ok(_) => Either::A(perform_request(req, &options).then(|res| match res {
Ok(body) => ok(body),
Err(err) => {
error!("internal server error: {:?}", err);
ok(Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.body(hyper::Body::empty())
.unwrap())
}
})),
Err(err) => Either::B(ok(err)),
})
}
2019-05-20 09:45:20 +02:00
fn wg_with_text(
wg_config_str: &str,
wg_output: ::std::process::Output,
) -> Result<Response<Body>, ExporterError> {
2019-05-20 10:09:02 +02:00
let pehm = peer_entry_hashmap_try_from(wg_config_str)?;
trace!("pehm == {:?}", pehm);
2019-05-20 09:45:20 +02:00
let wg_output_string = String::from_utf8(wg_output.stdout)?;
let wg = WireGuard::try_from(&wg_output_string as &str)?;
2019-05-20 10:21:32 +02:00
Ok(Response::new(Body::from(wg.render_with_names(Some(&pehm)))))
2019-05-20 09:45:20 +02:00
}
2019-04-23 23:06:35 +02:00
fn perform_request(
_req: Request<Body>,
2019-05-17 19:32:35 +02:00
options: &Options,
2019-04-23 23:06:35 +02:00
) -> impl Future<Item = Response<Body>, Error = ExporterError> {
trace!("perform_request");
2019-05-17 19:32:35 +02:00
// this is needed to satisfy the borrow checker
let options = options.clone();
2019-04-23 23:06:35 +02:00
done(
Command::new("wg")
.arg("show")
.arg("all")
.arg("dump")
.output(),
)
.from_err()
.and_then(|output| {
2019-05-17 19:32:35 +02:00
if let Some(extract_names_config_file) = options.extract_names_config_file {
Either::A(
2019-05-20 09:45:20 +02:00
done(::std::fs::read_to_string(extract_names_config_file))
2019-05-17 19:32:35 +02:00
.from_err()
2019-05-20 09:45:20 +02:00
.and_then(|wg_config_string| wg_with_text(&wg_config_string as &str, output)),
2019-05-17 19:32:35 +02:00
)
} else {
Either::B(
done(String::from_utf8(output.stdout))
2019-04-23 23:06:35 +02:00
.from_err()
2019-05-17 19:32:35 +02:00
.and_then(|output_str| {
trace!("{}", output_str);
done(WireGuard::try_from(&output_str as &str))
.from_err()
.and_then(|wg| ok(Response::new(Body::from(wg.render()))))
}),
)
}
2019-04-23 23:06:35 +02:00
})
}
fn main() {
let matches = clap::App::new("prometheus_wireguard_exporter")
2019-05-20 10:09:02 +02:00
.version("1.2.0")
2019-04-23 23:06:35 +02:00
.author("Francesco Cogno <francesco.cogno@outlook.com>")
.arg(
Arg::with_name("port")
.short("p")
2019-05-17 19:32:35 +02:00
.help("exporter port")
2019-04-23 23:06:35 +02:00
.default_value("9576")
.takes_value(true),
)
.arg(
Arg::with_name("verbose")
.short("v")
.help("verbose logging")
.takes_value(false),
)
2019-05-17 19:32:35 +02:00
.arg(
Arg::with_name("extract_names_config_file")
.short("n")
.help("If set, the exporter will look in the specified WireGuard config file for peer names (must be in [Peer] definition and be a comment)")
.takes_value(true))
2019-04-23 23:06:35 +02:00
.get_matches();
let options = Options::from_claps(&matches);
if options.verbose {
env::set_var("RUST_LOG", "prometheus_wireguard_exporter=trace");
} else {
env::set_var("RUST_LOG", "prometheus_wireguard_exporter=info");
}
env_logger::init();
info!("using options: {:?}", options);
let bind = matches.value_of("port").unwrap();
let bind = u16::from_str_radix(&bind, 10).expect("port must be a valid number");
let addr = ([0, 0, 0, 0], bind).into();
info!("starting exporter on {}", addr);
let new_svc = move || {
let options = options.clone();
service_fn(move |req| handle_request(req, options.clone()))
};
let server = Server::bind(&addr)
.serve(new_svc)
.map_err(|e| eprintln!("server error: {}", e));
hyper::rt::run(server);
}