use clap;
use std::fmt;
use std::net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket};
use std::ops::Index;

use crc32fast::hash;
use log::warn;
use nom::bytes::complete::{tag, take};
use nom::error::ParseError;
use nom::multi::{many, many0};
use nom::number::complete::{be_u128, be_u16, be_u32, be_u8};
use nom::{AsBytes, IResult, Parser};
use rand::RngCore;

// https://github.com/tailscale/tailscale/blob/main/net/stun/stun.go

const ATTR_NUM_SOFTWARE: u16 = 0x8022;
const ATTR_NUM_FINGERPRINT: u16 = 0x8028;
const ATTR_MAPPED_ADDRESS: u16 = 0x0001;
const ATTR_XOR_MAPPED_ADDRESS: u16 = 0x0020;
// This alternative attribute type is not
// mentioned in the RFC, but the shift into
// the "comprehension-optional" range seems
// like an easy mistake for a server to make.
// And servers appear to send it.
const ATTR_XOR_MAPPED_ADDRESS_ALT: u16 = 0x8020;
const ATTR_SOURCE_ADDRESS: u16 = 0x0004;
const ATTR_CHANGED_ADDRESS: u16 = 0x0005;
const ATTR_USERNAME: u16 = 0x0006;
const ATTR_MESSAGE_INTEGRITY: u16 = 0x0008;
const ATTR_ERROR_CODE: u16 = 0x0009;
const ATTR_UNKNOWN_ATTRIBUTES: u16 = 0x000a;
const ATTR_REALM: u16 = 0x0014;
const ATTR_NONCE: u16 = 0x0015;
const ATTR_ALTERNATE_SERVER: u16 = 0x8023;

const SOFTWARE: [u8; 8] = *b"tailnode";
const BINDING_REQUEST: [u8; 2] = [0x00, 0x01];
const MAGIC_COOKIE: [u8; 4] = [0x21, 0x12, 0xa4, 0x42];
const LEN_FINGERPRINT: u16 = 8;
const HEADER_LEN: u16 = 20;

#[derive(Debug, Clone)]
struct TxId([u8; 12]);

impl TxId {
    fn new() -> Self {
        let mut tx_id = [0; 12];
        rand::rng().fill_bytes(&mut tx_id);
        Self(tx_id)
    }

    fn from_bytes(bytes: &[u8]) -> Self {
        let mut tx_id = [0; 12];
        tx_id.copy_from_slice(bytes);
        Self(tx_id)
    }

    fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    fn make_request(&self) -> Vec<u8> {
        const LEN_ATTR_SOFTWARE: u16 = 4 + SOFTWARE.len() as u16;
        let mut buf =
            Vec::with_capacity((HEADER_LEN + LEN_ATTR_SOFTWARE + LEN_FINGERPRINT) as usize);
        buf.extend(&BINDING_REQUEST);
        buf.extend((LEN_ATTR_SOFTWARE + LEN_FINGERPRINT).to_be_bytes());
        buf.extend(&MAGIC_COOKIE);
        buf.extend(self.as_bytes());
        buf.extend(ATTR_NUM_SOFTWARE.to_be_bytes());
        buf.extend((SOFTWARE.len() as u16).to_be_bytes());
        buf.extend(&SOFTWARE);

        let fp = fingerprint(&buf);
        buf.extend(ATTR_NUM_FINGERPRINT.to_be_bytes());
        buf.extend((4 as u16).to_be_bytes());
        buf.extend(&fp.to_be_bytes());

        buf
    }
}

impl From<&TxId> for u128 {
    fn from(value: &TxId) -> Self {
        let mut padded = [0u8; 16];
        padded[4..].copy_from_slice(&value.0);
        u128::from_be_bytes(padded)
    }
}

fn fingerprint(msg: &[u8]) -> u32 {
    hash(msg) ^ 0x5354554e
}

#[derive(clap::Parser)]
#[command(version, about, long_about = None)]
#[command(about = "Test a Tailscale derp node's stun service")]
struct Cli {
    host: String,
    port: Option<u16>,
    #[clap(long, short, default_value_t = false)]
    debug: bool,
}

#[derive(Debug, Clone, Copy)]
enum StunClass {
    Request = 0,
    Indication = 1,
    SuccessResponse = 2,
    ErrorResponse = 3,
}

#[derive(Debug, Clone, Copy)]
enum StunMethod {
    Binding = 1,
}

#[derive(Debug, Clone)]
enum StunAttribute {
    MappedAddress((IpAddr, u16)),
    XorMappedAddress((IpAddr, u16)),
    SourceAddress((IpAddr, u16)),
    ChangedAddress((IpAddr, u16)),
    Username(String),
    MessageIntegrity([u8; 20]),
    Fingerprint(u32),
    ErrorCode((u16, String)),
    Realm(String),
    Nonce(String),
    UnknownAttributes(Vec<u16>),
    Software(String),
    AlternateServer((IpAddr, u16)),
    Unknown((u16, Vec<u8>)),
}

fn addr_family(addr: &IpAddr) -> &'static str {
    match addr {
        IpAddr::V4(_) => "IPv4",
        IpAddr::V6(_) => "IPv6",
    }
}

impl fmt::Display for StunAttribute {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StunAttribute::MappedAddress((addr, port)) => {
                write!(
                    f,
                    "  MappedAddress ({}) {}:{}",
                    addr_family(addr),
                    addr,
                    port
                )
            }
            StunAttribute::SourceAddress((addr, port)) => {
                write!(
                    f,
                    "  SourceAddress ({}) {}:{}",
                    addr_family(addr),
                    addr,
                    port
                )
            }
            StunAttribute::ChangedAddress((addr, port)) => {
                write!(
                    f,
                    "  ChangedAddress ({}) {}:{}",
                    addr_family(addr),
                    addr,
                    port
                )
            }
            StunAttribute::XorMappedAddress((addr, port)) => {
                write!(
                    f,
                    "  XorMappedAddress ({}) {}:{}",
                    addr_family(addr),
                    addr,
                    port
                )
            }
            StunAttribute::Username(username) => writeln!(f, "  Username {}", username),
            StunAttribute::MessageIntegrity(msg_integrity) => {
                write!(f, "  MessageIntegrity {:?}", msg_integrity)
            }
            StunAttribute::Fingerprint(fingerprint) => {
                write!(f, "  Fingerprint 0x{:08x}", fingerprint)
            }
            StunAttribute::ErrorCode((err_num, error)) => {
                write!(f, "  ErrorCode {} ({})", err_num, error)
            }
            StunAttribute::Realm(realm) => writeln!(f, "  Realm {}", realm),
            StunAttribute::Nonce(nonce) => writeln!(f, "  Nonce {}", nonce),
            StunAttribute::UnknownAttributes(unknown_attrs) => {
                write!(f, "  UnknownAttributes {:?}", unknown_attrs)
            }
            StunAttribute::Software(software) => writeln!(f, "  Software {}", software),
            StunAttribute::AlternateServer((addr, port)) => {
                write!(
                    f,
                    "  AlternateServer ({}) {}:{}",
                    addr_family(addr),
                    addr,
                    port
                )
            }
            StunAttribute::Unknown((attr_type, data)) => {
                write!(f, "  Unknown ({}) {:?}", attr_type, data)
            }
        }
    }
}

#[derive(Debug)]
struct StunMessageType {
    class: StunClass,
    method: StunMethod,
}

#[derive(Debug)]
struct StunHeader {
    msg_type: StunMessageType,
    msg_length: u16,
    tx_id: TxId,
}

impl fmt::Display for StunHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "  Header")?;
        writeln!(
            f,
            "    MessageType class={:?} (0x{:02x}) method={:?} (0x{:04x})",
            self.msg_type.class,
            self.msg_type.class as usize,
            self.msg_type.method,
            self.msg_type.method as usize
        )
    }
}

#[derive(Debug)]
struct StunAttributes(Vec<StunAttribute>);

#[derive(Debug)]
struct StunMessage {
    h: StunHeader,
    attributes: StunAttributes,
}

impl fmt::Display for StunMessage {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        writeln!(f, "StunMessage")?;
        write!(f, "{}", self.h)?;
        write!(f, "{}", self.attributes)
    }
}

impl fmt::Display for StunAttributes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "  Attributes")?;
        for attr in &self.0 {
            writeln!(f, "    {}", attr)?;
        }
        Ok(())
    }
}

fn take_txid<'a, E: ParseError<&'a [u8]>>(bytes: &'a [u8]) -> IResult<&[u8], TxId, E> {
    let (bytes, tx_id) = take(12usize)(bytes)?;
    Ok((bytes, TxId::from_bytes(tx_id)))
}

fn parse_stun_message<'a, E: ParseError<&'a [u8]>>(
    bytes: &'a [u8],
) -> IResult<&'a [u8], StunMessage, E> {
    let (bytes, h) = parse_stun_header(bytes)?;

    let (bytes, attributes) = many0(parse_stun_attribute::<&[u8], E>(&h.tx_id)).parse(bytes)?;
    let attributes = StunAttributes(attributes.iter().filter_map(|i| i.clone()).collect());

    Ok((bytes, StunMessage { h, attributes }))
}

fn parse_stun_message_type<'a, E: ParseError<&'a [u8]>>(
    bytes: &'a [u8],
) -> IResult<&'a [u8], StunMessageType, E> {
    let (bytes, msg_type_raw) = be_u16(bytes)?;
    if msg_type_raw & 0b11000000 != 0 {
        panic!("Invalid STUN message type");
    }
    let class_raw = ((msg_type_raw & (1 << 4)) >> 4) | ((msg_type_raw & (1 << 8)) >> 7);
    let class = match class_raw {
        0 => StunClass::Request,
        1 => StunClass::Indication,
        2 => StunClass::SuccessResponse,
        3 => StunClass::ErrorResponse,
        _ => panic!("Invalid STUN message class"),
    };
    let method_raw = msg_type_raw & 0x0f | msg_type_raw & 0xe0 >> 1 | msg_type_raw & 0x3e >> 2;
    let method = match method_raw {
        1 => StunMethod::Binding,
        _ => panic!("Invalid STUN message method"),
    };
    Ok((bytes, StunMessageType { class, method }))
}

fn parse_stun_address<I, E: ParseError<I>>(bytes: I) -> IResult<I, (IpAddr, u16), E>
where
    I: nom::Input<Item = u8>,
{
    let (bytes, _) = take(1usize)(bytes)?;
    let (bytes, family) = be_u8(bytes)?;
    let (bytes, port) = be_u16(bytes)?;
    let (bytes, addr) = match family {
        0x01 => {
            let (bytes, val) = be_u32(bytes)?;
            (bytes, (IpAddr::V4(val.into()), port))
        }
        0x02 => {
            let (bytes, val) = be_u128(bytes)?;
            (bytes, (IpAddr::V6(val.into()), port))
        }
        _ => panic!("Invalid address family"),
    };
    Ok((bytes, addr))
}

fn parse_stun_xor_address<I, E: ParseError<I>>(
    bytes: I,
    tx_id: &TxId,
) -> IResult<I, (IpAddr, u16), E>
where
    I: nom::Input<Item = u8>,
{
    let (bytes, addr) = parse_stun_address(bytes)?;
    let xor_port = addr.1 ^ 0x2112;
    let xor_addr = match addr.0 {
        IpAddr::V4(v4) => {
            let v4 = u32::from(v4);
            let xor_v4 = v4 ^ 0x2112a442;
            IpAddr::V4(xor_v4.into())
        }
        IpAddr::V6(v6) => {
            let v6 = u128::from(v6);
            let xor_v6: u128 = v6 ^ (0x2112a442 << 96 | u128::from(tx_id));
            IpAddr::V6(xor_v6.into())
        }
    };
    Ok((bytes, (xor_addr, xor_port)))
}

fn parse_stun_attribute<I, E: ParseError<I>>(
    tx_id: &TxId,
) -> impl Fn(I) -> IResult<I, Option<StunAttribute>, E>
where
    I: nom::Input<Item = u8> + nom::Compare<I> + AsBytes,
{
    let tx_id = tx_id.clone();
    move |bytes| parse_stun_attribute_impl(bytes, &tx_id)
}

fn parse_stun_attribute_impl<I, E: ParseError<I>>(
    bytes: I,
    tx_id: &TxId,
) -> IResult<I, Option<StunAttribute>, E>
where
    I: nom::Input<Item = u8> + nom::Compare<I> + AsBytes,
{
    let (bytes, attr_type) = be_u16(bytes)?;
    let (bytes, attr_len) = be_u16(bytes)?;
    let (bytes, attr_data) = take(attr_len)(bytes)?;

    if attr_len == 0 {
        return Ok((bytes, None));
    }

    let attr = match attr_type {
        ATTR_MAPPED_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::MappedAddress(addr)
        }
        ATTR_SOURCE_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::SourceAddress(addr)
        }
        ATTR_CHANGED_ADDRESS => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::ChangedAddress(addr)
        }
        ATTR_XOR_MAPPED_ADDRESS | ATTR_XOR_MAPPED_ADDRESS_ALT => {
            let (_residual, addr) = parse_stun_xor_address(attr_data, tx_id)?;

            StunAttribute::XorMappedAddress(addr)
        }
        ATTR_USERNAME => {
            let username = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Username(username)
        }
        ATTR_MESSAGE_INTEGRITY => {
            let mut msg_integrity = [0u8; 20];
            msg_integrity.copy_from_slice(attr_data.as_bytes());
            StunAttribute::MessageIntegrity(msg_integrity)
        }
        ATTR_ERROR_CODE => {
            let (attr_data, zeros) = take(2usize)(attr_data)?;
            if zeros.iter_elements().any(|b| b != 0) {
                panic!("Invalid STUN error code");
            }
            let (attr_data, err_num) = be_u16(attr_data)?;
            if err_num & 0b1111100000000000 != 0 {
                panic!("Invalid STUN error code");
            }
            let error = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::ErrorCode((err_num, error))
        }
        ATTR_UNKNOWN_ATTRIBUTES => {
            let (_residual, unknown_attrs) =
                many((attr_len / 2) as usize, be_u16).parse(attr_data)?;
            StunAttribute::UnknownAttributes(unknown_attrs)
        }
        ATTR_NUM_FINGERPRINT => {
            let (_residual, fingerprint) = be_u32(attr_data)?;
            StunAttribute::Fingerprint(fingerprint)
        }
        ATTR_REALM => {
            let realm = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Realm(realm)
        }
        ATTR_NONCE => {
            let nonce = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Nonce(nonce)
        }
        ATTR_NUM_SOFTWARE => {
            let software = String::from_iter(attr_data.iter_elements().map(|b| b as char));
            StunAttribute::Software(software)
        }
        ATTR_ALTERNATE_SERVER => {
            let (_residual, addr) = parse_stun_address(attr_data)?;
            StunAttribute::AlternateServer(addr)
        }

        t => {
            warn!("Unknown STUN attribute type {t}");
            StunAttribute::Unknown((t, attr_data.iter_elements().collect()))
        }
    };

    Ok((bytes, Some(attr)))
}

fn parse_stun_header<'a, E: ParseError<&'a [u8]>>(
    bytes: &'a [u8],
) -> IResult<&'a [u8], StunHeader, E> {
    let (bytes, msg_type) = parse_stun_message_type(bytes)?;
    let (bytes, msg_length) = be_u16(bytes)?;
    let (bytes, _) = tag(MAGIC_COOKIE.as_slice())(bytes)?;
    let (bytes, tx_id) = take_txid(bytes)?;
    Ok((
        bytes,
        StunHeader {
            msg_type,
            msg_length,
            tx_id,
        },
    ))
}

fn main() {
    let cli = <Cli as clap::Parser>::parse();

    let dest = cli.host + ":" + &cli.port.unwrap_or(3478).to_string();
    let socket = UdpSocket::bind("[::]:0").expect("Unable to bind a UDP socket");
    socket
        .connect(dest)
        .expect("Unable to connect to the destination");
    if cli.debug {
        println!(
            "Connected to {:?} from {:?}",
            socket.peer_addr(),
            socket.local_addr()
        );
    }

    let req = TxId::new().make_request();
    if cli.debug {
        println!("Sending request {:?}", &req);
    }
    socket.send(&req).expect("Unable to send request");

    let mut buf = [0u8; 1500];

    if let Ok(received) = socket.recv(&mut buf) {
        if cli.debug {
            println!("Received response: {:?}", &buf[..received]);
        }
        let (_residual, msg) = parse_stun_message::<nom::error::Error<&[u8]>>(&buf).unwrap();
        println!("Parsed message from {}:", socket.peer_addr().unwrap());
        println!("{}", msg);
    } else if let Err(e) = socket.recv(&mut buf) {
        println!("recv function failed: {e:?}");
    }
}