nixos-exporter/src/main.rs

use axum::{
    extract::Query,
    http::StatusCode,
    response::{IntoResponse, Response},
    routing::get,
    Router,
};

use std::collections::HashMap;
use std::net::SocketAddr;
use std::str::FromStr;

#[derive(PartialEq)]
enum OperationMode {
    Exporter,
    Validator,
}

fn parse_nix_store_path(path: std::path::PathBuf) -> Result<(String, String), String> {
    let (hash, name) = path.file_name()
        .ok_or_else(String::default)?
        .to_str()
        .ok_or_else(String::default)?
        .split_once("-")
        .ok_or_else(String::default)?;
    return Ok((hash.to_string(), name.to_string()));
}

#[tokio::main]
async fn main() {
    let mut listen = String::from("[::]:9152");
    let mut operationmode = OperationMode::Exporter;
    let mut args = std::env::args();
    args.next();
    loop {
        let arg = if let Some(arg) = args.next() {
            arg
        } else {
            break;
        };

        match arg.as_str() {
            "--help" | "-h" => {
                println!("Prometheus exporter for NixOS systems");
                println!("Use --listen <addr:port> bind the web service.");
                println!("Output will be on /metrics endpoint. HTTP 500 if something broke while scraping.");
                std::process::exit(0);
            }
            "--listen" => {
                listen = args.next().unwrap();
            }
            "exporter" => {
                operationmode = OperationMode::Exporter;
            }
            "validator" => {
                operationmode = OperationMode::Validator;
            }
            unknown => {
                println!("unknown option: {}", unknown);
                std::process::exit(1)
            }
        }
    }

    let mut app = Router::new();
    if operationmode == OperationMode::Exporter {
        println!("Running NixOS Exporter in Exporter mode");
        app = app.route("/metrics", get(metrics));
    } else if operationmode == OperationMode::Validator {
        println!("Running NixOS Exporter in Validator mode");
        app = app.route("/metrics", get(check));
    }

    let addr = SocketAddr::from_str(&listen).unwrap();
    println!("listening on http://{}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

fn get_current_system() -> Result<(String, String), String> {
    let symlink = match std::fs::read_link("/run/current-system") {
        Ok(symlink) => symlink,
        Err(err) => return Err(err.to_string()),
    };

    let (hash, name) = parse_nix_store_path(symlink)?;

    Ok((String::from(hash), String::from(name)))
}

async fn metrics() -> Response {
    let current_system = get_current_system();
    let (hash, name) = match current_system {
        Ok((hash, name)) => (hash, name),
        Err(err) => {
            println!("failed: {}", err);
            return (StatusCode::INTERNAL_SERVER_ERROR, "").into_response();
        }
    };

    (
        StatusCode::OK,
        format!("nixos_current_system_hash{{hash=\"{}\"}} 1\nnixos_current_system_name{{name=\"{}\"}} 1\n", hash, name)
    ).into_response()
}

async fn check(Query(params): Query<HashMap<String, String>>) -> Response {
    let target = match params.get("target") {
        Some(target) => target,
        None => {
            return (StatusCode::NOT_FOUND, "specify target").into_response();
        },
    };

    let client = reqwest::Client::new();

    let prometheus_req = match client.get(format!("https://prometheus.monitoring.clerie.de/api/v1/query?query=nixos_nixos_current_system_hash{{job=%22nixos-exporter%22,instance=%22{}.mon.clerie.de:9152%22}}", target))
        .header("Accept", "application/json")
        .send().await {
            Ok(req) => req,
            Err(_) => {
                return (StatusCode::INTERNAL_SERVER_ERROR, "Promehteus can't get reached").into_response();
            },
        };


    match prometheus_req.status() {
        reqwest::StatusCode::OK => (),
        _ => {
            return (StatusCode::NOT_FOUND, "Target does not exist in Hydra").into_response();
        },
    }

    let prometheus_body = match prometheus_req.json::<serde_json::Value>().await {
        Ok(body) => body,
        Err(_) => {
            return (StatusCode::INTERNAL_SERVER_ERROR, "Invalid response from Hydra").into_response();
        },
    };

    let current_system_hash = match prometheus_body["data"]["result"][0]["metric"]["hash"].as_str() {
        Some(nix_store_path) => nix_store_path,
        _ => {
            return (StatusCode::INTERNAL_SERVER_ERROR, "No buildoutput found in Hydra").into_response();
        },
    };

    let hydra_req = match client.get(format!("https://hydra.clerie.de/job/nixfiles/nixfiles/nixosConfigurations.{}/latest", target))
        .header("Accept", "application/json")
        .send().await {
            Ok(req) => req,
            Err(_) => {
                return (StatusCode::INTERNAL_SERVER_ERROR, "Hydra can't get reached").into_response();
            },
        };

    match hydra_req.status() {
        reqwest::StatusCode::OK => (),
        _ => {
            return (StatusCode::NOT_FOUND, "Target does not exist in Hydra").into_response();
        },
    }

    let hydra_body = match hydra_req.json::<serde_json::Value>().await {
        Ok(body) => body,
        Err(_) => {
            return (StatusCode::INTERNAL_SERVER_ERROR, "Invalid response from Hydra").into_response();
        },
    };

    let nix_store_path = match hydra_body["buildoutputs"]["out"]["path"].as_str() {
        Some(nix_store_path) => nix_store_path,
        _ => {
            return (StatusCode::INTERNAL_SERVER_ERROR, "No buildoutput found in Hydra").into_response();
        },
    };

    let (hydra_system_hash, _) = match parse_nix_store_path(std::path::PathBuf::from(nix_store_path)) {
        Ok((hash, name)) => (hash, name),
        Err(_) => {
            return (StatusCode::INTERNAL_SERVER_ERROR, "Invalid store path returned by Hydra").into_response();
        },
    };

    let mut status = "0";

    if current_system_hash == hydra_system_hash {
        status = "1";
    }

    return (
        StatusCode::OK,
        format!("nixos_current_system_valid{{target=\"{}.net.clerie.de:9152\"}} {}\n", target, status)
    ).into_response();
}