use pcap::Capture;
use stable_eyre::eyre::{eyre, Result};
use std::collections::HashSet;
use std::convert::TryInto;

mod ieee80211 {
    use int_enum::IntEnum;
    use std::convert::TryFrom;
    use std::convert::TryInto;
    use std::fmt;

    // data/extension frames are irrelevant so just ignore their subtype
    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
    pub enum FrameType {
        Management(ManagementSubtype),
        Control(ControlSubtype),
        Data,
        Extension
    }

    #[repr(u8)]
    #[derive(Clone, Copy, Debug, Eq, PartialEq, IntEnum)]
    pub enum ManagementSubtype {
        AssociationRequest = 0b0000,
        AssociationResponse = 0b0001,
        ReassociationRequest = 0b0010,
        ReassociationResponse = 0b0011,
        ProbeRequest = 0b0100,
        ProbeResponse = 0b0101,
        Beacon = 0b1000,
        Disassociation = 0b1010,
        Authentication = 0b1011,
        Deauthentication = 0b1100,
        Other = 0b1111
    }

    #[repr(u8)]
    #[derive(Clone, Copy, Debug, Eq, PartialEq, IntEnum)]
    pub enum ControlSubtype {
        Rts = 0b1011,
        Cts = 0b1100,
        Ack = 0b1101,
        Other = 0b1111
    }

    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
    pub struct FrameControl {
        pub version: u8,
        pub ftype: FrameType,
        pub to_ds: bool,
        pub from_ds: bool,
        pub more_fragments: bool,
        pub retry: bool,
        pub more_data: bool,
        pub protected: bool,
        pub htc_order: bool,
        subtype: u8
    }

    // read a bit out of a number (0 = LSB)
    fn read_bit(x: u8, bit: u8) -> bool {
        x >> bit & 0b1 == 1
    }

    pub fn parse_frame_control(version_type_subtype: u8, flags: u8) -> Option<FrameControl> {
        let subtype = version_type_subtype >> 4 & 0b1111;
        let ftype = match version_type_subtype >> 2 & 0b11 {
            0b00 => FrameType::Management(ManagementSubtype::try_from(subtype).unwrap_or(ManagementSubtype::Other)),
            0b01 => FrameType::Control(ControlSubtype::try_from(subtype).unwrap_or(ControlSubtype::Other)),
            0b10 => FrameType::Data,
            0b11 => FrameType::Extension,
            _ => unreachable!()
        };
        //println!("{:?} {:08b}", ftype, flags);
        Some(FrameControl {
            version: version_type_subtype & 0b11,
            ftype: ftype,
            to_ds: read_bit(flags, 0),
            from_ds: read_bit(flags, 1),
            more_fragments: read_bit(flags, 2),
            retry: read_bit(flags, 3),
            more_data: read_bit(flags, 5),
            protected: read_bit(flags, 6),
            htc_order: read_bit(flags, 7),
            subtype: subtype
        })
    }

    #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
    pub struct MacAddress([u8; 6]);

    impl fmt::Display for MacAddress {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x}", self.0[0], self.0[1], self.0[2], self.0[3], self.0[4], self.0[5])
        }
    }

    impl MacAddress {
        pub fn is_local(&self) -> bool { read_bit(self.0[0], 1) }
        pub fn is_group(&self) -> bool { read_bit(self.0[0], 0) }
    }

    pub fn read_addresses(frame: &[u8]) -> Vec<MacAddress> {
        let mut out = Vec::with_capacity(4);
        if frame.len() >= 10 { out.push(MacAddress(frame[4..10].try_into().unwrap())); }
        if frame.len() >= 16 { out.push(MacAddress(frame[10..16].try_into().unwrap())); }
        if frame.len() >= 22 { out.push(MacAddress(frame[16..22].try_into().unwrap())); }
        //if frame.len() >= 30 { out.push(MacAddress(frame[24..30].try_into().unwrap())); }
        out
    }
}

fn main() -> Result<()> {
    let cap = Capture::from_device("mon0")?.immediate_mode(true);
    let mut cap = cap.open()?;
    if cap.get_datalink() != pcap::Linktype::IEEE802_11_RADIOTAP { return Err(eyre!("device link type wrong")) }

    let mut seen = HashSet::new();

    while let Ok(packet) = cap.next() {
        //println!("{:?}", packet);
        match radiotap::parse(packet.data) {
            Ok(header) => {
                let data = &packet.data[header.length()..];
                match ieee80211::parse_frame_control(data[0], data[1]) {
                    Some(ieee80211::FrameControl { ftype: ieee80211::FrameType::Control(ieee80211::ControlSubtype::Other), .. }) => {},
                    Some(x) => {
                        let sequence = if data.len() >= 24 {
                            let sn_fst = data[22] as u16;
                            let sn_snd = data[23] as u16;
                            let seq = (sn_snd << 4) + (sn_fst >> 4);
                            Some((seq, sn_fst & 0b1111))
                        } else { None };
                        let addrs = ieee80211::read_addresses(data);
                        for addr in &addrs {
                            if !seen.contains(addr) && !addr.is_group() && !addr.is_local() {
                                println!("attained new MAC {} from {:?} {:?} {:?} local={} group={}", addr, x, addrs, sequence, addr.is_local(), addr.is_group());
                                seen.insert(addr.clone());
                            }
                        }
                        //println!("{:?} {:?}\nADDR={:?} SEQ={:?}\n", header.antenna_signal.unwrap().into_inner(), x, addrs, sequence);
                    },
                    _ => {}
                }
            },
            // This happens sometimes. I don't know why, since it really shouldn't (the radiotap headers are applied by the kernel). Thusly, just ignore it and hope it goes away.
            Err(e) => {
                eprintln!("header decoding failed: {:?}", e);
            }
        }
    }
    Ok(())
}