porta-rs 0.1.1

Zero-trust CGNAT bypass via WireGuard tunnel
use anyhow::{Context, Result};
use clap::Args;
use std::fs;
use std::io::{self, Write};
use std::path::PathBuf;

use crate::config::{ClientConfig, ServerConfig, Role};
use crate::systemd::service::SystemdManager;
use crate::wireguard::keys;

#[derive(Args)]
pub struct SetupArgs {
    /// Role to configure: client (behind CGNAT) or server (public IP)
    #[arg(short, long)]
    pub role: Option<Role>,

    /// Skip systemd service installation
    #[arg(long)]
    pub no_systemd: bool,
}

pub fn execute(args: &SetupArgs) -> Result<()> {
    let role = match &args.role {
        Some(r) => *r,
        None => prompt_role()?,
    };

    match role {
        Role::Client => setup_client(args)?,
        Role::Server => setup_server(args)?,
    }

    Ok(())
}

fn prompt_role() -> Result<Role> {
    println!("\n\x1b[1;36m╔══════════════════════════════════════════════════════════════╗\x1b[0m");
    println!("\x1b[1;36m║                    PORTA SETUP WIZARD                        ║\x1b[0m");
    println!("\x1b[1;36m╚══════════════════════════════════════════════════════════════╝\x1b[0m\n");

    println!("Select the role for this machine:\n");
    println!("  \x1b[1;32m1)\x1b[0m Client - Machine behind CGNAT (no public IP)");
    println!("  \x1b[1;32m2)\x1b[0m Server - Machine with public IP\n");

    print!("Enter choice [1-2]: ");
    io::stdout().flush()?;

    let mut input = String::new();
    io::stdin().read_line(&mut input)?;

    match input.trim() {
        "1" => Ok(Role::Client),
        "2" => Ok(Role::Server),
        _ => Err(anyhow::anyhow!("Invalid choice. Enter 1 or 2.")),
    }
}

fn prompt_input(prompt: &str, default: &str) -> Result<String> {
    print!("{} [{}]: ", prompt, default);
    io::stdout().flush()?;

    let mut input = String::new();
    io::stdin().read_line(&mut input)?;

    let value = input.trim();
    if value.is_empty() {
        Ok(default.to_string())
    } else {
        Ok(value.to_string())
    }
}

fn setup_client(args: &SetupArgs) -> Result<()> {
    println!("\n\x1b[1;32m--- Client Setup (Machine behind CGNAT) ---\x1b[0m\n");

    let server_address = prompt_input("Server public IP address", "")?;
    if server_address.is_empty() {
        return Err(anyhow::anyhow!("Server address is required"));
    }

    let server_port: u16 = prompt_input("Server WireGuard port", "51820")?.parse()?;
    let local_ip = prompt_input("Local VPN IP", "10.0.0.2/32")?;
    let server_ip = prompt_input("Server VPN IP", "10.0.0.1/32")?;
    let client_id = prompt_input("Client identifier", "client-01")?;
    let interface = prompt_input("WireGuard interface name", "wg-porta")?;

    // Generate keys
    println!("\nGenerating WireGuard keypair...");
    let keypair = keys::generate_keypair()?;

    // Generate HMAC secret
    println!("Generating HMAC authentication secret...");
    let hmac_secret = keys::generate_hmac_secret();

    let config = ClientConfig {
        role: Role::Client,
        server_address,
        server_port,
        local_interface: interface.clone(),
        local_ip,
        server_ip,
        client_id: client_id.clone(),
        hmac_secret,
        public_key: keypair.public_key.clone(),
        ports: Vec::new(),
    };

    let config_dir = dirs::config_dir()
        .context("Failed to determine config directory")?
        .join("porta");

    fs::create_dir_all(&config_dir)?;

    let config_path = config_dir.join("client.toml");
    let toml_content = toml::to_string_pretty(&config)?;
    fs::write(&config_path, &toml_content)?;

    // Save private key separately with restricted permissions
    let key_path = config_dir.join("client.key");
    fs::write(&key_path, &keypair.private_key)?;

    println!("\n\x1b[1;32m✓ Configuration saved to {}\x1b[0m", config_path.display());
    println!("\x1b[1;32m✓ Private key saved to {}\x1b[0m", key_path.display());

    // Install systemd service
    if !args.no_systemd {
        match SystemdManager::new().install("client") {
            Ok(()) => {
                println!("\x1b[1;32m✓ Systemd service installed\x1b[0m");
            }
            Err(e) => {
                println!("\x1b[1;33m⚠ Failed to install systemd service: {}\x1b[0m", e);
                println!("  You can install it manually later with:");
                println!("  \x1b[1msudo systemctl enable --now porta-client\x1b[0m");
            }
        }
    }

    println!("\n\x1b[1;33mNext steps:\x1b[0m");
    println!("1. Copy this public key to the server's allowed_clients:");
    println!("   \x1b[1m{}\x1b[0m", keypair.public_key);
    println!("\n2. On the server, run:");
    println!("   \x1b[1mporta setup --role server\x1b[0m");
    println!("\n3. Add port forwarding rules:");
    println!("   \x1b[1mporta add-port --remote 80 --local 8080 --protocol tcp --desc \"Web\"\x1b[0m");
    println!("\n4. Start the service:");
    println!("   \x1b[1msudo systemctl enable --now porta-client\x1b[0m\n");

    Ok(())
}

fn setup_server(args: &SetupArgs) -> Result<()> {
    println!("\n\x1b[1;32m--- Server Setup (Machine with public IP) ---\x1b[0m\n");

    let listen_port: u16 = prompt_input("WireGuard listen port", "51820")?.parse()?;
    let interface = prompt_input("WireGuard interface name", "wg-porta")?;
    let ip = prompt_input("Server VPN IP with CIDR", "10.0.0.1/24")?;
    let data_dir = prompt_input("Data directory", "/var/lib/porta")?;

    // Generate keys
    println!("\nGenerating WireGuard keypair...");
    let keypair = keys::generate_keypair()?;

    let config = ServerConfig {
        role: Role::Server,
        listen_port,
        interface: interface.clone(),
        ip,
        data_dir,
        public_key: keypair.public_key.clone(),
        allowed_clients: Vec::new(),
    };

    let config_dir = PathBuf::from("/etc/porta");
    fs::create_dir_all(&config_dir)?;

    let config_path = config_dir.join("server.toml");
    let toml_content = toml::to_string_pretty(&config)?;
    fs::write(&config_path, &toml_content)?;

    // Save private key
    let key_path = config_dir.join("server.key");
    fs::write(&key_path, &keypair.private_key)?;

    println!("\n\x1b[1;32m✓ Configuration saved to {}\x1b[0m", config_path.display());
    println!("\x1b[1;32m✓ Private key saved to {}\x1b[0m", key_path.display());

    // Install systemd service
    if !args.no_systemd {
        match SystemdManager::new().install("server") {
            Ok(()) => {
                println!("\x1b[1;32m✓ Systemd service installed\x1b[0m");
            }
            Err(e) => {
                println!("\x1b[1;33m⚠ Failed to install systemd service: {}\x1b[0m", e);
                println!("  You can install it manually later with:");
                println!("  \x1b[1msudo systemctl enable --now porta-server\x1b[0m");
            }
        }
    }

    println!("\n\x1b[1;33mServer public key:\x1b[0m");
    println!("   \x1b[1m{}\x1b[0m", keypair.public_key);

    println!("\n\x1b[1;33mNext steps:\x1b[0m");
    println!("1. Add client public keys to the server config:");
    println!("   Edit /etc/porta/server.toml and add to allowed_clients");
    println!("\n2. Start the service:");
    println!("   \x1b[1msudo systemctl enable --now porta-server\x1b[0m\n");

    Ok(())
}