//! Host interface configuration
//!
//! Reference:
//! * WireGuard [Cross-platform Userspace Implementation](https://www.wireguard.com/xplatform/)

use std::{
    collections::HashMap,
    fmt,
    io::{self, BufRead, BufReader, Read},
    net::SocketAddr,
    str::FromStr,
    time::{Duration, SystemTime},
};

#[cfg(target_os = "linux")]
use netlink_packet_wireguard::WireguardAttribute;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

use super::{key::Key, peer::Peer};

#[cfg(target_os = "linux")]
const WGDEVICE_F_REPLACE_PEERS: u32 = 1;

/// WireGuard host representation.
#[derive(Clone, Default)]
#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
pub struct Host {
    pub listen_port: u16,
    pub private_key: Option<Key>,
    // FwMark with value of 0, removes the setting from WireGuard interface.
    pub(super) fwmark: Option<u32>,
    pub peers: HashMap<Key, Peer>,
}

// Implement `Debug` manually to avoid exposing private keys.
impl fmt::Debug for Host {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Host")
            .field("listen_port", &self.listen_port)
            .field("fwmark", &self.fwmark)
            .field("peers", &self.peers)
            .finish_non_exhaustive()
    }
}

impl Host {
    /// Create new `Host` with a given `listen_port` and `private_key`.
    #[must_use]
    pub fn new(listen_port: u16, private_key: Key) -> Self {
        Self {
            listen_port,
            private_key: Some(private_key),
            fwmark: None,
            peers: HashMap::new(),
        }
    }

    #[must_use]
    pub fn as_uapi(&self) -> String {
        let mut output = format!("listen_port={}\n", self.listen_port);
        if let Some(key) = &self.private_key {
            output.push_str("private_key=");
            output.push_str(&key.to_lower_hex());
            output.push('\n');
        }
        if let Some(fwmark) = &self.fwmark {
            output.push_str("fwmark=");
            output.push_str(&fwmark.to_string());
            output.push('\n');
        }
        output.push_str("replace_peers=true\n");
        for peer in self.peers.values() {
            output.push_str(peer.as_uapi_update().as_ref());
        }

        output
    }

    // TODO: use custom Error
    pub fn parse_uapi(buf: impl Read) -> io::Result<Self> {
        let reader = BufReader::new(buf);
        let mut host = Self::default();
        let mut peer_ref = None;

        for line_result in reader.lines() {
            let line = match line_result {
                Ok(line) => line,
                Err(err) => {
                    error!("Error parsing buffer line: {err}");
                    continue;
                }
            };
            if let Some((keyword, value)) = line.split_once('=') {
                match keyword {
                    "listen_port" => host.listen_port = value.parse().unwrap_or_default(),
                    "fwmark" => host.fwmark = value.parse().ok(),
                    "private_key" => host.private_key = Key::decode(value).ok(),
                    // "public_key" starts new peer definition
                    "public_key" => {
                        if let Ok(key) = Key::decode(value) {
                            let peer = Peer::new(key.clone());
                            host.peers.insert(key.clone(), peer);
                            peer_ref = host.peers.get_mut(&key);
                        } else {
                            peer_ref = None;
                        }
                    }
                    "preshared_key" => {
                        if let Some(ref mut peer) = peer_ref {
                            peer.preshared_key = Key::decode(value).ok();
                        }
                    }
                    "protocol_version" => {
                        if let Some(ref mut peer) = peer_ref {
                            peer.protocol_version = value.parse().ok();
                        }
                    }
                    "endpoint" => {
                        if let Some(ref mut peer) = peer_ref {
                            peer.endpoint = SocketAddr::from_str(value).ok();
                        }
                    }
                    "persistent_keepalive_interval" => {
                        if let Some(ref mut peer) = peer_ref {
                            peer.persistent_keepalive_interval = value.parse().ok();
                        }
                    }
                    "allowed_ip" => {
                        if let Some(ref mut peer) = peer_ref {
                            if let Ok(addr) = value.parse() {
                                peer.allowed_ips.push(addr);
                            }
                        }
                    }
                    "last_handshake_time_sec" => {
                        if let Some(ref mut peer) = peer_ref {
                            let handshake =
                                peer.last_handshake.get_or_insert(SystemTime::UNIX_EPOCH);
                            *handshake += Duration::from_secs(value.parse().unwrap_or_default());
                        }
                    }
                    "last_handshake_time_nsec" => {
                        if let Some(ref mut peer) = peer_ref {
                            let handshake =
                                peer.last_handshake.get_or_insert(SystemTime::UNIX_EPOCH);
                            *handshake += Duration::from_nanos(value.parse().unwrap_or_default());
                        }
                    }
                    "rx_bytes" => {
                        if let Some(ref mut peer) = peer_ref {
                            peer.rx_bytes = value.parse().unwrap_or_default();
                        }
                    }
                    "tx_bytes" => {
                        if let Some(ref mut peer) = peer_ref {
                            peer.tx_bytes = value.parse().unwrap_or_default();
                        }
                    }
                    // "errno" ends config
                    "errno" => {
                        if let Ok(errno) = value.parse::<u32>() {
                            if errno == 0 {
                                // Break here, or BufReader will wait for EOF.
                                break;
                            }
                        }
                        return Err(io::Error::other("error reading UAPI"));
                    }
                    _ => error!("Unknown UAPI keyword {keyword}"),
                }
            }
        }

        Ok(host)
    }
}

#[cfg(target_os = "linux")]
impl Host {
    pub(crate) fn append_nlas(&mut self, nlas: &[WireguardAttribute]) {
        for nla in nlas {
            match nla {
                WireguardAttribute::PrivateKey(value) => self.private_key = Some(Key::new(*value)),
                WireguardAttribute::ListenPort(value) => self.listen_port = *value,
                WireguardAttribute::Fwmark(value) => self.fwmark = Some(*value),
                WireguardAttribute::Peers(nlas) => {
                    for nla in nlas {
                        let peer = Peer::from_nlas(nla);
                        // On tunnels with a lot of routes, the peer info may be split across multiple
                        // entries. We need to merge them here.
                        // https://github.com/DefGuard/client/issues/617
                        // https://github.com/DefGuard/wireguard-rs/issues/101
                        if let Some(existing_peer) = self.peers.get_mut(&peer.public_key) {
                            existing_peer.rx_bytes += peer.rx_bytes;
                            existing_peer.tx_bytes += peer.tx_bytes;
                            existing_peer.last_handshake =
                                match (existing_peer.last_handshake, peer.last_handshake) {
                                    (Some(x), Some(y)) => Some(x.max(y)),
                                    (Some(x), None) => Some(x),
                                    (None, Some(y)) => Some(y),
                                    (None, None) => None,
                                };
                            if peer.preshared_key.is_some() {
                                existing_peer.preshared_key = peer.preshared_key;
                            }
                            if peer.protocol_version.is_some() {
                                existing_peer.protocol_version = peer.protocol_version;
                            }
                            if peer.endpoint.is_some() {
                                existing_peer.endpoint = peer.endpoint;
                            }
                            if peer.persistent_keepalive_interval.is_some() {
                                existing_peer.persistent_keepalive_interval =
                                    peer.persistent_keepalive_interval;
                            }
                            existing_peer.allowed_ips.extend(peer.allowed_ips);
                        } else {
                            self.peers.insert(peer.public_key.clone(), peer);
                        }
                    }
                }
                _ => (),
            }
        }
    }

    #[must_use]
    pub(crate) fn as_nlas(&self, ifname: &str) -> Vec<WireguardAttribute> {
        let mut nlas = vec![
            WireguardAttribute::IfName(ifname.into()),
            WireguardAttribute::ListenPort(self.listen_port),
        ];
        if let Some(key) = &self.private_key {
            nlas.push(WireguardAttribute::PrivateKey(key.as_array()));
        }
        if let Some(fwmark) = &self.fwmark {
            nlas.push(WireguardAttribute::Fwmark(*fwmark));
        }
        nlas.push(WireguardAttribute::Flags(WGDEVICE_F_REPLACE_PEERS));

        // IMPORTANT: To avoid buffer overflow, do not add peers here.
        // let peers = self.peers.values().map(Peer::as_nlas_peer).collect();
        // nlas.push(WireguardAttribute::Peers(peers));

        nlas
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use super::*;

    #[test]
    fn test_parse_config() {
        let uapi_output =
            b"private_key=000102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            listen_port=7301\n\
            public_key=100102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            preshared_key=0000000000000000000000000000000000000000000000000000000000000000\n\
            protocol_version=1\n\
            last_handshake_time_sec=0\n\
            last_handshake_time_nsec=0\n\
            tx_bytes=0\n\
            rx_bytes=0\n\
            persistent_keepalive_interval=0\n\
            allowed_ip=10.6.0.12/32\n\
            public_key=200102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            preshared_key=0000000000000000000000000000000000000000000000000000000000000000\n\
            protocol_version=1\n\
            endpoint=83.11.218.160:51421\n\
            last_handshake_time_sec=1654631933\n\
            last_handshake_time_nsec=862977251\n\
            tx_bytes=52759980\n\
            rx_bytes=3683056\n\
            persistent_keepalive_interval=0\n\
            allowed_ip=10.6.0.25/32\n\
            public_key=300102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            preshared_key=0000000000000000000000000000000000000000000000000000000000000000\n\
            protocol_version=1\n\
            endpoint=31.135.163.194:37712\n\
            last_handshake_time_sec=1654776419\n\
            last_handshake_time_nsec=732507856\n\
            tx_bytes=1009094476\n\
            rx_bytes=76734328\n\
            persistent_keepalive_interval=0\n\
            allowed_ip=10.6.0.23/32\n\
            errno=0\n";
        let buf = Cursor::new(uapi_output);
        let host = Host::parse_uapi(buf).unwrap();
        assert_eq!(host.listen_port, 7301);
        assert_eq!(host.peers.len(), 3);
    }

    #[test]
    fn test_host_uapi() {
        let key_str = "000102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f";
        let key = Key::decode(key_str).unwrap();

        let host = Host::new(12345, key);
        assert_eq!(
            "listen_port=12345\n\
            private_key=000102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            replace_peers=true\n",
            host.as_uapi()
        );
    }

    #[test]
    fn test_peer_uapi() {
        let key_str = "000102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f";
        let key = Key::decode(key_str).unwrap();

        let peer = Peer::new(key);
        assert_eq!(
            "public_key=000102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            replace_allowed_ips=true\n",
            peer.as_uapi_update()
        );

        let key_str = "00112233445566778899aaabbcbddeeff0e1d2c3b4a5968778695a4b3c2d1e0f";
        let key = Key::decode(key_str).unwrap();
        let peer = Peer::new(key);
        assert_eq!(
            "public_key=00112233445566778899aaabbcbddeeff0e1d2c3b4a5968778695a4b3c2d1e0f\n\
            remove=true\n",
            peer.as_uapi_remove()
        );
    }

    #[test]
    fn dg25_28_test_dont_expose_preshared_keys() {
        let preshared_key_str = "000102030405060708090a0b0c0d0e0ff0e1d2c3b4a5968778695a4b3c2d1e0f";
        let peer = Peer {
            public_key: Key::decode(
                "286ac5ff9b2f900259008172225da774031e8a3689d8f341667be157b2336970",
            )
            .unwrap(),
            preshared_key: Some(Key::decode(preshared_key_str).unwrap()),
            protocol_version: None,
            endpoint: None,
            last_handshake: None,
            tx_bytes: 0,
            rx_bytes: 0,
            persistent_keepalive_interval: None,
            allowed_ips: Vec::new(),
        };

        let debug = format!("{peer:?}");
        assert!(!debug.contains("preshared_key"));
        assert!(!debug.contains(preshared_key_str));
    }
}
