wgctrl 0.1.0

use std::{
    fmt,
    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
    ops::{Deref, DerefMut},
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use anyhow::{anyhow, Context, Result};
use base64::{engine::general_purpose, Engine as _};
use ipnet::IpNet;
use rand::Rng;
use rsln::types::message::RouteAttrIter;
use x25519_dalek::{PublicKey, StaticSecret};

use crate::constants::{WgAllowedIpAttr, WgDeviceAttr, WgPeerAttr};

/// `KEY_LEN` is the expected key length for a WireGuard key.
const KEY_LEN: usize = 32;
const ATTR_TYPE_MAST: u16 = 0x3fff;

/// `DeviceType` specifies the underlying implementation of a WireGuard device.
#[derive(Debug)]
pub enum DeviceType {
    LinuxKernel,
    Userspace,
    Unknown,
}

impl fmt::Display for DeviceType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DeviceType::LinuxKernel => write!(f, "linux-kernel"),
            DeviceType::Userspace => write!(f, "userspace"),
            DeviceType::Unknown => write!(f, "unknown"),
        }
    }
}

impl Default for DeviceType {
    fn default() -> Self {
        Self::LinuxKernel
    }
}

/// A `Key` is a public, private, or pre-shared secret key.
/// The Key constructor functions in this package can be used
/// to create Keys suitable for each of these applications.
#[derive(Default, PartialEq, Eq, Hash, Clone, Copy)]
pub struct Key([u8; KEY_LEN]);

impl fmt::Debug for Key {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Key([REDACTED])")
    }
}

impl TryFrom<&[u8]> for Key {
    type Error = anyhow::Error;

    fn try_from(bytes: &[u8]) -> Result<Self> {
        if bytes.len() != KEY_LEN {
            return Err(anyhow!("Incorrect key size: {}", bytes.len()));
        }

        let mut key = [0; KEY_LEN];
        key.copy_from_slice(bytes);

        Ok(Self(key))
    }
}

impl TryFrom<&str> for Key {
    type Error = anyhow::Error;

    fn try_from(s: &str) -> Result<Self> {
        let bytes = general_purpose::STANDARD.decode(s)?;
        Self::try_from(bytes.as_slice())
    }
}

impl Deref for Key {
    type Target = [u8; KEY_LEN];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Key {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl From<Key> for String {
    fn from(val: Key) -> Self {
        general_purpose::STANDARD.encode(*val)
    }
}

impl Key {
    /// `generate_key` generates a Key suitable for use as a pre-shared secret key from
    /// a cryptographically safe source.
    ///
    /// The output Key should not be used as a private key; use `generate_private_key` instead.
    pub fn generate_key() -> Result<Self> {
        let mut key = [0; KEY_LEN];
        let mut rng = rand::thread_rng();
        rng.fill(&mut key);
        Key::try_from(key.as_ref())
    }

    /// `generate_private_key` generates a Key suitable for use as a private key from a
    /// cryptographically safe source.
    pub fn generate_private_key() -> Result<Self> {
        let mut key = Key::generate_key()?;

        // Modify random bytes using algorithm described at:
        // https://cr.yp.to/ecdh.html.
        key[0] &= 248;
        key[31] &= 127;
        key[31] |= 64;
        Ok(key)
    }

    /// `public_key` computes a public key from the private key.
    pub fn public_key(&self) -> Key {
        let secret = StaticSecret::from(self.0);
        let public: PublicKey = (&secret).into();
        Self(*public.as_bytes())
    }

    /// `exchange` performs ECDH key exchange with a public key.
    pub fn exchange(&self, public_key: &Key) -> Key {
        let secret = StaticSecret::from(self.0);
        let public: PublicKey = (public_key.0).into();
        let shared = secret.diffie_hellman(&public);
        Self(*shared.as_bytes())
    }
}

/// `Peer` is a WireGuard peer to a Device.
#[derive(Default, Debug)]
pub struct Peer {
    /// `public_key` is the public key of a peer, computed from its private key.
    /// It is always present in a Peer.
    pub public_key: Key,

    /// `preshared_key` is an optional preshared key which may be used
    /// as an additional layer of security for peer communications.
    pub preshared_key: Option<Key>,

    /// `endpoint` is the most recent source address used for communication by this Peer.
    pub endpoint: Option<SocketAddr>,

    /// `persistent_keepalive_interval` specifies how often an "empty" packet is
    /// sent to a peer to keep a connection alive.
    pub persistent_keepalive_interval: Option<Duration>,

    /// `last_handshake_time` indicates the most recent time a handshake was performed with this peer.
    pub last_handshake_time: Option<SystemTime>,

    /// `rx_bytes` indicates the number of bytes received from this peer.
    pub rx_bytes: i64,

    /// `tx_bytes` indicates the number of bytes transmitted to this peer.
    pub tx_bytes: i64,

    /// `allowed_ips` specifies which IPv4 and IPv6 addresses this peer is allowed to communicate on.
    pub allowed_ips: Vec<IpNet>,

    /// `protocol_version` specifies which version of the WireGuard protocol is used for this Peer.
    pub protocol_version: Option<u16>,
}

impl TryFrom<&[u8]> for Peer {
    type Error = anyhow::Error;

    fn try_from(data: &[u8]) -> Result<Self> {
        let mut peer = Peer::default();

        for attr in RouteAttrIter::new(data) {
            let (kind, value) = attr?;
            let kind = kind & ATTR_TYPE_MAST;

            match kind {
                k if k == WgPeerAttr::PublicKey as u16 => {
                    peer.public_key = Key::try_from(value)?;
                }
                k if k == WgPeerAttr::PresharedKey as u16 => {
                    peer.preshared_key = Some(Key::try_from(value)?);
                }
                k if k == WgPeerAttr::Endpoint as u16 => {
                    peer.endpoint = Some(parse_sockaddr(value)?);
                }
                k if k == WgPeerAttr::PersistentKeepalive as u16 => {
                    let secs = parse_u16(value)?;
                    peer.persistent_keepalive_interval = if secs > 0 {
                        Some(Duration::from_secs(secs as u64))
                    } else {
                        None
                    };
                }
                k if k == WgPeerAttr::LastHandshakeTime as u16 => {
                    peer.last_handshake_time = parse_timespec(value);
                }
                k if k == WgPeerAttr::RxBytes as u16 => {
                    peer.rx_bytes = parse_u64(value)? as i64;
                }
                k if k == WgPeerAttr::TxBytes as u16 => {
                    peer.tx_bytes = parse_u64(value)? as i64;
                }
                k if k == WgPeerAttr::AllowedIps as u16 => {
                    for ip_attr in RouteAttrIter::new(value) {
                        let (_, ip_payload) = ip_attr?;
                        let allowed_ip = parse_allowed_ip(ip_payload)?;
                        peer.allowed_ips.push(allowed_ip);
                    }
                }
                k if k == WgPeerAttr::ProtocolVersion as u16 => {
                    peer.protocol_version = Some(parse_u32(value)? as u16);
                }
                _ => {}
            }
        }
        Ok(peer)
    }
}

/// `PeerConfig` is a WireGuard device peer configuration.
#[derive(Debug, Clone)]
pub struct PeerConfig {
    /// `public_key` specifies the public key of this peer.
    /// PublicKey is a mandatory field for all PeerConfigs.
    pub public_key: Key,

    /// `remove` specifies if the peer with this public key should be removed
    /// from a device's peer list.
    pub remove: bool,

    /// `update_only` specifies that an operation will only occur on this peer
    /// if the peer already exists as part of the interface.
    pub update_only: bool,

    /// `preshared_key` specifies a peer's preshared key configuration, if not none.
    pub preshared_key: Option<Key>,

    /// `endpoint` specifies the endpoint of this peer entry, if not none.
    pub endpoint: Option<SocketAddr>,

    /// `persistent_keepalive_interval` specifies the persistent keepalive interval
    /// for this peer, if not none.
    pub persistent_keepalive_interval: Option<Duration>,

    /// `replace_allowed_ips` specifies if the allowed IPs specified in this peer configuration
    /// should replace any existing ones, instead of appending them to the allowed IPs list.
    pub replace_allowed_ips: bool,

    /// `allowed_ips` specifies a list of allowed IP addresses in CIDR notation for this peer.
    pub allowed_ips: Vec<IpNet>,
}

/// `Device` is a WireGuard device.
#[derive(Default, Debug)]
pub struct Device {
    /// `name` is the name of the device.
    pub name: String,

    /// `device_type` specifies the underlying implementation of the device.
    pub device_type: DeviceType,

    /// `private_key` is the device's private key.
    pub private_key: Key,

    /// `public_key` is the device's public key, computed from its PrivateKey.
    pub public_key: Key,

    /// `listen_port` is the device's network listening port.
    pub listen_port: u16,

    /// `firewall_mark` is the device's current firewall mark.
    ///
    /// The firewall mark can be used in conjunction with firewall software to
    /// take action on outgoing WireGuard packets.
    pub firewall_mark: u32,

    /// `peers` is the list of network peers associated with this device.
    pub peers: Vec<Peer>,
}

impl TryFrom<&[u8]> for Device {
    type Error = anyhow::Error;

    fn try_from(payload: &[u8]) -> Result<Self> {
        if payload.len() < 4 {
            return Err(anyhow!("Short payload"));
        }

        let buf = &payload[4..]; // Skip Genl Header
        let mut dev = Device::default();

        for attr in RouteAttrIter::new(buf) {
            let (kind, value) = attr?;
            let kind = kind & ATTR_TYPE_MAST;

            match kind {
                k if k == WgDeviceAttr::IfName as u16 => {
                    dev.name = parse_string(value);
                }
                k if k == WgDeviceAttr::PrivateKey as u16 => {
                    dev.private_key = Key::try_from(value)?;
                }
                k if k == WgDeviceAttr::PublicKey as u16 => {
                    dev.public_key = Key::try_from(value)?;
                }
                k if k == WgDeviceAttr::ListenPort as u16 => {
                    dev.listen_port = parse_u16(value)?;
                }
                k if k == WgDeviceAttr::Fwmark as u16 => {
                    dev.firewall_mark = parse_u32(value)?;
                }
                k if k == WgDeviceAttr::Peers as u16 => {
                    for peer_attr in RouteAttrIter::new(value) {
                        let (_, peer_payload) = peer_attr?;
                        let peer = Peer::try_from(peer_payload)?;
                        dev.peers.push(peer);
                    }
                }
                _ => {}
            }
        }

        Ok(dev)
    }
}

/// `Config` is a WireGuard device configuration.
#[derive(Debug, Clone)]
pub struct Config {
    /// `private_key` specifies a private key configuration, if not none.
    pub private_key: Option<Key>,

    /// `listen_port` specifies a device's listening port, if not none.
    pub listen_port: Option<u16>,

    /// `firewall_mark` specifies a device's firewall mark, if not none.
    pub firewall_mark: Option<u32>,

    /// `replace_peers` specifies if the Peers in this configuration should replace
    /// the existing peer list, instead of appending them to the existing list.
    pub replace_peers: bool,

    /// `peers` specifies a list of peer configurations to apply to a device.
    pub peers: Vec<PeerConfig>,
}

/// C-style Null-terminated string parsing
fn parse_string(data: &[u8]) -> String {
    let trimmed = match data.iter().position(|&c| c == 0) {
        Some(pos) => &data[..pos],
        None => data,
    };
    String::from_utf8_lossy(trimmed).into_owned()
}

/// Parsing `struct sockaddr` (Handling IPv4/IPv6 & Endianness)
fn parse_sockaddr(data: &[u8]) -> Result<SocketAddr> {
    if data.len() < 2 {
        return Err(anyhow!("Sockaddr data too short"));
    }

    // Family is usually in Host Byte Order on Linux
    let family = u16::from_ne_bytes([data[0], data[1]]);

    match family as i32 {
        libc::AF_INET => {
            // struct sockaddr_in (IPv4)
            if data.len() < 8 {
                return Err(anyhow!("IPv4 sockaddr too short"));
            }

            let port = u16::from_be_bytes([data[2], data[3]]);
            let ip_bytes: [u8; 4] = data[4..8]
                .try_into()
                .map_err(|_| anyhow!("Failed to convert IPv4 address bytes"))?;

            Ok(SocketAddr::V4(SocketAddrV4::new(
                Ipv4Addr::from(ip_bytes),
                port,
            )))
        }
        libc::AF_INET6 => {
            // struct sockaddr_in6 (IPv6)
            if data.len() < 28 {
                return Err(anyhow!("IPv6 sockaddr too short"));
            }

            let port = u16::from_be_bytes([data[2], data[3]]);
            let ip_bytes: [u8; 16] = data[8..24]
                .try_into()
                .map_err(|_| anyhow!("Failed to convert IPv6 address bytes"))?;
            let scope_id = u32::from_ne_bytes([data[24], data[25], data[26], data[27]]);

            Ok(SocketAddr::V6(SocketAddrV6::new(
                Ipv6Addr::from(ip_bytes),
                port,
                0, // flowinfo is ignored
                scope_id,
            )))
        }
        _ => Err(anyhow!("Unsupported address family: {}", family)),
    }
}

/// Parsing Nested Attributes for AllowedIPs (Uses RouteAttrIter!)
fn parse_allowed_ip(data: &[u8]) -> Result<IpNet> {
    let mut ip_addr = None;
    let mut cidr = 0u8;

    for attr in RouteAttrIter::new(data) {
        let (raw_kind, value) = attr?;
        let kind = raw_kind & 0x3fff;

        match kind {
            k if k == WgAllowedIpAttr::IpAddr as u16 => {
                if value.len() == 4 {
                    let b: [u8; 4] = value
                        .try_into()
                        .map_err(|_| anyhow!("Invalid IPv4 length inside AllowedIPs"))?;
                    ip_addr = Some(std::net::IpAddr::V4(Ipv4Addr::from(b)));
                } else if value.len() == 16 {
                    let b: [u8; 16] = value
                        .try_into()
                        .map_err(|_| anyhow!("Invalid IPv6 length inside AllowedIPs"))?;
                    ip_addr = Some(std::net::IpAddr::V6(Ipv6Addr::from(b)));
                }
            }
            k if k == WgAllowedIpAttr::CidrMask as u16 => {
                if !value.is_empty() {
                    cidr = value[0];
                }
            }
            _ => {}
        }
    }

    let ip = ip_addr.ok_or_else(|| anyhow!("Missing IP address in AllowedIPs"))?;
    IpNet::new(ip, cidr).context("Invalid CIDR")
}

/// Parsing `struct timespec64`
/// Returns Option because failing to parse time is usually not critical enough to crash
fn parse_timespec(data: &[u8]) -> Option<SystemTime> {
    // Linux kernel sends __kernel_timespec (64-bit sec, 64-bit nsec)
    if data.len() < 16 {
        return None;
    }

    let mut sec_bytes = [0u8; 8];
    sec_bytes.copy_from_slice(&data[0..8]);
    let sec = i64::from_ne_bytes(sec_bytes);

    let mut nsec_bytes = [0u8; 8];
    nsec_bytes.copy_from_slice(&data[8..16]);
    let nsec = i64::from_ne_bytes(nsec_bytes);

    if sec == 0 && nsec == 0 {
        return None;
    }

    Some(UNIX_EPOCH + Duration::new(sec as u64, nsec as u32))
}

// Primitives (Using std, no extra deps, no unwrap)
fn parse_u16(data: &[u8]) -> Result<u16> {
    if data.len() < 2 {
        return Err(anyhow!("u16 data too short"));
    }
    Ok(u16::from_ne_bytes([data[0], data[1]]))
}

fn parse_u32(data: &[u8]) -> Result<u32> {
    if data.len() < 4 {
        return Err(anyhow!("u32 data too short"));
    }
    Ok(u32::from_ne_bytes([data[0], data[1], data[2], data[3]]))
}

fn parse_u64(data: &[u8]) -> Result<u64> {
    if data.len() < 8 {
        return Err(anyhow!("u64 data too short"));
    }
    Ok(u64::from_ne_bytes(
        data[0..8]
            .try_into()
            .map_err(|_| anyhow!("Failed to parse u64"))?,
    ))
}

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

    #[test]
    fn test_prepared_keys() {
        // Keys generated via "wg genkey" and "wg pubkey" for comparison
        // with this Rust implementation.
        let private = "GHuMwljFfqd2a7cs6BaUOmHflK23zME8VNvC5B37S3k=";
        let public = "aPxGwq8zERHQ3Q1cOZFdJ+cvJX5Ka4mLN38AyYKYF10=";

        let priv_key = Key::try_from(private).unwrap();
        let public_key = priv_key.public_key();

        assert_eq!(private, Into::<String>::into(priv_key));
        assert_eq!(public, Into::<String>::into(public_key));
    }

    #[test]
    fn test_key_exchange() {
        let alice = Key::generate_private_key().unwrap();
        let bob = Key::generate_private_key().unwrap();

        let alice_pub = alice.public_key();
        let bob_pub = bob.public_key();

        let alice_shared = alice.exchange(&bob_pub);
        let bob_shared = bob.exchange(&alice_pub);

        assert_eq!(*alice_shared, *bob_shared);
    }

    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_linux_integration_full() {
        use std::process::Command;

        // 0. Preparation: Check root privileges
        let euid = unsafe { libc::geteuid() };
        if euid != 0 {
            eprintln!("SKIPPING: Root privileges required for integration test");
            return;
        }

        let test_ifname = "wg_test_tmp";

        // Cleanup potential leftover from previous runs
        let _ = Command::new("ip").args(["link", "del", test_ifname]).output();

        // 1. Setup: Create real WireGuard interface for testing
        let status = Command::new("ip")
            .args(["link", "add", test_ifname, "type", "wireguard"])
            .status();

        if status.is_err() || !status.unwrap().success() {
            eprintln!("SKIPPING: Could not create wireguard interface (kernel module missing?)");
            return;
        }

        // Drop guard to ensure cleanup (RAII pattern)
        struct InterfaceGuard<'a>(&'a str);
        impl<'a> Drop for InterfaceGuard<'a> {
            fn drop(&mut self) {
                let _ = Command::new("ip")
                    .args(["link", "del", self.0])
                    .output();
            }
        }
        let _guard = InterfaceGuard(test_ifname);

        // 2. Execution: Create Client
        let mut client = crate::client::Client::new().expect("Failed to create Netlink Client");

        // 3. Verification A: Get existing device (should succeed)
        let device = client
            .get_device(test_ifname)
            .expect("Should find created device");
        assert_eq!(device.name, test_ifname);
        println!("Success: Found existing device '{}'", device.name);

        // 4. Verification B: Get non-existent device (should fail)
        let err = client.get_device("wg_imaginary_99");
        assert!(err.is_err());
        println!("Success: Correctly failed to find non-existent device");

        // 5. Verification C: Get non-WireGuard interface (should fail)
        // 'lo' (Loopback) always exists but is not WireGuard
        let err_loopback = client.get_device("lo");
        assert!(err_loopback.is_err());
        println!("Success: Correctly refused non-WireGuard interface 'lo'");
    }

    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_linux_client_is_permission() {
        // 1. Check if running as non-root
        let euid = unsafe { libc::geteuid() };
        if euid == 0 {
            println!("SKIPPING: Test must be run without elevated privileges (uid != 0)");
            return;
        }

        // 2. Create Client
        // Note: Client::new() attempts to resolve the wireguard family ID.
        // If this fails, we skip.
        let mut client = match crate::client::Client::new() {
            Ok(c) => c,
            Err(e) => {
                println!("Skipping: Failed to create client (generic netlink not available?): {}", e);
                return;
            }
        };

        // 3. Call get_device("wgnotexist0")
        // Check for permission denied as unprivileged user.
        let err = client.get_device("wgnotexist0");
        
        match err {
            Ok(_) => panic!("Expected error, got success"),
            Err(e) => {
                // Verify if it resembles a permission error.
                // In Rust with anyhow, we check the message or downcast.
                // Common permission errors: "Permission denied" (EACCES), "Operation not permitted" (EPERM).
                let msg = e.to_string();
                if msg.to_lowercase().contains("permission denied") || msg.to_lowercase().contains("operation not permitted") {
                    println!("Success: Got permission error: {}", msg);
                } else {
                    // Fail if it's not a permission error (e.g. just "Device not found")
                    // Note: If the kernel allows reading config by non-root users (some configs do),
                    // this test might naturally fail (finding "Device not found").
                    // But we follow the RIIR instruction.
                    panic!("expected permission denied, but got: {}", msg);
                }
            }
        }
    }

    #[cfg(target_os = "linux")]
    #[test]
    #[ignore]
    fn test_linux_client_devices_empty() {
        use std::process::Command;

        // 1. Check root privileges
        let euid = unsafe { libc::geteuid() };
        if euid != 0 {
            eprintln!("SKIPPING: Root privileges required");
            return;
        }

        // 2. Check for existing WireGuard interfaces to ensure clean state
        // This command returns output only if WG interfaces exist.
        let output = Command::new("ip")
            .args(["-o", "link", "show", "type", "wireguard"])
            .output()
            .expect("Failed to run ip link");

        if !output.stdout.is_empty() {
            eprintln!("SKIPPING: Existing WireGuard interfaces found. Cannot strictly test 'Empty' case without potentially affecting system state.");
            // We could run assertions that devices.len() == count, but the test name implies "Empty".
            return;
        }

        // 3. Create Client
        // Note: Client creation might fail if WireGuard module is not loaded/available.
        let mut client = match crate::client::Client::new() {
            Ok(c) => c,
            Err(e) => {
                eprintln!("Skipping: Failed to create client: {}", e);
                return;
            }
        };

        // 4. Call list_devices
        let devices = client.list_devices().expect("Failed to list devices");

        // 5. Assert Empty
        assert!(devices.is_empty(), "Expected no devices, got {}", devices.len());
        println!("Success: list_devices returned empty list as expected");
    }
}