dragoon-server 0.1.0

//! `dragoon-server` — public-relay server for remote-executor.

use std::io::{IsTerminal, Read};
use std::path::PathBuf;

use anyhow::{Context, Result};
use base64::{engine::general_purpose::STANDARD as B64, Engine};
use clap::{Parser, Subcommand};
use dragoon_proto::pubkey::parse_pubkey_blob;
use dragoon_server::{
    app::{build_state, create_app},
    audit, settings::Settings, users_repo, workers_repo,
};
use rusqlite::Connection;

#[derive(Parser, Debug)]
#[command(name = "dragoon-server", version, about = "remote-executor server")]
struct Cli {
    #[command(subcommand)]
    command: Cmd,
}

#[derive(Subcommand, Debug)]
enum Cmd {
    /// Run the HTTP server (axum + tokio).
    Run(RunArgs),
    /// User management.
    User {
        #[command(subcommand)]
        cmd: UserCmd,
    },
    /// Worker token management.
    Token {
        #[command(subcommand)]
        cmd: TokenCmd,
    },
    /// Tail recent audit-log entries.
    Audit {
        #[arg(long, short = 'n', default_value_t = 20)]
        n: i64,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
}

#[derive(Parser, Debug)]
struct RunArgs {
    #[arg(long, short = 'c')]
    config: PathBuf,
}

#[derive(Subcommand, Debug)]
enum UserCmd {
    /// Create a new user. Prints TOTP secret + 10 recovery codes.
    Create {
        #[arg(long)]
        name: String,
        #[arg(long)]
        password: Option<String>,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
    /// Revoke a user (soft delete; sessions stop working immediately).
    Revoke {
        #[arg(long)]
        name: String,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
    /// Reset a user's password.
    Passwd {
        #[arg(long)]
        name: String,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
    /// Manage SSH public keys for a user.
    Pubkey {
        #[command(subcommand)]
        cmd: PubkeyCmd,
    },
}

#[derive(Subcommand, Debug)]
enum PubkeyCmd {
    Add {
        #[arg(long)]
        user: String,
        #[arg(long)]
        label: Option<String>,
        #[arg(long)]
        pub_: PathBuf,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
    List {
        #[arg(long)]
        user: String,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
    Revoke {
        #[arg(long)]
        user: String,
        #[arg(long)]
        fingerprint: String,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
}

#[derive(Subcommand, Debug)]
enum TokenCmd {
    Issue {
        #[arg(long)]
        worker: String,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
    Revoke {
        #[arg(long)]
        worker: String,
        #[arg(long, conflicts_with = "data_dir")]
        config: Option<PathBuf>,
        #[arg(long)]
        data_dir: Option<PathBuf>,
    },
}

fn settings_from(config: &Option<PathBuf>, data_dir: &Option<PathBuf>) -> Result<Settings> {
    if let Some(p) = config {
        Settings::from_toml_path(p)
    } else if let Some(d) = data_dir {
        Ok(Settings::for_test(d.clone()))
    } else {
        anyhow::bail!("either --config or --data-dir is required")
    }
}

fn open_db(settings: &Settings) -> Result<Connection> {
    std::fs::create_dir_all(&settings.data_dir)?;
    let conn = dragoon_server::db::connect(settings.db_path())?;
    dragoon_server::db::bootstrap(&conn)?;
    Ok(conn)
}

fn read_password_or_prompt(explicit: Option<String>) -> Result<String> {
    if let Some(p) = explicit {
        return Ok(p);
    }
    if !std::io::stdin().is_terminal() {
        let mut s = String::new();
        std::io::stdin().read_to_string(&mut s)?;
        return Ok(s.trim_end_matches(['\n', '\r']).to_string());
    }
    Ok(rpassword::prompt_password("Password: ")?)
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    match cli.command {
        Cmd::Run(args) => {
            let settings = Settings::from_toml_path(&args.config)?;
            let runtime = tokio::runtime::Builder::new_multi_thread()
                .enable_all()
                .build()?;
            runtime.block_on(serve(settings))
        }
        Cmd::User { cmd } => match cmd {
            UserCmd::Create {
                name, password, config, data_dir,
            } => {
                let settings = settings_from(&config, &data_dir)?;
                let conn = open_db(&settings)?;
                let pw = read_password_or_prompt(password)?;
                let (uid, secret, codes) = users_repo::create_user(&conn, &name, &pw)?;
                audit::log(
                    &conn,
                    Some("server-cli"),
                    "user.create",
                    Some(&name),
                    None,
                    &serde_json::json!({}),
                )?;
                println!("User created: id={uid}");
                println!("TOTP secret (otpauth://totp/{}:{}?secret={}):", settings.totp_issuer, name, secret);
                println!("  {secret}");
                println!("Recovery codes (one-shot; store somewhere safe):");
                for c in codes {
                    println!("  {c}");
                }
                Ok(())
            }
            UserCmd::Revoke { name, config, data_dir } => {
                let settings = settings_from(&config, &data_dir)?;
                let conn = open_db(&settings)?;
                users_repo::revoke_user(&conn, &name)?;
                audit::log(&conn, Some("server-cli"), "user.revoke", Some(&name), None, &serde_json::json!({}))?;
                println!("revoked: {name}");
                Ok(())
            }
            UserCmd::Passwd { name, config, data_dir } => {
                let settings = settings_from(&config, &data_dir)?;
                let conn = open_db(&settings)?;
                let new = read_password_or_prompt(None)?;
                users_repo::set_password(&conn, &name, &new)?;
                audit::log(&conn, Some("server-cli"), "user.passwd", Some(&name), None, &serde_json::json!({}))?;
                println!("password updated");
                Ok(())
            }
            UserCmd::Pubkey { cmd } => match cmd {
                PubkeyCmd::Add { user, label, pub_, config, data_dir } => {
                    let settings = settings_from(&config, &data_dir)?;
                    let conn = open_db(&settings)?;
                    let u = users_repo::get_user(&conn, &user)?
                        .ok_or_else(|| anyhow::anyhow!("unknown user: {user}"))?;
                    let text = std::fs::read_to_string(&pub_)
                        .with_context(|| format!("read {}", pub_.display()))?;
                    let parts: Vec<&str> = text.split_whitespace().collect();
                    if parts.len() < 2 {
                        anyhow::bail!("not an OpenSSH public-key line");
                    }
                    let blob = B64.decode(parts[1].as_bytes())?;
                    let _parsed = parse_pubkey_blob(&blob)?;
                    let fp = users_repo::add_pubkey(&conn, u.id, &blob, label.as_deref())?;
                    audit::log(
                        &conn, Some("server-cli"), "pubkey.add", Some(&user),
                        Some(&fp), &serde_json::json!({}),
                    )?;
                    println!("{fp}");
                    Ok(())
                }
                PubkeyCmd::List { user, config, data_dir } => {
                    let settings = settings_from(&config, &data_dir)?;
                    let conn = open_db(&settings)?;
                    let u = users_repo::get_user(&conn, &user)?
                        .ok_or_else(|| anyhow::anyhow!("unknown user: {user}"))?;
                    for k in users_repo::list_pubkeys(&conn, u.id)? {
                        let suffix = if k.revoked_at.is_some() { " (revoked)" } else { "" };
                        println!("{}  {}  {}{}", k.fingerprint, k.alg, k.label.as_deref().unwrap_or(""), suffix);
                    }
                    Ok(())
                }
                PubkeyCmd::Revoke { user, fingerprint, config, data_dir } => {
                    let settings = settings_from(&config, &data_dir)?;
                    let conn = open_db(&settings)?;
                    let u = users_repo::get_user(&conn, &user)?
                        .ok_or_else(|| anyhow::anyhow!("unknown user: {user}"))?;
                    let ok = users_repo::revoke_pubkey(&conn, u.id, &fingerprint)?;
                    audit::log(
                        &conn, Some("server-cli"), "pubkey.revoke", Some(&user),
                        Some(&fingerprint), &serde_json::json!({}),
                    )?;
                    println!("{}", if ok { "revoked" } else { "not found" });
                    Ok(())
                }
            },
        },
        Cmd::Token { cmd } => match cmd {
            TokenCmd::Issue { worker, config, data_dir } => {
                let settings = settings_from(&config, &data_dir)?;
                let conn = open_db(&settings)?;
                let code = workers_repo::create_or_replace_register_code(
                    &conn, &worker, settings.worker_register_code_ttl_sec,
                )?;
                audit::log(&conn, Some("server-cli"), "worker.token.issue", Some(&worker), None, &serde_json::json!({}))?;
                println!("register code (valid {} s): {code}", settings.worker_register_code_ttl_sec);
                Ok(())
            }
            TokenCmd::Revoke { worker, config, data_dir } => {
                let settings = settings_from(&config, &data_dir)?;
                let conn = open_db(&settings)?;
                workers_repo::revoke(&conn, &worker)?;
                audit::log(&conn, Some("server-cli"), "worker.token.revoke", Some(&worker), None, &serde_json::json!({}))?;
                println!("revoked");
                Ok(())
            }
        },
        Cmd::Audit { n, config, data_dir } => {
            let settings = settings_from(&config, &data_dir)?;
            let conn = open_db(&settings)?;
            for ev in audit::tail(&conn, n)?.into_iter().rev() {
                println!(
                    "{} {} {} {} fp={} {}",
                    ev.ts,
                    ev.actor.as_deref().unwrap_or("-"),
                    ev.action,
                    ev.target.as_deref().unwrap_or("-"),
                    ev.key_fingerprint.as_deref().unwrap_or("-"),
                    ev.metadata,
                );
            }
            Ok(())
        }
    }
}

async fn serve(settings: Settings) -> Result<()> {
    let bind = format!("{}:{}", settings.bind_host, settings.bind_port);
    let listener = tokio::net::TcpListener::bind(&bind).await?;
    let state = build_state(settings)?;
    let app = create_app(state);
    tracing::info!("dragoon-server listening on {bind}");
    axum::serve(listener, app).await?;
    Ok(())
}