use anyhow::{Result, bail};
use std::net::IpAddr;

/// The only link-layer type the IP peek understands. Anything else matches
/// nothing -- see `matches`.
pub const LINKTYPE_ETHERNET: u32 = 1;

const ETH_HEADER_LEN: usize = 14;
const ETHERTYPE_IPV4: u16 = 0x0800;
const ETHERTYPE_IPV6: u16 = 0x86DD;
const ETHERTYPE_VLAN: u16 = 0x8100;
const ETHERTYPE_QINQ: u16 = 0x88A8;
const ETHERTYPE_QINQ_ALT: u16 = 0x9100;
/// 802.1Q and one layer of QinQ. Deeper stacking is real but vanishingly
/// rare in capture files, and the bound keeps a corrupt/truncated packet
/// from walking off into the payload looking for an ethertype.
const MAX_VLAN_TAGS: usize = 2;

const IPV4_HEADER_MIN: usize = 20;
const IPV6_HEADER_LEN: usize = 40;

/// A set of IP addresses/prefixes to keep. A packet matches if *either* its
/// source or destination address falls in *any* entry.
///
/// Sized for the expected 1-4 entries: two plain vectors, split by address
/// family so a v4 packet never compares against a v6 entry, scanned
/// linearly. No sorting (a prefix match isn't an equality test, so ordering
/// buys nothing) and no hashing.
#[derive(Debug, Clone)]
pub struct IpFilter {
    /// Set when no `--ip` was given: every packet is kept, without looking
    /// at it at all. This is the default so that the filtering code path is
    /// the *only* code path -- see the chunker.
    match_all: bool,
    /// (network, mask) pairs, network already masked at parse time.
    v4: Vec<(u32, u32)>,
    v6: Vec<(u128, u128)>,
}

impl IpFilter {
    /// The no-`--ip` default: keeps everything, matches without inspecting.
    pub fn match_all() -> Self {
        Self { match_all: true, v4: Vec::new(), v6: Vec::new() }
    }

    /// Parses `--ip` specs: a bare address (`1.1.1.1`, `2606:4700::1111`)
    /// or a CIDR prefix (`10.0.0.0/8`). An empty list yields `match_all`.
    pub fn parse(specs: &[String]) -> Result<Self> {
        if specs.is_empty() {
            return Ok(Self::match_all());
        }
        let mut f = Self { match_all: false, v4: Vec::new(), v6: Vec::new() };
        for spec in specs {
            let (addr_str, prefix_str) = match spec.split_once('/') {
                Some((a, p)) => (a, Some(p)),
                None => (spec.as_str(), None),
            };
            let addr: IpAddr = addr_str
                .parse()
                .map_err(|_| anyhow::anyhow!("--ip {spec:?}: not an IP address"))?;
            let width = if addr.is_ipv4() { 32u32 } else { 128u32 };
            let prefix = match prefix_str {
                Some(p) => p
                    .parse::<u32>()
                    .map_err(|_| anyhow::anyhow!("--ip {spec:?}: prefix length is not a number"))?,
                None => width,
            };
            if prefix > width {
                bail!("--ip {spec:?}: prefix length {prefix} exceeds {width} bits");
            }
            match addr {
                IpAddr::V4(v4) => {
                    let mask = if prefix == 0 { 0 } else { u32::MAX << (32 - prefix) };
                    f.v4.push((u32::from(v4) & mask, mask));
                }
                IpAddr::V6(v6) => {
                    let mask = if prefix == 0 { 0 } else { u128::MAX << (128 - prefix) };
                    f.v6.push((u128::from(v6) & mask, mask));
                }
            }
        }
        Ok(f)
    }

    /// Whether this filter keeps everything, i.e. no `--ip` was given.
    /// Callers use it to skip work (and warnings) that only a real filter
    /// needs.
    pub fn is_match_all(&self) -> bool {
        self.match_all
    }

    /// Whether to keep `packet` -- the captured bytes of one pcap record,
    /// starting at the link-layer header.
    ///
    /// Everything unmatched is dropped: non-IP frames (ARP, LLDP), packets
    /// truncated by snaplen before the addresses, and any link type other
    /// than Ethernet.
    pub fn matches(&self, packet: &[u8], linktype: u32) -> bool {
        if self.match_all {
            return true;
        }
        if linktype != LINKTYPE_ETHERNET || packet.len() < ETH_HEADER_LEN {
            return false;
        }
        let mut ethertype = be16(&packet[12..14]);
        let mut off = ETH_HEADER_LEN;
        let mut tags = 0;
        // A VLAN tag is TPID(2) + TCI(2); the TPID is the ethertype we just
        // read, so the *next* type field sits 2 bytes into the tag.
        while matches!(ethertype, ETHERTYPE_VLAN | ETHERTYPE_QINQ | ETHERTYPE_QINQ_ALT) {
            tags += 1;
            if tags > MAX_VLAN_TAGS || packet.len() < off + 4 {
                return false;
            }
            ethertype = be16(&packet[off + 2..off + 4]);
            off += 4;
        }
        match ethertype {
            ETHERTYPE_IPV4 => {
                if packet.len() < off + IPV4_HEADER_MIN || packet[off] >> 4 != 4 {
                    return false;
                }
                let src = be32(&packet[off + 12..off + 16]);
                let dst = be32(&packet[off + 16..off + 20]);
                self.v4.iter().any(|&(net, mask)| src & mask == net || dst & mask == net)
            }
            ETHERTYPE_IPV6 => {
                if packet.len() < off + IPV6_HEADER_LEN || packet[off] >> 4 != 6 {
                    return false;
                }
                let src = be128(&packet[off + 8..off + 24]);
                let dst = be128(&packet[off + 24..off + 40]);
                self.v6.iter().any(|&(net, mask)| src & mask == net || dst & mask == net)
            }
            _ => false,
        }
    }
}

fn be16(b: &[u8]) -> u16 {
    u16::from_be_bytes(b.try_into().unwrap())
}

fn be32(b: &[u8]) -> u32 {
    u32::from_be_bytes(b.try_into().unwrap())
}

fn be128(b: &[u8]) -> u128 {
    u128::from_be_bytes(b.try_into().unwrap())
}

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

    fn specs(list: &[&str]) -> Vec<String> {
        list.iter().map(|s| s.to_string()).collect()
    }

    fn filter(list: &[&str]) -> IpFilter {
        IpFilter::parse(&specs(list)).unwrap()
    }

    /// An Ethernet frame with `vlan_tags` 802.1Q tags in front of `ethertype`.
    fn eth(vlan_tags: usize, ethertype: u16, payload: &[u8]) -> Vec<u8> {
        let mut f = vec![0u8; 12]; // dst + src MAC
        for _ in 0..vlan_tags {
            f.extend_from_slice(&ETHERTYPE_VLAN.to_be_bytes());
            f.extend_from_slice(&[0x00, 0x64]); // TCI: VLAN 100
        }
        f.extend_from_slice(&ethertype.to_be_bytes());
        f.extend_from_slice(payload);
        f
    }

    fn ipv4(src: &str, dst: &str) -> Vec<u8> {
        let mut p = vec![0u8; IPV4_HEADER_MIN];
        p[0] = 0x45; // version 4, IHL 5
        let s: std::net::Ipv4Addr = src.parse().unwrap();
        let d: std::net::Ipv4Addr = dst.parse().unwrap();
        p[12..16].copy_from_slice(&s.octets());
        p[16..20].copy_from_slice(&d.octets());
        p
    }

    fn ipv6(src: &str, dst: &str) -> Vec<u8> {
        let mut p = vec![0u8; IPV6_HEADER_LEN];
        p[0] = 0x60; // version 6
        let s: std::net::Ipv6Addr = src.parse().unwrap();
        let d: std::net::Ipv6Addr = dst.parse().unwrap();
        p[8..24].copy_from_slice(&s.octets());
        p[24..40].copy_from_slice(&d.octets());
        p
    }

    fn v4_packet(src: &str, dst: &str) -> Vec<u8> {
        eth(0, ETHERTYPE_IPV4, &ipv4(src, dst))
    }

    #[test]
    fn match_all_keeps_everything_including_non_ip_and_odd_linktypes() {
        let f = IpFilter::match_all();
        assert!(f.is_match_all());
        assert!(f.matches(&[], 113));
        assert!(f.matches(&eth(0, 0x0806, &[1, 2, 3]), LINKTYPE_ETHERNET));
        // An empty --ip list is the same thing.
        assert!(IpFilter::parse(&[]).unwrap().is_match_all());
    }

    #[test]
    fn exact_v4_matches_either_direction_only() {
        let f = filter(&["1.1.1.1"]);
        assert!(!f.is_match_all());
        assert!(f.matches(&v4_packet("1.1.1.1", "9.9.9.9"), LINKTYPE_ETHERNET));
        assert!(f.matches(&v4_packet("9.9.9.9", "1.1.1.1"), LINKTYPE_ETHERNET));
        assert!(!f.matches(&v4_packet("9.9.9.9", "8.8.8.8"), LINKTYPE_ETHERNET));
        // Neighbouring address must not match a bare (/32) entry.
        assert!(!f.matches(&v4_packet("1.1.1.2", "9.9.9.9"), LINKTYPE_ETHERNET));
    }

    #[test]
    fn exact_v6_matches_either_direction() {
        let f = filter(&["2606:4700::1111"]);
        let pkt = eth(0, ETHERTYPE_IPV6, &ipv6("2001:db8::1", "2606:4700::1111"));
        assert!(f.matches(&pkt, LINKTYPE_ETHERNET));
        let miss = eth(0, ETHERTYPE_IPV6, &ipv6("2001:db8::1", "2606:4700::1112"));
        assert!(!f.matches(&miss, LINKTYPE_ETHERNET));
    }

    #[test]
    fn cidr_boundaries_are_exact() {
        let f = filter(&["10.0.0.0/8"]);
        assert!(f.matches(&v4_packet("10.0.0.0", "9.9.9.9"), LINKTYPE_ETHERNET));
        assert!(f.matches(&v4_packet("10.255.255.255", "9.9.9.9"), LINKTYPE_ETHERNET));
        assert!(!f.matches(&v4_packet("11.0.0.0", "9.9.9.9"), LINKTYPE_ETHERNET));
        assert!(!f.matches(&v4_packet("9.255.255.255", "9.9.9.9"), LINKTYPE_ETHERNET));

        // /0 keeps every packet of that family, but still not non-IP.
        let all = filter(&["0.0.0.0/0"]);
        assert!(all.matches(&v4_packet("203.0.113.9", "9.9.9.9"), LINKTYPE_ETHERNET));
        assert!(!all.matches(&eth(0, 0x0806, &[0u8; 28]), LINKTYPE_ETHERNET));

        let v6 = filter(&["2606:4700::/32"]);
        assert!(v6.matches(&eth(0, ETHERTYPE_IPV6, &ipv6("2606:4700:dead::1", "::1")), LINKTYPE_ETHERNET));
        assert!(!v6.matches(&eth(0, ETHERTYPE_IPV6, &ipv6("2606:4701::1", "::1")), LINKTYPE_ETHERNET));
    }

    #[test]
    fn families_do_not_cross_match() {
        // A v4-only filter must never keep a v6 packet, and vice versa --
        // the split vectors mean the comparison is never even attempted.
        let v4only = filter(&["1.1.1.1"]);
        assert!(!v4only.matches(&eth(0, ETHERTYPE_IPV6, &ipv6("::1", "::2")), LINKTYPE_ETHERNET));
        let v6only = filter(&["::1"]);
        assert!(!v6only.matches(&v4_packet("1.1.1.1", "9.9.9.9"), LINKTYPE_ETHERNET));
    }

    #[test]
    fn vlan_tags_are_skipped_single_and_double() {
        let f = filter(&["1.1.1.1"]);
        let one = eth(1, ETHERTYPE_IPV4, &ipv4("1.1.1.1", "9.9.9.9"));
        let two = eth(2, ETHERTYPE_IPV4, &ipv4("9.9.9.9", "1.1.1.1"));
        assert!(f.matches(&one, LINKTYPE_ETHERNET));
        assert!(f.matches(&two, LINKTYPE_ETHERNET));
        // Beyond the tag bound we stop rather than hunt through the payload.
        let three = eth(3, ETHERTYPE_IPV4, &ipv4("1.1.1.1", "9.9.9.9"));
        assert!(!f.matches(&three, LINKTYPE_ETHERNET));
    }

    #[test]
    fn unparseable_packets_are_dropped() {
        let f = filter(&["1.1.1.1"]);
        // ARP.
        assert!(!f.matches(&eth(0, 0x0806, &[0u8; 28]), LINKTYPE_ETHERNET));
        // Truncated before the addresses (snaplen too small).
        let full = v4_packet("1.1.1.1", "9.9.9.9");
        assert!(!f.matches(&full[..ETH_HEADER_LEN + 12], LINKTYPE_ETHERNET));
        // Shorter than an Ethernet header, and empty.
        assert!(!f.matches(&full[..8], LINKTYPE_ETHERNET));
        assert!(!f.matches(&[], LINKTYPE_ETHERNET));
        // Right ethertype, wrong version nibble.
        let mut bad_version = full.clone();
        bad_version[ETH_HEADER_LEN] = 0x65;
        assert!(!f.matches(&bad_version, LINKTYPE_ETHERNET));
    }

    #[test]
    fn non_ethernet_linktypes_match_nothing() {
        let f = filter(&["1.1.1.1"]);
        let pkt = v4_packet("1.1.1.1", "9.9.9.9");
        // 113 = LINUX_SLL, what `tcpdump -i any` produces; 101 = RAW; 0 = NULL.
        for linktype in [0u32, 101, 113, 276] {
            assert!(!f.matches(&pkt, linktype), "linktype {linktype} should match nothing");
        }
    }

    #[test]
    fn bad_specs_are_rejected_at_parse() {
        for spec in ["1.1.1.1/33", "::1/129", "not-an-ip", "1.1.1.1/x", "1.1.1.256"] {
            assert!(
                IpFilter::parse(&specs(&[spec])).is_err(),
                "{spec:?} should be rejected"
            );
        }
        // /0 and full-width prefixes are fine.
        assert!(IpFilter::parse(&specs(&["1.1.1.1/32", "0.0.0.0/0", "::/0", "::1/128"])).is_ok());
    }

    #[test]
    fn multiple_entries_are_all_considered() {
        let f = filter(&["1.1.1.1", "8.8.8.8", "10.0.0.0/8", "2606:4700::/32"]);
        assert!(f.matches(&v4_packet("8.8.8.8", "9.9.9.9"), LINKTYPE_ETHERNET));
        assert!(f.matches(&v4_packet("9.9.9.9", "10.1.2.3"), LINKTYPE_ETHERNET));
        assert!(f.matches(&eth(0, ETHERTYPE_IPV6, &ipv6("::1", "2606:4700::1")), LINKTYPE_ETHERNET));
        assert!(!f.matches(&v4_packet("9.9.9.9", "203.0.113.1"), LINKTYPE_ETHERNET));
    }
}