#![no_std]
#![no_main]

use aya_ebpf::{bindings::xdp_action, macros::{map, xdp}, maps::{HashMap, RingBuf}, programs::XdpContext};
use network_types::{eth::{EthHdr, EtherType}, ip::{Ipv4Hdr, IpProto}, tcp::TcpHdr, udp::UdpHdr};

#[repr(C)]
#[derive(Clone, Copy)]
pub struct ConnKey {
    pub port_lo: u16,  // lower port (client ephemeral)
    pub port_hi: u16,  // higher port (server, e.g. 8443)
}

#[repr(C)]
pub struct TlsEvent {
    pub src_ip: u32,
    pub dst_ip: u32,
    pub src_port: u16,
    pub dst_port: u16,
    pub tls_len: u16,
    pub _pad: u16,
}

#[map]
static TLS_EVENTS: RingBuf = RingBuf::with_byte_size(256 * 1024, 0);

#[map]
static ALLOWED_CONNS: HashMap<ConnKey, u8> = HashMap::with_max_entries(4096, 0);

#[map]
static BLOCKED_CONNS: HashMap<ConnKey, u8> = HashMap::with_max_entries(4096, 0);

// Magic word: "hell0123" (8 bytes)
const MAGIC: [u8; 8] = *b"hell0123";

#[xdp]
pub fn packet_detector(ctx: XdpContext) -> u32 {
    match try_process(&ctx) {
        Ok(action) => action,
        Err(_) => xdp_action::XDP_PASS
    }
}

fn try_process(ctx: &XdpContext) -> Result<u32, ()> {
    let eth = ptr::<EthHdr>(ctx, 0)?;
    if unsafe { (*eth).ether_type } != EtherType::Ipv4.into() { return Ok(xdp_action::XDP_PASS); }

    let ip = ptr::<Ipv4Hdr>(ctx, EthHdr::LEN)?;
    let src_ip = u32::from_be_bytes(unsafe { (*ip).src_addr });
    let ip_hdr_len = (unsafe { (*ip).vihl } & 0x0F) as usize * 4;

    match unsafe { (*ip).proto } {
        IpProto::Tcp => process_tcp(ctx, ip, ip_hdr_len, src_ip),
        IpProto::Udp => process_udp(ctx, ip, ip_hdr_len, src_ip),
        _ => Ok(xdp_action::XDP_PASS),
    }
}

fn make_conn_key(src_port: u16, dst_port: u16) -> ConnKey {
    if src_port < dst_port {
        ConnKey { port_lo: src_port, port_hi: dst_port }
    } else {
        ConnKey { port_lo: dst_port, port_hi: src_port }
    }
}

fn process_tcp(ctx: &XdpContext, ip: *const Ipv4Hdr, ip_hdr_len: usize, src_ip: u32) -> Result<u32, ()> {
    let tcp = ptr::<TcpHdr>(ctx, EthHdr::LEN + ip_hdr_len)?;
    let payload_off = EthHdr::LEN + ip_hdr_len + (unsafe { (*tcp).doff() } as usize) * 4;

    let src_port = u16::from_be_bytes(unsafe { (*tcp).source });
    let dst_port = u16::from_be_bytes(unsafe { (*tcp).dest });

    // Only filter TLS ports (443/8443)
    let is_tls_port = dst_port == 443 || src_port == 443 || dst_port == 8443 || src_port == 8443;
    if !is_tls_port {
        return Ok(xdp_action::XDP_PASS);
    }

    let conn_key = make_conn_key(src_port, dst_port);

    // Explicitly blocked (invalid cert) - drop
    if unsafe { BLOCKED_CONNS.get(&conn_key) }.is_some() {
        return Ok(xdp_action::XDP_DROP);
    }

    // Already validated (valid cert) - allow
    if unsafe { ALLOWED_CONNS.get(&conn_key) }.is_some() {
        return Ok(xdp_action::XDP_PASS);
    }

    // Check if this packet has a TLS cert - send to userspace for validation
    if is_tls_cert(ctx, payload_off) {
        const SZ: usize = 512;
        if let Some(tls_ptr) = bounds(ctx, payload_off, SZ) {
            if let Some(mut entry) = TLS_EVENTS.reserve::<[u8; 16 + SZ]>(0) {
                let p = entry.as_mut_ptr() as *mut u8;
                unsafe {
                    core::ptr::write(p as *mut TlsEvent, TlsEvent {
                        src_ip,
                        dst_ip: u32::from_be_bytes((*ip).dst_addr),
                        src_port,
                        dst_port,
                        tls_len: SZ as u16,
                        _pad: 0,
                    });
                    for i in 0..SZ { *p.add(16 + i) = *tls_ptr.add(i); }
                }
                entry.submit(0);
            }
        }
    }

    // Allow packet through - handshake continues, userspace will decide later
    Ok(xdp_action::XDP_PASS)
}

fn process_udp(ctx: &XdpContext, ip: *const Ipv4Hdr, ip_hdr_len: usize, src_ip: u32) -> Result<u32, ()> {
    let udp = ptr::<UdpHdr>(ctx, EthHdr::LEN + ip_hdr_len)?;
    let payload_off = EthHdr::LEN + ip_hdr_len + UdpHdr::LEN;

    // Check for magic word at payload start
    if let Some(data) = bounds(ctx, payload_off, MAGIC.len()) {
        let mut matched = true;
        for i in 0..MAGIC.len() {
            if unsafe { *data.add(i) } != MAGIC[i] { matched = false; break; }
        }
        if matched {
            // Send event to userspace (reuse TlsEvent struct)
            if let Some(mut entry) = TLS_EVENTS.reserve::<[u8; 16]>(0) {
                let p = entry.as_mut_ptr() as *mut u8;
                unsafe {
                    core::ptr::write(p as *mut TlsEvent, TlsEvent {
                        src_ip,
                        dst_ip: u32::from_be_bytes((*ip).dst_addr),
                        src_port: u16::from_be_bytes((*udp).src),
                        dst_port: u16::from_be_bytes((*udp).dst),
                        tls_len: 0, // 0 = UDP magic match
                        _pad: 0,
                    });
                }
                entry.submit(0);
            }
            return Ok(xdp_action::XDP_PASS);
        }
        // No match - drop the packet
        return Ok(xdp_action::XDP_DROP);
    }
    Ok(xdp_action::XDP_PASS)
}

#[inline(always)]
fn is_tls_cert(ctx: &XdpContext, off: usize) -> bool {
    if let Some(p) = bounds(ctx, off, 6) {
        unsafe { *p == 0x16 && *p.add(1) == 0x03 && *p.add(2) <= 0x03 && *p.add(5) == 0x0B }
    } else { false }
}

#[inline(always)]
fn bounds(ctx: &XdpContext, off: usize, len: usize) -> Option<*const u8> {
    if ctx.data() + off + len > ctx.data_end() { None } else { Some((ctx.data() + off) as *const u8) }
}

#[inline(always)]
fn ptr<T>(ctx: &XdpContext, off: usize) -> Result<*const T, ()> {
    bounds(ctx, off, core::mem::size_of::<T>()).map(|p| p as *const T).ok_or(())
}

#[panic_handler]
fn panic(_: &core::panic::PanicInfo) -> ! { unsafe { core::hint::unreachable_unchecked() } }