archy/core/archipelago/src/vpn.rs

//! VPN integration: Tailscale and WireGuard management.
//!
//! Manages VPN connections, generates WireGuard configs, and monitors
//! VPN interface status for remote access to the Archipelago node.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;
use tokio::fs;

const VPN_CONFIG_FILE: &str = "vpn-config.json";

/// VPN provider type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum VpnProvider {
    Tailscale,
    Wireguard,
}

/// Persisted VPN configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VpnConfig {
    pub provider: VpnProvider,
    pub enabled: bool,
    #[serde(default)]
    pub tailscale_auth_key: Option<String>,
    #[serde(default)]
    pub wireguard_config: Option<WireGuardConfig>,
    #[serde(default)]
    pub configured_at: Option<String>,
}

/// WireGuard configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WireGuardConfig {
    pub private_key: String,
    pub public_key: String,
    pub address: String,
    pub dns: String,
    #[serde(default)]
    pub peers: Vec<WireGuardPeer>,
}

/// A WireGuard peer.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WireGuardPeer {
    pub public_key: String,
    pub endpoint: String,
    pub allowed_ips: String,
    #[serde(default)]
    pub persistent_keepalive: Option<u16>,
}

/// Current VPN status (gathered at runtime, not persisted).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VpnStatus {
    pub connected: bool,
    pub provider: Option<String>,
    pub interface: Option<String>,
    pub ip_address: Option<String>,
    pub hostname: Option<String>,
    #[serde(default)]
    pub peers_connected: u32,
    #[serde(default)]
    pub bytes_in: u64,
    #[serde(default)]
    pub bytes_out: u64,
}

impl Default for VpnConfig {
    fn default() -> Self {
        Self {
            provider: VpnProvider::Tailscale,
            enabled: false,
            tailscale_auth_key: None,
            wireguard_config: None,
            configured_at: None,
        }
    }
}

pub async fn load_config(data_dir: &Path) -> Result<VpnConfig> {
    let path = data_dir.join(VPN_CONFIG_FILE);
    if !path.exists() {
        return Ok(VpnConfig::default());
    }
    let content = fs::read_to_string(&path)
        .await
        .context("Failed to read VPN config")?;
    let config: VpnConfig = serde_json::from_str(&content).unwrap_or_default();
    Ok(config)
}

pub async fn save_config(data_dir: &Path, config: &VpnConfig) -> Result<()> {
    fs::create_dir_all(data_dir).await.context("Failed to create data dir")?;
    let content = serde_json::to_string_pretty(config).context("Failed to serialize VPN config")?;
    fs::write(data_dir.join(VPN_CONFIG_FILE), content)
        .await
        .context("Failed to write VPN config")?;
    Ok(())
}

/// Generate a WireGuard keypair using the `wg` command.
pub async fn generate_wireguard_keypair() -> Result<(String, String)> {
    let privkey_output = tokio::process::Command::new("wg")
        .arg("genkey")
        .output()
        .await
        .context("Failed to run wg genkey — is wireguard-tools installed?")?;

    if !privkey_output.status.success() {
        anyhow::bail!(
            "wg genkey failed: {}",
            String::from_utf8_lossy(&privkey_output.stderr)
        );
    }

    let private_key = String::from_utf8(privkey_output.stdout)
        .context("Invalid UTF-8 from wg genkey")?
        .trim()
        .to_string();

    let mut child = tokio::process::Command::new("wg")
        .arg("pubkey")
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .context("Failed to spawn wg pubkey")?;

    if let Some(mut stdin) = child.stdin.take() {
        use tokio::io::AsyncWriteExt;
        stdin.write_all(private_key.as_bytes()).await
            .context("Failed to write private key to wg stdin")?;
    }

    let output = child.wait_with_output().await
        .context("wg pubkey process failed")?;

    if !output.status.success() {
        anyhow::bail!("wg pubkey failed: {}", String::from_utf8_lossy(&output.stderr));
    }

    let public_key = String::from_utf8(output.stdout)
        .context("wg pubkey output is not valid UTF-8")?
        .trim()
        .to_string();

    Ok((private_key, public_key))
}

/// Generate a WireGuard configuration file content.
pub fn generate_wireguard_conf(config: &WireGuardConfig) -> String {
    let mut conf = format!(
        "[Interface]\nPrivateKey = {}\nAddress = {}\nDNS = {}\n",
        config.private_key, config.address, config.dns
    );

    for peer in &config.peers {
        conf.push_str(&format!(
            "\n[Peer]\nPublicKey = {}\nEndpoint = {}\nAllowedIPs = {}\n",
            peer.public_key, peer.endpoint, peer.allowed_ips
        ));
        if let Some(ka) = peer.persistent_keepalive {
            conf.push_str(&format!("PersistentKeepalive = {}\n", ka));
        }
    }

    conf
}

/// Get the current VPN status by checking network interfaces.
pub async fn get_status() -> VpnStatus {
    // Check for Tailscale interface
    if let Ok(tailscale) = get_tailscale_status().await {
        return tailscale;
    }

    // Check for WireGuard interface
    if let Ok(wg) = get_wireguard_status().await {
        return wg;
    }

    VpnStatus {
        connected: false,
        provider: None,
        interface: None,
        ip_address: None,
        hostname: None,
        peers_connected: 0,
        bytes_in: 0,
        bytes_out: 0,
    }
}

async fn get_tailscale_status() -> Result<VpnStatus> {
    // Check if tailscale0 interface exists
    let output = tokio::process::Command::new("ip")
        .args(["addr", "show", "tailscale0"])
        .output()
        .await
        .context("Failed to check tailscale0")?;

    if !output.status.success() {
        anyhow::bail!("No tailscale0 interface");
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let ip = stdout
        .lines()
        .find(|l| l.contains("inet ") && !l.contains("inet6"))
        .and_then(|l| {
            l.split_whitespace()
                .nth(1)
                .map(|ip| ip.split('/').next().unwrap_or(ip).to_string())
        });

    // Try to get hostname from tailscale status
    let hostname = tokio::process::Command::new("sh")
        .arg("-c")
        .arg("podman exec tailscale tailscale status --self --json 2>/dev/null | grep -o '\"DNSName\":\"[^\"]*\"' | head -1 | cut -d'\"' -f4")
        .output()
        .await
        .ok()
        .and_then(|o| {
            let s = String::from_utf8_lossy(&o.stdout).trim().to_string();
            if s.is_empty() { None } else { Some(s.trim_end_matches('.').to_string()) }
        });

    // Get peer count
    let peers = tokio::process::Command::new("sh")
        .arg("-c")
        .arg("podman exec tailscale tailscale status 2>/dev/null | grep -c 'active' || echo 0")
        .output()
        .await
        .ok()
        .and_then(|o| {
            String::from_utf8_lossy(&o.stdout)
                .trim()
                .parse::<u32>()
                .ok()
        })
        .unwrap_or(0);

    Ok(VpnStatus {
        connected: ip.is_some(),
        provider: Some("tailscale".to_string()),
        interface: Some("tailscale0".to_string()),
        ip_address: ip,
        hostname,
        peers_connected: peers,
        bytes_in: 0,
        bytes_out: 0,
    })
}

async fn get_wireguard_status() -> Result<VpnStatus> {
    let output = tokio::process::Command::new("ip")
        .args(["addr", "show", "wg0"])
        .output()
        .await
        .context("Failed to check wg0")?;

    if !output.status.success() {
        anyhow::bail!("No wg0 interface");
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let ip = stdout
        .lines()
        .find(|l| l.contains("inet ") && !l.contains("inet6"))
        .and_then(|l| {
            l.split_whitespace()
                .nth(1)
                .map(|ip| ip.split('/').next().unwrap_or(ip).to_string())
        });

    // Get peer count from wg show
    let peers = tokio::process::Command::new("sh")
        .arg("-c")
        .arg("wg show wg0 peers 2>/dev/null | wc -l")
        .output()
        .await
        .ok()
        .and_then(|o| {
            String::from_utf8_lossy(&o.stdout)
                .trim()
                .parse::<u32>()
                .ok()
        })
        .unwrap_or(0);

    // Get transfer stats
    let (bytes_in, bytes_out) = tokio::process::Command::new("sh")
        .arg("-c")
        .arg("wg show wg0 transfer 2>/dev/null | awk '{i+=$2; o+=$3} END {print i, o}'")
        .output()
        .await
        .ok()
        .and_then(|o| {
            let s = String::from_utf8_lossy(&o.stdout);
            let mut parts = s.trim().split_whitespace();
            let i = parts.next().and_then(|v| v.parse::<u64>().ok()).unwrap_or(0);
            let o = parts.next().and_then(|v| v.parse::<u64>().ok()).unwrap_or(0);
            Some((i, o))
        })
        .unwrap_or((0, 0));

    Ok(VpnStatus {
        connected: ip.is_some(),
        provider: Some("wireguard".to_string()),
        interface: Some("wg0".to_string()),
        ip_address: ip,
        hostname: None,
        peers_connected: peers,
        bytes_in,
        bytes_out,
    })
}

/// Configure Tailscale with an auth key (triggers tailscale up).
pub async fn configure_tailscale(auth_key: &str, data_dir: &Path) -> Result<()> {
    let mut config = load_config(data_dir).await?;
    config.provider = VpnProvider::Tailscale;
    config.enabled = true;
    config.tailscale_auth_key = Some(auth_key.to_string());
    config.configured_at = Some(chrono::Utc::now().to_rfc3339());
    save_config(data_dir, &config).await?;

    // Run tailscale up with auth key in the container
    let output = tokio::process::Command::new("podman")
        .args([
            "exec", "tailscale", "tailscale", "up",
            "--authkey", auth_key,
            "--accept-routes",
        ])
        .output()
        .await
        .context("Failed to run tailscale up")?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        anyhow::bail!("tailscale up failed: {}", stderr);
    }

    Ok(())
}

/// Configure WireGuard with generated keys and optional peer.
pub async fn configure_wireguard(
    data_dir: &Path,
    address: &str,
    dns: &str,
    peer: Option<WireGuardPeer>,
) -> Result<WireGuardConfig> {
    let (private_key, public_key) = generate_wireguard_keypair().await?;

    let wg_config = WireGuardConfig {
        private_key: private_key.clone(),
        public_key: public_key.clone(),
        address: address.to_string(),
        dns: dns.to_string(),
        peers: peer.map_or_else(Vec::new, |p| vec![p]),
    };

    let mut config = load_config(data_dir).await?;
    config.provider = VpnProvider::Wireguard;
    config.enabled = true;
    config.wireguard_config = Some(wg_config.clone());
    config.configured_at = Some(chrono::Utc::now().to_rfc3339());
    save_config(data_dir, &config).await?;

    // Write WireGuard config file
    let conf_content = generate_wireguard_conf(&wg_config);
    let wg_dir = data_dir.join("wireguard");
    fs::create_dir_all(&wg_dir).await.context("Failed to create wireguard dir")?;
    fs::write(wg_dir.join("wg0.conf"), &conf_content)
        .await
        .context("Failed to write wg0.conf")?;

    Ok(wg_config)
}

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

    #[test]
    fn test_vpn_config_default() {
        let config = VpnConfig::default();
        assert!(!config.enabled);
        assert_eq!(config.provider, VpnProvider::Tailscale);
        assert!(config.tailscale_auth_key.is_none());
    }

    #[test]
    fn test_vpn_config_serialization() {
        let config = VpnConfig {
            provider: VpnProvider::Wireguard,
            enabled: true,
            tailscale_auth_key: None,
            wireguard_config: Some(WireGuardConfig {
                private_key: "privkey".to_string(),
                public_key: "pubkey".to_string(),
                address: "10.0.0.1/24".to_string(),
                dns: "1.1.1.1".to_string(),
                peers: vec![],
            }),
            configured_at: Some("2026-01-01T00:00:00Z".to_string()),
        };
        let json = serde_json::to_string(&config).unwrap();
        let parsed: VpnConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.provider, VpnProvider::Wireguard);
        assert!(parsed.enabled);
        assert!(parsed.wireguard_config.is_some());
    }

    #[test]
    fn test_generate_wireguard_conf() {
        let config = WireGuardConfig {
            private_key: "test_privkey".to_string(),
            public_key: "test_pubkey".to_string(),
            address: "10.0.0.2/24".to_string(),
            dns: "1.1.1.1, 8.8.8.8".to_string(),
            peers: vec![WireGuardPeer {
                public_key: "peer_pubkey".to_string(),
                endpoint: "vpn.example.com:51820".to_string(),
                allowed_ips: "0.0.0.0/0".to_string(),
                persistent_keepalive: Some(25),
            }],
        };
        let conf = generate_wireguard_conf(&config);
        assert!(conf.contains("[Interface]"));
        assert!(conf.contains("PrivateKey = test_privkey"));
        assert!(conf.contains("Address = 10.0.0.2/24"));
        assert!(conf.contains("[Peer]"));
        assert!(conf.contains("PublicKey = peer_pubkey"));
        assert!(conf.contains("PersistentKeepalive = 25"));
    }

    #[tokio::test]
    async fn test_load_config_default_when_no_file() {
        let dir = tempfile::tempdir().unwrap();
        let config = load_config(dir.path()).await.unwrap();
        assert!(!config.enabled);
    }

    #[tokio::test]
    async fn test_save_and_load_config_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let config = VpnConfig {
            provider: VpnProvider::Tailscale,
            enabled: true,
            tailscale_auth_key: Some("tskey-auth-test".to_string()),
            wireguard_config: None,
            configured_at: Some("2026-03-10T00:00:00Z".to_string()),
        };
        save_config(dir.path(), &config).await.unwrap();
        let loaded = load_config(dir.path()).await.unwrap();
        assert!(loaded.enabled);
        assert_eq!(loaded.tailscale_auth_key, Some("tskey-auth-test".to_string()));
    }

    #[test]
    fn test_vpn_status_default() {
        let status = VpnStatus {
            connected: false,
            provider: None,
            interface: None,
            ip_address: None,
            hostname: None,
            peers_connected: 0,
            bytes_in: 0,
            bytes_out: 0,
        };
        assert!(!status.connected);
    }
}