//! TLS Certificate Validator / UDP Magic Detector - eBPF-based

use std::collections::HashSet;
use std::mem::size_of;
use std::net::Ipv4Addr;

use anyhow::{Context, Result};
use aya::maps::{HashMap as AyaHashMap, RingBuf};
use aya::programs::{Xdp, XdpFlags};
use aya::{include_bytes_aligned, Bpf};
use log::{info, warn};
use tls_parser::{parse_tls_plaintext, TlsMessage, TlsMessageHandshake};
use tokio::signal;

use packet_detector::validator::CertValidator;

// this has to be the exact same as the struct in kernelspace
#[repr(C)]
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct ConnKey {
    port_lo: u16,
    port_hi: u16,
}

unsafe impl aya::Pod for ConnKey {}

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 }
    }
}

// this has to be the exact same as the struct in kernelspace
#[repr(C)]
#[derive(Clone, Copy)]
struct Event {
    src_ip: u32,
    dst_ip: u32,
    src_port: u16,
    dst_port: u16,
    tls_len: u16,
    _pad: u16,
}

unsafe impl aya::Pod for Event {}

const EVENT_SIZE: usize = size_of::<Event>();

fn ip(n: u32) -> Ipv4Addr { 
    Ipv4Addr::from(n.to_be_bytes()) 
}

fn extract_certs(tls_data: &[u8]) -> Option<Vec<Vec<u8>>> {
    if tls_data.len() < 6 || tls_data[0] != 0x16 || tls_data[5] != 0x0B { return None; }
    let (_, rec) = parse_tls_plaintext(tls_data).ok()?;
    for msg in &rec.msg {
        if let TlsMessage::Handshake(TlsMessageHandshake::Certificate(c)) = msg {
            return Some(c.cert_chain.iter().map(|x| x.data.to_vec()).collect());
        }
    }
    None
}

enum Decision {
    Allow(ConnKey),
    Block(ConnKey),
    Skip,
}

fn handle_event(data: &[u8], validator: Option<&CertValidator>) -> Decision {
    if data.len() < EVENT_SIZE { return Decision::Skip; }
    let ev: Event = unsafe { std::ptr::read(data.as_ptr() as *const _) };
    let addr = format!("{}:{} -> {}:{}", ip(ev.src_ip), ev.src_port, ip(ev.dst_ip), ev.dst_port);
    let conn_key = make_conn_key(ev.src_port, ev.dst_port);

    if ev.tls_len == 0 {
        info!("UDP magic from {}", addr);
        return Decision::Allow(conn_key);
    }

    let Some(v) = validator else { return Decision::Skip };
    let end = EVENT_SIZE + ev.tls_len as usize;
    if end > data.len() { return Decision::Skip; }
    let Some(certs) = extract_certs(&data[EVENT_SIZE..end]) else { return Decision::Skip };
    let result = v.validate(&certs);
    info!("{}: {}", addr, result.subject);

    if result.valid {
        info!("ALLOW conn {}:{} (signed by {})", conn_key.port_lo, conn_key.port_hi, result.issuer);
        Decision::Allow(conn_key)
    } else {
        warn!("BLOCK conn {}:{} - {}", conn_key.port_lo, conn_key.port_hi, result.error.unwrap_or_default());
        Decision::Block(conn_key)
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    rustls::crypto::ring::default_provider().install_default().ok();
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: {} <interface> [ca-cert.pem]", args[0]);
        std::process::exit(1);
    }

    let iface = &args[1];
    let validator = args.get(2).map(|p| CertValidator::with_ca_file(p)).transpose()?;
    info!("Mode: {}", if validator.is_some() { "TLS cert validation" } else { "UDP magic detection" });

    let mut bpf = Bpf::load(include_bytes_aligned!("../../target/bpfel-unknown-none/release/packet-detector"))?;
    let program: &mut Xdp = bpf.program_mut("packet_detector").unwrap().try_into()?;
    program.load()?;
    program.attach(iface, XdpFlags::default()).context("XDP attach failed")?;
    info!("XDP attached to {}", iface);

    let mut allowed: AyaHashMap<_, ConnKey, u8> = AyaHashMap::try_from(bpf.take_map("ALLOWED_CONNS").unwrap())?;
    let mut blocked: AyaHashMap<_, ConnKey, u8> = AyaHashMap::try_from(bpf.take_map("BLOCKED_CONNS").unwrap())?;
    let mut ring: RingBuf<_> = RingBuf::try_from(bpf.take_map("TLS_EVENTS").unwrap())?;
    let mut allowed_count = 0u32;
    let mut blocked_count = 0u32;

    println!("\nRunning on {} - Ctrl+C to stop\n", iface);

    loop {
        tokio::select! {
            _ = signal::ctrl_c() => break,
            _ = tokio::time::sleep(tokio::time::Duration::from_millis(10)) => {
                while let Some(item) = ring.next() {
                    match handle_event(item.as_ref(), validator.as_ref()) {
                        Decision::Allow(key) => {
                            allowed.insert(key, 1, 0)?;
                            allowed_count += 1;
                        }
                        Decision::Block(key) => {
                            blocked.insert(key, 1, 0)?;
                            blocked_count += 1;
                        }
                        Decision::Skip => {}
                    }
                }
            }
        }
    }
    println!("\nAllowed: {}, Blocked: {}", allowed_count, blocked_count);
    Ok(())
}