rustunnel/src/generate.rs

use std::path::Path;

use anyhow::Context;
use rcgen::{
    BasicConstraints, Certificate, CertificateParams, CertifiedIssuer, DistinguishedName, DnType,
    IsCa, KeyPair, SanType, SerialNumber,
};

use crate::config::Config;
use crate::redact::Redacted;

#[derive(Debug, Clone)]
pub struct GeneratedMaterial {
    pub ca_cert_pem: String,
    pub ca_key_pem: String,
    pub server_cert_pem: String,
    pub server_key_pem: String,
    pub client_cert_pem: String,
    pub client_key_pem: String,
    pub auth_token: String,
}

/// Generate all credential material and write it to the output directory.
///
/// Creates:
/// - `ca.pem` — CA certificate (PEM)
/// - `ca.key` — CA private key (PEM, redacted in logs)
/// - `server.crt` — Server certificate signed by CA (PEM)
/// - `server.key` — Server private key (PEM, redacted in logs)
/// - `client.crt` — Client certificate signed by CA (PEM)
/// - `client.key` — Client private key (PEM, redacted in logs)
/// - `token.txt` — Auth/session token (redacted in logs)
/// - `config.json` — Configuration file with paths and defaults
pub fn generate(
    out_dir: &Path,
    ca_name: &str,
    server_name: &str,
    client_name: &str,
) -> anyhow::Result<()> {
    std::fs::create_dir_all(out_dir)
        .with_context(|| format!("failed to create output directory: {}", out_dir.display()))?;

    let material = generate_material(ca_name, server_name, client_name)?;
    let ca_cert_path = out_dir.join("ca.pem");
    let ca_key_path = out_dir.join("ca.key");
    std::fs::write(&ca_cert_path, &material.ca_cert_pem).with_context(|| {
        format!(
            "failed to write CA certificate to {}",
            ca_cert_path.display()
        )
    })?;
    std::fs::write(&ca_key_path, &material.ca_key_pem)
        .with_context(|| format!("failed to write CA key to {}", ca_key_path.display()))?;
    tracing::info!("Generated CA certificate at {}", ca_cert_path.display());
    tracing::info!("Generated CA key at {}", ca_key_path.display());

    let server_cert_path = out_dir.join("server.crt");
    let server_key_path = out_dir.join("server.key");
    std::fs::write(&server_cert_path, &material.server_cert_pem).with_context(|| {
        format!(
            "failed to write server cert to {}",
            server_cert_path.display()
        )
    })?;
    std::fs::write(&server_key_path, &material.server_key_pem).with_context(|| {
        format!(
            "failed to write server key to {}",
            server_key_path.display()
        )
    })?;
    tracing::info!(
        "Generated server certificate at {}",
        server_cert_path.display()
    );
    tracing::info!("Generated server key at {}", server_key_path.display());

    let client_cert_path = out_dir.join("client.crt");
    let client_key_path = out_dir.join("client.key");
    std::fs::write(&client_cert_path, &material.client_cert_pem).with_context(|| {
        format!(
            "failed to write client cert to {}",
            client_cert_path.display()
        )
    })?;
    std::fs::write(&client_key_path, &material.client_key_pem).with_context(|| {
        format!(
            "failed to write client key to {}",
            client_key_path.display()
        )
    })?;
    tracing::info!(
        "Generated client certificate at {}",
        client_cert_path.display()
    );
    tracing::info!("Generated client key at {}", client_key_path.display());

    let auth_token_redacted = Redacted::new(material.auth_token.clone());
    let token_path = out_dir.join("token.txt");
    std::fs::write(&token_path, &material.auth_token)
        .with_context(|| format!("failed to write token to {}", token_path.display()))?;
    tracing::info!(
        "Generated auth token at {} (value: {})",
        token_path.display(),
        auth_token_redacted
    );

    // Generate config
    let config = Config {
        listen_address: "0.0.0.0".to_string(),
        listen_port: 4180,
        socks_address: "127.0.0.1".to_string(),
        socks_port: 1180,
        ca_cert_path: "ca.pem".to_string(),
        server_cert_path: "server.crt".to_string(),
        server_key_path: "server.key".to_string(),
        client_cert_path: "client.crt".to_string(),
        client_key_path: "client.key".to_string(),
        auth_token_path: "token.txt".to_string(),
    };
    let config_path = out_dir.join("config.json");
    config
        .write_to(&config_path)
        .with_context(|| format!("failed to write config to {}", config_path.display()))?;
    tracing::info!("Generated config at {}", config_path.display());

    tracing::info!(
        "Generated credential material in {} ({} files)",
        out_dir.display(),
        list_generated_files(out_dir)
    );

    Ok(())
}

pub fn generate_material(
    ca_name: &str,
    server_name: &str,
    client_name: &str,
) -> anyhow::Result<GeneratedMaterial> {
    generate_material_with_server_host(ca_name, server_name, client_name, None)
}

pub fn generate_material_with_server_host(
    ca_name: &str,
    server_name: &str,
    client_name: &str,
    server_host: Option<&str>,
) -> anyhow::Result<GeneratedMaterial> {
    let ca = generate_ca(ca_name)?;
    let (server_cert, server_key) = generate_server_cert(&ca, server_name, server_host)?;
    let (client_cert, client_key) = generate_client_cert(&ca, client_name)?;
    Ok(GeneratedMaterial {
        ca_cert_pem: ca.as_ref().pem(),
        ca_key_pem: ca.key().serialize_pem(),
        server_cert_pem: server_cert.pem(),
        server_key_pem: server_key.serialize_pem(),
        client_cert_pem: client_cert.pem(),
        client_key_pem: client_key.serialize_pem(),
        auth_token: generate_auth_token(),
    })
}

fn generate_ca(cn: &str) -> anyhow::Result<CertifiedIssuer<'static, KeyPair>> {
    let mut params = CertificateParams::default();
    let mut dn = DistinguishedName::new();
    dn.push(DnType::CommonName, cn);
    dn.push(DnType::OrganizationName, "rustunnel");
    params.distinguished_name = dn;
    params.is_ca = IsCa::Ca(BasicConstraints::Constrained(0));
    params.serial_number = Some(SerialNumber::from(1));

    let key_pair = KeyPair::generate()?;
    let ca = CertifiedIssuer::self_signed(params, key_pair)?;
    Ok(ca)
}

fn generate_server_cert(
    ca: &CertifiedIssuer<'_, KeyPair>,
    cn: &str,
    extra_host: Option<&str>,
) -> anyhow::Result<(Certificate, KeyPair)> {
    let mut params = CertificateParams::default();
    let mut dn = DistinguishedName::new();
    dn.push(DnType::CommonName, cn);
    dn.push(DnType::OrganizationName, "rustunnel");
    params.distinguished_name = dn;
    params.is_ca = IsCa::NoCa;
    params.subject_alt_names = vec![
        SanType::DnsName("localhost".try_into().unwrap()),
        SanType::DnsName("rustunnel".try_into().unwrap()),
        SanType::IpAddress(std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1))),
    ];
    if let Some(host) = extra_host
        && host != "0.0.0.0"
        && host != "::"
    {
        if let Ok(ip) = host.parse::<std::net::IpAddr>() {
            params.subject_alt_names.push(SanType::IpAddress(ip));
        } else if let Ok(dns) = host.try_into() {
            params.subject_alt_names.push(SanType::DnsName(dns));
        }
    }
    params.serial_number = Some(SerialNumber::from(2));

    let key_pair = KeyPair::generate()?;
    let cert = params.signed_by(&key_pair, ca)?;
    Ok((cert, key_pair))
}

fn generate_client_cert(
    ca: &CertifiedIssuer<'_, KeyPair>,
    cn: &str,
) -> anyhow::Result<(Certificate, KeyPair)> {
    let mut params = CertificateParams::default();
    let mut dn = DistinguishedName::new();
    dn.push(DnType::CommonName, cn);
    dn.push(DnType::OrganizationName, "rustunnel");
    params.distinguished_name = dn;
    params.is_ca = IsCa::NoCa;
    params.serial_number = Some(SerialNumber::from(3));

    let key_pair = KeyPair::generate()?;
    let cert = params.signed_by(&key_pair, ca)?;
    Ok((cert, key_pair))
}

fn generate_auth_token() -> String {
    // Generate a 32-byte random token as hex
    // Use rand::random to avoid `gen` keyword conflict in Rust 2024
    let bytes: [u8; 32] = rand::random();
    hex::encode(bytes)
}

fn list_generated_files(dir: &Path) -> usize {
    std::fs::read_dir(dir)
        .map(|entries| entries.filter_map(|e| e.ok()).count())
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::redact;

    fn temp_dir() -> std::path::PathBuf {
        // Use process ID + monotonic counter to avoid collisions when tests run in parallel
        static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
        let unique = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        std::env::temp_dir().join(format!(
            "rustunnel_generate_test_{}_{}",
            std::process::id(),
            unique
        ))
    }

    #[test]
    fn generate_creates_expected_files() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        let expected_files = [
            "ca.pem",
            "ca.key",
            "server.crt",
            "server.key",
            "client.crt",
            "client.key",
            "token.txt",
            "config.json",
        ];

        for file in &expected_files {
            let path = dir.join(file);
            assert!(path.exists(), "Expected file not found: {}", file);
            let metadata = std::fs::metadata(&path).unwrap();
            assert!(metadata.len() > 0, "File is empty: {}", file);
        }

        // Clean up
        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_creates_valid_ca_cert() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        let ca_cert = std::fs::read_to_string(dir.join("ca.pem")).unwrap();
        assert!(ca_cert.contains("-----BEGIN CERTIFICATE-----"));
        assert!(ca_cert.contains("-----END CERTIFICATE-----"));

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_creates_private_keys() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        let ca_key = std::fs::read_to_string(dir.join("ca.key")).unwrap();
        assert!(
            redact::is_sensitive(&ca_key),
            "CA key should be detected as sensitive"
        );

        let server_key = std::fs::read_to_string(dir.join("server.key")).unwrap();
        assert!(
            redact::is_sensitive(&server_key),
            "Server key should be detected as sensitive"
        );

        let client_key = std::fs::read_to_string(dir.join("client.key")).unwrap();
        assert!(
            redact::is_sensitive(&client_key),
            "Client key should be detected as sensitive"
        );

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_token_is_non_empty() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        let token = std::fs::read_to_string(dir.join("token.txt")).unwrap();
        assert!(
            token.len() > 16,
            "Token should be at least 16 chars, got {}",
            token.len()
        );

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_config_is_valid_json() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        let config_str = std::fs::read_to_string(dir.join("config.json")).unwrap();
        let config: Config = serde_json::from_str(&config_str).unwrap();
        assert_eq!(config.listen_address, "0.0.0.0");
        assert_eq!(config.listen_port, 4180);
        assert_eq!(config.socks_address, "127.0.0.1");
        assert_eq!(config.socks_port, 1180);

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_no_hidden_files() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        // Ensure no hidden or stealth files are created
        for entry in std::fs::read_dir(&dir).unwrap() {
            let entry = entry.unwrap();
            let name = entry.file_name().to_string_lossy().to_string();
            assert!(!name.starts_with('.'), "Hidden file found: {}", name);
            assert!(!name.starts_with('~'), "Backup file found: {}", name);
        }

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_cross_platform_paths() {
        let dir = temp_dir();
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        let config_str = std::fs::read_to_string(dir.join("config.json")).unwrap();
        let config: Config = serde_json::from_str(&config_str).unwrap();

        // Config paths should be relative (cross-platform safe)
        assert!(!config.ca_cert_path.contains('/'));
        assert!(!config.server_cert_path.contains('/'));
        assert!(!config.client_cert_path.contains('/'));

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_custom_names() {
        let dir = temp_dir();
        generate(&dir, "custom-ca", "custom-server", "custom-client").unwrap();

        // Files should still be created with standard names
        assert!(dir.join("ca.pem").exists());
        assert!(dir.join("server.crt").exists());
        assert!(dir.join("client.crt").exists());

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_creates_nested_output_dir() {
        let dir = temp_dir().join("nested/deep/dir");
        generate(&dir, "test-ca", "test-server", "test-client").unwrap();

        assert!(dir.join("ca.pem").exists());
        assert!(dir.join("config.json").exists());

        // Clean up the full nested path
        if let Some(root) = dir
            .parent()
            .and_then(|p| p.parent())
            .and_then(|p| p.parent())
        {
            std::fs::remove_dir_all(root).ok();
        }
    }

    #[test]
    fn generated_certs_have_correct_cn() {
        let dir = temp_dir();
        generate(&dir, "my-ca", "my-server", "my-client").unwrap();

        // Verify certs contain expected CN
        let ca_cert = std::fs::read_to_string(dir.join("ca.pem")).unwrap();
        // The cert content will be base64-encoded, so we just verify structure
        assert!(ca_cert.contains("-----BEGIN CERTIFICATE-----"));

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn generate_idempotent() {
        let dir = temp_dir();
        // Generate twice should succeed
        generate(&dir, "ca1", "srv1", "cli1").unwrap();
        generate(&dir, "ca2", "srv2", "cli2").unwrap();

        // All files should still exist
        assert!(dir.join("ca.pem").exists());
        assert!(dir.join("config.json").exists());

        std::fs::remove_dir_all(&dir).ok();
    }
}