acaua/crates/acaua-cli/src/web.rs

use acaua_core::fusion::Alert;
use acaua_core::storage::{DatasetSummary, Storage};
use acaua_core::update::UpdateStatus;
use acaua_core::zeek::{self, AnalysisReport};
use acaua_core::{update, Result};
use axum::body::Body;
use axum::extract::State;
use axum::http::{header, StatusCode, Uri};
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{Json, Router};
use rust_embed::RustEmbed;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::net::TcpListener;
use tower_http::compression::CompressionLayer;
use tower_http::cors::CorsLayer;
use tower_http::trace::TraceLayer;

const LOGO: &[u8] = include_bytes!("../../../logo.png");

#[derive(RustEmbed)]
#[folder = "../../apps/gui/ui/"]
struct Assets;

#[derive(Clone)]
struct AppState {
    storage: Arc<Storage>,
}

#[derive(Serialize)]
struct Health {
    status: &'static str,
    version: &'static str,
}

#[derive(Serialize)]
struct Status {
    version: &'static str,
    mode: &'static str,
    dataset: DatasetSummary,
}

#[derive(Serialize)]
struct ErrorBody {
    error: String,
}

#[derive(Deserialize)]
struct AnalyzeRequest {
    path: String,
}

async fn health() -> Json<Health> {
    Json(Health {
        status: "ok",
        version: env!("CARGO_PKG_VERSION"),
    })
}

async fn status(State(state): State<AppState>) -> Json<Status> {
    Json(Status {
        version: env!("CARGO_PKG_VERSION"),
        mode: "serve",
        dataset: state
            .storage
            .summary()
            .unwrap_or_else(|_| DatasetSummary::empty()),
    })
}

async fn alerts() -> Json<Vec<Alert>> {
    Json(Vec::new())
}

async fn analyze(
    State(state): State<AppState>,
    Json(request): Json<AnalyzeRequest>,
) -> std::result::Result<Json<AnalysisReport>, (StatusCode, Json<ErrorBody>)> {
    zeek::analyze_dir(&request.path, state.storage.root())
        .map(Json)
        .map_err(|error| {
            (
                StatusCode::BAD_REQUEST,
                Json(ErrorBody {
                    error: error.to_string(),
                }),
            )
        })
}

async fn update_check() -> std::result::Result<Json<UpdateStatus>, (StatusCode, Json<ErrorBody>)> {
    update::check(env!("CARGO_PKG_VERSION"))
        .map(Json)
        .map_err(|error| {
            (
                StatusCode::BAD_GATEWAY,
                Json(ErrorBody {
                    error: error.to_string(),
                }),
            )
        })
}

async fn logo() -> impl IntoResponse {
    ([(header::CONTENT_TYPE, "image/png")], Body::from(LOGO))
}

async fn asset(uri: Uri) -> Response {
    let path = uri.path().trim_start_matches('/');
    let path = if path.is_empty() { "index.html" } else { path };
    match Assets::get(path) {
        Some(file) => {
            let mime = mime_guess::from_path(path).first_or_octet_stream();
            Response::builder()
                .header(header::CONTENT_TYPE, mime.as_ref())
                .body(Body::from(file.data.into_owned()))
                .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
        }
        None => (StatusCode::NOT_FOUND, "not found").into_response(),
    }
}

async fn shutdown_signal() {
    let interrupt = async {
        let _ = tokio::signal::ctrl_c().await;
    };
    #[cfg(unix)]
    let terminate = async {
        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
            Ok(mut stream) => {
                stream.recv().await;
            }
            Err(_) => std::future::pending::<()>().await,
        }
    };
    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();
    tokio::select! {
        _ = interrupt => {}
        _ = terminate => {}
    }
}

pub async fn serve(bind: &str, data_dir: PathBuf) -> Result<()> {
    let _ = tracing_subscriber::fmt().try_init();
    std::fs::create_dir_all(&data_dir)?;
    let state = AppState {
        storage: Arc::new(Storage::open(&data_dir)?),
    };
    let app = Router::new()
        .route("/api/health", get(health))
        .route("/api/status", get(status))
        .route("/api/alerts", get(alerts))
        .route("/api/analyze", post(analyze))
        .route("/api/update", get(update_check))
        .route("/logo.png", get(logo))
        .fallback(asset)
        .with_state(state)
        .layer(CompressionLayer::new())
        .layer(CorsLayer::new())
        .layer(TraceLayer::new_for_http());
    let listener = TcpListener::bind(bind).await?;
    println!("Acauã listening on http://{bind}");
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;
    Ok(())
}