AWG2_harness/traffic_core/src/wg.rs

use std::net::Ipv4Addr;

use crate::{
    config,
    netns::{DockerExecTransport, NsTransport, NsenterTransport, nsenter_available},
};

#[derive(Debug, Clone)]
pub struct Peer {
    pub ipv4: Ipv4Addr,
}

#[derive(Debug, Clone)]
pub struct PeerRuntime {
    pub public_key: String,
    pub allowed_ips: Vec<String>,
    pub endpoint: Option<String>,
    pub latest_handshake_epoch: Option<u64>,
    pub transfer_rx_bytes: Option<u64>,
    pub transfer_tx_bytes: Option<u64>,
}

pub fn discover_peers(pid: u32, container: &str, peer_cidr: &str) -> anyhow::Result<Vec<Peer>> {
    let boxed: Box<dyn NsTransport> = if nsenter_available() {
        Box::new(NsenterTransport { pid })
    } else {
        Box::new(DockerExecTransport {
            container: container.to_string(),
        })
    };

    discover_peers_with_transport(boxed.as_ref(), peer_cidr)
}

pub fn discover_peers_with_transport(
    transport: &dyn NsTransport,
    peer_cidr: &str,
) -> anyhow::Result<Vec<Peer>> {
    let mut peers = match transport.run(&["wg", "show", "all", "allowed-ips"]) {
        Ok(primary) => parse_allowed_ips(&primary, peer_cidr)?,
        Err(_) => Vec::new(),
    };
    if peers.is_empty() {
        let fallback = transport
            .run(&["wg", "show"])
            .map_err(|e| anyhow::anyhow!("wg show fallback failed: {e}"))?;
        peers = parse_fallback_wg_show(&fallback, peer_cidr)?;
    }

    peers.sort_by_key(|p| p.ipv4.to_string());
    peers.dedup_by_key(|p| p.ipv4);
    Ok(peers)
}

pub fn discover_peer_runtime(
    pid: u32,
    container: &str,
    peer_cidr: &str,
) -> anyhow::Result<Vec<PeerRuntime>> {
    let boxed: Box<dyn NsTransport> = if nsenter_available() {
        Box::new(NsenterTransport { pid })
    } else {
        Box::new(DockerExecTransport {
            container: container.to_string(),
        })
    };
    discover_peer_runtime_with_transport(boxed.as_ref(), peer_cidr)
}

pub fn discover_peer_runtime_with_transport(
    transport: &dyn NsTransport,
    peer_cidr: &str,
) -> anyhow::Result<Vec<PeerRuntime>> {
    let raw = transport
        .run(&["wg", "show", "all", "dump"])
        .map_err(|e| anyhow::anyhow!("wg show all dump failed: {e}"))?;
    parse_wg_show_dump_runtime_peers(&raw, peer_cidr)
}

fn parse_allowed_ips(raw: &str, peer_cidr: &str) -> anyhow::Result<Vec<Peer>> {
    let mut out = Vec::new();
    for line in raw.lines() {
        let cols: Vec<&str> = line.split_whitespace().collect();
        if cols.len() < 3 {
            continue;
        }
        for tok in cols[2].split(',') {
            let tok = tok.trim();
            if let Some((ip, pfx)) = tok.split_once('/') {
                if pfx != "32" {
                    continue;
                }
                if let Ok(v4) = ip.parse::<Ipv4Addr>() {
                    if config::ip_in_cidr(v4, peer_cidr)? {
                        out.push(Peer { ipv4: v4 });
                    }
                }
            }
        }
    }
    Ok(out)
}

fn parse_fallback_wg_show(raw: &str, peer_cidr: &str) -> anyhow::Result<Vec<Peer>> {
    let mut out = Vec::new();
    for line in raw.lines() {
        let t = line.trim();
        if !t.starts_with("allowed ips:") {
            continue;
        }
        let rest = t.trim_start_matches("allowed ips:").trim();
        for tok in rest.split(',') {
            let tok = tok.trim();
            if let Some((ip, pfx)) = tok.split_once('/') {
                if pfx != "32" {
                    continue;
                }
                if let Ok(v4) = ip.parse::<Ipv4Addr>() {
                    if config::ip_in_cidr(v4, peer_cidr)? {
                        out.push(Peer { ipv4: v4 });
                    }
                }
            }
        }
    }
    Ok(out)
}

fn parse_wg_show_dump_runtime_peers(
    raw: &str,
    peer_cidr: &str,
) -> anyhow::Result<Vec<PeerRuntime>> {
    let mut out = Vec::new();
    for line in raw.lines() {
        let cols: Vec<&str> = line.split('\t').collect();
        if cols.len() < 8 {
            continue;
        }

        // peer lines include endpoint/allowed_ips/handshake/rx/tx; interface lines do not match this contract
        let public_key = cols.get(1).map(|s| s.trim()).unwrap_or_default();
        let allowed = cols.get(4).map(|s| s.trim()).unwrap_or_default();
        if public_key.is_empty() || allowed.is_empty() || allowed == "(none)" {
            continue;
        }

        let filtered_allowed = filter_allowed_ips_in_cidr(allowed, peer_cidr)?;
        if filtered_allowed.is_empty() {
            continue;
        }

        let endpoint = cols
            .get(3)
            .map(|s| s.trim())
            .filter(|s| !s.is_empty() && *s != "(none)")
            .map(|s| s.to_string());

        let latest_handshake_epoch = cols
            .get(5)
            .and_then(|s| s.trim().parse::<u64>().ok())
            .filter(|v| *v > 0);
        let transfer_rx_bytes = cols.get(6).and_then(|s| s.trim().parse::<u64>().ok());
        let transfer_tx_bytes = cols.get(7).and_then(|s| s.trim().parse::<u64>().ok());

        out.push(PeerRuntime {
            public_key: public_key.to_string(),
            allowed_ips: filtered_allowed,
            endpoint,
            latest_handshake_epoch,
            transfer_rx_bytes,
            transfer_tx_bytes,
        });
    }
    Ok(out)
}

fn filter_allowed_ips_in_cidr(raw: &str, peer_cidr: &str) -> anyhow::Result<Vec<String>> {
    let mut out = Vec::new();
    for tok in raw.split(',') {
        let tok = tok.trim();
        let Some((ip, pfx)) = tok.split_once('/') else {
            continue;
        };
        if pfx != "32" {
            continue;
        }
        let Ok(v4) = ip.parse::<Ipv4Addr>() else {
            continue;
        };
        if config::ip_in_cidr(v4, peer_cidr)? {
            out.push(v4.to_string());
        }
    }
    out.sort();
    out.dedup();
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;
    use std::sync::Mutex;

    struct TestTransport {
        replies: Mutex<BTreeMap<String, anyhow::Result<String>>>,
    }

    impl TestTransport {
        fn new(entries: Vec<(&str, anyhow::Result<&str>)>) -> Self {
            let mut replies = BTreeMap::new();
            for (cmd, out) in entries {
                let mapped = out.map(|s| s.to_string());
                replies.insert(cmd.to_string(), mapped);
            }
            Self {
                replies: Mutex::new(replies),
            }
        }
    }

    impl crate::netns::NsTransport for TestTransport {
        fn run(&self, args: &[&str]) -> anyhow::Result<String> {
            let key = args.join(" ");
            let mut lock = self.replies.lock().expect("lock");
            lock.remove(&key)
                .unwrap_or_else(|| Err(anyhow::anyhow!("missing stub for: {key}")))
        }
    }

    #[test]
    fn parse_primary() {
        let raw = "wg0 key1 10.8.1.10/32\nwg0 key2 10.8.1.2/32";
        let peers = parse_allowed_ips(raw, "10.8.1.0/24").expect("parse");
        assert_eq!(peers.len(), 2);
    }

    #[test]
    fn parse_fallback() {
        let raw = "allowed ips: 10.8.1.4/32, 10.8.1.5/32\nallowed ips: 192.168.0.1/32";
        let peers = parse_fallback_wg_show(raw, "10.8.1.0/24").expect("parse");
        assert_eq!(peers.len(), 2);
    }

    #[test]
    fn cidr_filtering_keeps_only_target_subnet() {
        let raw = "wg0 key1 10.8.1.10/32\nwg0 key2 10.9.1.2/32";
        let peers = parse_allowed_ips(raw, "10.8.1.0/24").expect("parse");
        assert_eq!(peers.len(), 1);
        assert_eq!(peers[0].ipv4.to_string(), "10.8.1.10");
    }

    #[test]
    fn lexicographic_sorting_parity() {
        let raw = "wg0 key1 10.8.1.2/32\nwg0 key2 10.8.1.10/32";
        let mut peers = parse_allowed_ips(raw, "10.8.1.0/24").expect("parse");
        peers.sort_by_key(|p| p.ipv4.to_string());
        assert_eq!(peers[0].ipv4.to_string(), "10.8.1.10");
        assert_eq!(peers[1].ipv4.to_string(), "10.8.1.2");
    }

    #[test]
    fn primary_failure_falls_back_to_wg_show() {
        let transport = TestTransport::new(vec![
            ("wg show all allowed-ips", Err(anyhow::anyhow!("wg broken"))),
            ("wg show", Ok("allowed ips: 10.8.1.4/32")),
        ]);

        let peers = discover_peers_with_transport(&transport, "10.8.1.0/24").expect("peers");
        assert_eq!(peers.len(), 1);
        assert_eq!(peers[0].ipv4.to_string(), "10.8.1.4");
    }

    #[test]
    fn empty_primary_output_falls_back_to_wg_show() {
        let transport = TestTransport::new(vec![
            ("wg show all allowed-ips", Ok("")),
            ("wg show", Ok("allowed ips: 10.8.1.4/32, 10.8.1.5/32")),
        ]);

        let peers = discover_peers_with_transport(&transport, "10.8.1.0/24").expect("peers");
        assert_eq!(peers.len(), 2);
        assert_eq!(peers[0].ipv4.to_string(), "10.8.1.4");
        assert_eq!(peers[1].ipv4.to_string(), "10.8.1.5");
    }

    #[test]
    fn parses_wg_show_dump_runtime_peers() {
        let raw = "wg0\tpriv\tpub\t51820\toff\nwg0\tpubkey1\t(none)\t1.2.3.4:51820\t10.8.1.10/32\t1715000000\t123\t456\t0\nwg0\tpubkey2\t(none)\t(none)\t10.8.1.27/32,192.168.1.5/32\t0\t0\t0\t0";
        let peers = parse_wg_show_dump_runtime_peers(raw, "10.8.1.0/24").expect("parse");
        assert_eq!(peers.len(), 2);
        assert_eq!(peers[0].allowed_ips, vec!["10.8.1.10".to_string()]);
        assert_eq!(peers[1].allowed_ips, vec!["10.8.1.27".to_string()]);
        assert_eq!(peers[1].latest_handshake_epoch, None);
    }

    #[test]
    fn maps_allowed_ip_to_runtime_peer() {
        let raw = "wg0\tpubkey1\t(none)\t1.2.3.4:51820\t10.8.1.10/32,10.8.1.11/32\t1715000000\t123\t456\t0";
        let peers = parse_wg_show_dump_runtime_peers(raw, "10.8.1.0/24").expect("parse");
        assert_eq!(peers.len(), 1);
        assert_eq!(peers[0].allowed_ips.len(), 2);
    }

    #[test]
    fn propagates_wg_dump_command_failure() {
        let transport = TestTransport::new(vec![(
            "wg show all dump",
            Err(anyhow::anyhow!("wg broken")),
        )]);
        let err = discover_peer_runtime_with_transport(&transport, "10.8.1.0/24")
            .err()
            .expect("err");
        assert!(err.to_string().contains("wg show all dump failed"));
    }
}