nx9-wg/crates/nx9-wg-api/tests/test_wiregui_capabilities.rs

use chrono::{Duration, Utc};
use nx9_wg_api::{AppState, DiagnosticsService, IpAllocator, ReconciliationEngine};
use nx9_wg_core::types::diagnostics::DiagnosticSubsystem;
use nx9_wg_core::types::firewall::{
    FirewallAction, FirewallDirection, FirewallProtocol, FirewallRule,
};
use nx9_wg_core::types::network::Network;
use nx9_wg_core::types::wireguard::{
    Interface, Peer, PeerProfile, PeerState, PeerType, WireGuardPrivateKey, WireGuardPublicKey,
};
use nx9_wg_db::Store;
use nx9_wg_network::{NetworkEngine, SimulatedNetworkEngine};
use nx9_wireguard::{SimulatedWireGuardEngine, WireGuardEngine};
use std::net::IpAddr;
use std::sync::Arc;
use uuid::Uuid;

async fn setup_test_context() -> (
    AppState,
    Arc<SimulatedWireGuardEngine>,
    Arc<SimulatedNetworkEngine>,
    Arc<ReconciliationEngine>,
) {
    let store = Store::connect_in_memory().await.expect("connect DB");
    store.migrate().await.expect("migrate DB");

    let state = AppState::new(store);
    let wg_engine = Arc::new(SimulatedWireGuardEngine::new());
    let net_engine = Arc::new(SimulatedNetworkEngine::new());
    let reconciler = Arc::new(ReconciliationEngine::new(
        state.clone(),
        wg_engine.clone(),
        net_engine.clone(),
    ));

    (state, wg_engine, net_engine, reconciler)
}

#[tokio::test]
async fn test_automatic_ip_allocation() {
    let (state, _, _, _) = setup_test_context().await;
    let now = Utc::now().naive_utc();

    let net_id = Uuid::new_v4();
    let network = Network {
        id: net_id,
        name: "Test-V4-Subnet".to_string(),
        cidr: "10.50.0.0/24".parse().unwrap(),
        enabled: true,
        description: None,
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_network(&network)
        .await
        .expect("create net");

    let iface_id = Uuid::new_v4();
    let iface = Interface {
        id: iface_id,
        name: "wg50".to_string(),
        private_key: WireGuardPrivateKey::new(
            "cGFzc3dvcmRkZXZlbG9wbWVudGtleTEyMzQ1Njc4OTAxMg==".to_string(),
        ),
        public_key: WireGuardPublicKey::new(
            "cHVibGlja2V5ZGV2ZWxvcG1lbnRrZXkxMjM0NTY3ODkwMTI=".to_string(),
        ),
        listen_port: 51850,
        address_v4: "10.50.0.1/24".parse().unwrap(),
        address_v6: None,
        mtu: Some(1420),
        dns: None,
        enabled: true,
        pre_up: None,
        post_up: None,
        pre_down: None,
        post_down: None,
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_interface(&iface)
        .await
        .expect("create iface");

    // First allocation: 10.50.0.1 is interface -> next available is 10.50.0.2/32
    let ip1 = IpAllocator::allocate_next_ip(&state.store, &network, Some(&iface), None)
        .await
        .expect("allocate ip1");
    assert_eq!(ip1.to_string(), "10.50.0.2/32");

    // Create a peer with this allocated IP
    let peer1 = Peer {
        id: Uuid::new_v4(),
        interface_id: iface_id,
        name: "peer-1".to_string(),
        peer_type: PeerType::RoadWarrior,
        state: PeerState::Active,
        public_key: WireGuardPublicKey::new(
            "peer1pubkey12345678901234567890123456789012=".to_string(),
        ),
        private_key: None,
        preshared_key: None,
        endpoint: None,
        allowed_ips: ip1.to_string(),
        server_allowed_ips: None,
        address_v4: Some(ip1),
        address_v6: None,
        dns: None,
        mtu: None,
        persistent_keepalive: None,
        profile: PeerProfile::FullTunnel,
        expires_at: None,
        last_handshake_at: None,
        created_at: now,
        updated_at: now,
    };
    state.store.create_peer(&peer1).await.expect("create peer1");

    // Second allocation: next should be 10.50.0.3/32
    let ip2 = IpAllocator::allocate_next_ip(&state.store, &network, Some(&iface), None)
        .await
        .expect("allocate ip2");
    assert_eq!(ip2.to_string(), "10.50.0.3/32");

    // List available IPs: first should be 10.50.0.3
    let available = IpAllocator::list_available_ips(&state.store, &network, Some(&iface), 5)
        .await
        .expect("list available");
    assert_eq!(available.len(), 5);
    assert_eq!(available[0], "10.50.0.3".parse::<IpAddr>().unwrap());
    assert_eq!(available[1], "10.50.0.4".parse::<IpAddr>().unwrap());

    // List allocations: should show peer1
    let allocs = IpAllocator::list_allocations(&state.store, &network)
        .await
        .expect("list allocs");
    assert_eq!(allocs.len(), 1);
    assert_eq!(allocs[0].ip_address, "10.50.0.2/32");
    assert_eq!(allocs[0].peer_name.as_deref(), Some("peer-1"));
}

#[tokio::test]
async fn test_peer_expiration_lifecycle() {
    let (state, wg_engine, _net_engine, reconciler) = setup_test_context().await;
    let now = Utc::now().naive_utc();

    let iface_id = Uuid::new_v4();
    let iface = Interface {
        id: iface_id,
        name: "wg60".to_string(),
        private_key: WireGuardPrivateKey::new(
            "cGFzc3dvcmRkZXZlbG9wbWVudGtleTEyMzQ1Njc4OTAxMg==".to_string(),
        ),
        public_key: WireGuardPublicKey::new(
            "cHVibGlja2V5ZGV2ZWxvcG1lbnRrZXkxMjM0NTY3ODkwMTI=".to_string(),
        ),
        listen_port: 51860,
        address_v4: "10.60.0.1/24".parse().unwrap(),
        address_v6: None,
        mtu: Some(1420),
        dns: None,
        enabled: true,
        pre_up: None,
        post_up: None,
        pre_down: None,
        post_down: None,
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_interface(&iface)
        .await
        .expect("create iface");

    // Peer with expiration in the past
    let expired_peer_id = Uuid::new_v4();
    let expired_peer = Peer {
        id: expired_peer_id,
        interface_id: iface_id,
        name: "expired-peer".to_string(),
        peer_type: PeerType::RoadWarrior,
        state: PeerState::Active, // marked active initially
        public_key: WireGuardPublicKey::new(
            "expiredpubkey123456789012345678901234567890=".to_string(),
        ),
        private_key: None,
        preshared_key: None,
        endpoint: None,
        allowed_ips: "10.60.0.5/32".to_string(),
        server_allowed_ips: None,
        address_v4: Some("10.60.0.5/32".parse().unwrap()),
        address_v6: None,
        dns: None,
        mtu: None,
        persistent_keepalive: None,
        profile: PeerProfile::FullTunnel,
        expires_at: Some(now - Duration::hours(1)), // expired 1 hour ago
        last_handshake_at: None,
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_peer(&expired_peer)
        .await
        .expect("create peer");

    // Active peer without expiration
    let active_peer_id = Uuid::new_v4();
    let active_peer = Peer {
        id: active_peer_id,
        interface_id: iface_id,
        name: "active-peer".to_string(),
        peer_type: PeerType::RoadWarrior,
        state: PeerState::Active,
        public_key: WireGuardPublicKey::new(
            "activepubkey1234567890123456789012345678901=".to_string(),
        ),
        private_key: None,
        preshared_key: None,
        endpoint: None,
        allowed_ips: "10.60.0.6/32".to_string(),
        server_allowed_ips: None,
        address_v4: Some("10.60.0.6/32".parse().unwrap()),
        address_v6: None,
        dns: None,
        mtu: None,
        persistent_keepalive: None,
        profile: PeerProfile::FullTunnel,
        expires_at: Some(now + Duration::days(30)),
        last_handshake_at: None,
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_peer(&active_peer)
        .await
        .expect("create peer");

    // Run reconciliation sweep
    let swept = reconciler.sweep_expired_peers().await.expect("sweep");
    assert_eq!(swept, 1);

    // Verify expired_peer transitioned to Expired
    let p1 = state
        .store
        .get_peer(expired_peer_id)
        .await
        .expect("get")
        .unwrap();
    assert_eq!(p1.state, PeerState::Expired);

    // Verify active_peer remains Active
    let p2 = state
        .store
        .get_peer(active_peer_id)
        .await
        .expect("get")
        .unwrap();
    assert_eq!(p2.state, PeerState::Active);

    // Reconcile apply ensures only active peers are synced to WireGuard kernel engine
    let rep = reconciler.apply().await.expect("apply");
    assert!(rep.success);

    let stats = wg_engine
        .get_interface_stats("wg60")
        .await
        .unwrap()
        .unwrap();
    // Only active peer should be live in interface
    assert_eq!(stats.peers.len(), 1);
    assert_eq!(stats.peers[0].public_key, active_peer.public_key.as_str());
}

#[tokio::test]
async fn test_peer_firewall_and_port_ranges() {
    let (state, _, net_engine, reconciler) = setup_test_context().await;
    let now = Utc::now().naive_utc();

    let iface_id = Uuid::new_v4();
    let iface = Interface {
        id: iface_id,
        name: "wg70".to_string(),
        private_key: WireGuardPrivateKey::new(
            "cGFzc3dvcmRkZXZlbG9wbWVudGtleTEyMzQ1Njc4OTAxMg==".to_string(),
        ),
        public_key: WireGuardPublicKey::new(
            "cHVibGlja2V5ZGV2ZWxvcG1lbnRrZXkxMjM0NTY3ODkwMTI=".to_string(),
        ),
        listen_port: 51870,
        address_v4: "10.70.0.1/24".parse().unwrap(),
        address_v6: None,
        mtu: Some(1420),
        dns: None,
        enabled: true,
        pre_up: None,
        post_up: None,
        pre_down: None,
        post_down: None,
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_interface(&iface)
        .await
        .expect("create iface");

    let peer_id = Uuid::new_v4();
    let peer = Peer {
        id: peer_id,
        interface_id: iface_id,
        name: "dev-peer".to_string(),
        peer_type: PeerType::RoadWarrior,
        state: PeerState::Active,
        public_key: WireGuardPublicKey::new(
            "devpeerpubkey1234567890123456789012345678901=".to_string(),
        ),
        private_key: None,
        preshared_key: None,
        endpoint: None,
        allowed_ips: "10.70.0.10/32".to_string(),
        server_allowed_ips: None,
        address_v4: Some("10.70.0.10/32".parse().unwrap()),
        address_v6: None,
        dns: None,
        mtu: None,
        persistent_keepalive: None,
        profile: PeerProfile::FullTunnel,
        expires_at: None,
        last_handshake_at: None,
        created_at: now,
        updated_at: now,
    };
    state.store.create_peer(&peer).await.expect("create peer");

    // Peer-specific rule with multi-port and TCP/UDP protocol
    let rule = FirewallRule {
        id: Uuid::new_v4(),
        name: "Allow Dev Ports".to_string(),
        interface_id: Some(iface_id),
        peer_id: Some(peer_id),
        direction: FirewallDirection::Forward,
        action: FirewallAction::Accept,
        protocol: FirewallProtocol::TcpUdp,
        source: None,
        destination: None,
        source_port: None,
        destination_port: None,
        port_range: Some("8000-8100".to_string()),
        priority: 10,
        enabled: true,
        description: Some("Peer port range".to_string()),
        created_at: now,
        updated_at: now,
    };
    state
        .store
        .create_firewall_rule(&rule)
        .await
        .expect("create rule");

    // Apply reconciliation to compile ruleset
    reconciler.apply().await.expect("apply");

    let ruleset = net_engine
        .get_active_nftables_ruleset()
        .await
        .expect("get ruleset");
    assert!(ruleset.contains("table inet nx9_wg"));
    // Resolved peer IP 10.70.0.10, protocol meta l4proto { tcp, udp }, and port range 8000-8100
    assert!(ruleset.contains("ip saddr 10.70.0.10"));
    assert!(ruleset.contains("meta l4proto { tcp, udp }"));
    assert!(ruleset.contains("th dport 8000-8100 accept"));
}

#[tokio::test]
async fn test_native_diagnostics_subsystem() {
    let (state, wg_engine, net_engine, reconciler) = setup_test_context().await;
    let diag = DiagnosticsService::new(state, wg_engine, net_engine, reconciler);

    let all_reports = diag.diagnose_all().await.expect("diagnose all");
    assert!(!all_reports.is_empty());

    let sys_report = diag
        .run_diagnostic(DiagnosticSubsystem::System, None)
        .await
        .expect("diag system");
    assert_eq!(sys_report.len(), 1);
    assert_eq!(sys_report[0].subsystem, "system");

    let fwd_report = diag
        .run_diagnostic(DiagnosticSubsystem::Forwarding, None)
        .await
        .expect("diag fwd");
    assert_eq!(fwd_report.len(), 1);
    assert_eq!(fwd_report[0].subsystem, "forwarding");
}