// Copyright 2025-2026 Dillution <hskimse1@gmail.com>.
//
// This file is part of DPIBreak.
//
// DPIBreak is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation, either version 3 of the License, or (at your
// option) any later version.
//
// DPIBreak is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
// for more details.
//
// You should have received a copy of the GNU General Public License
// along with DPIBreak. If not, see <https://www.gnu.org/licenses/>.

use anyhow::Result;
use etherparse::{IpSlice, TcpSlice};
use anyhow::anyhow;
#[cfg(target_os = "linux")]
use std::sync::atomic::Ordering;

#[cfg(debug_assertions)]
use crate::log_println;
#[cfg(debug_assertions)]
use crate::log;
use crate::opt;
use crate::platform;
use crate::tls;

mod fake;
mod hoptab;

#[cfg(debug_assertions)]
use log::LogLevel;

struct PktView<'a> {
    ip: IpSlice<'a>,
    tcp: TcpSlice<'a>
}

impl<'a> PktView<'a> {
    #[inline]
    fn from_raw(raw: &'a [u8]) -> Result<Self> {
        let ip = IpSlice::from_slice(raw)?;
        let tcp = TcpSlice::from_slice(ip.payload().payload)?;

        Ok(Self { ip, tcp })
    }

    fn ttl(&self) -> u8 {
        use etherparse::IpSlice;

        match &self.ip {
            IpSlice::Ipv4(v4) => v4.header().ttl(),
            IpSlice::Ipv6(v6) => v6.header().hop_limit()
        }
    }

    fn saddr(&self) -> std::net::IpAddr {
        self.ip.source_addr()
    }

    fn daddr(&self) -> std::net::IpAddr {
        self.ip.destination_addr()
    }
}

/// Write TCP/IP packet (payload = view.tcp.payload[start..Some(end)])
/// to out_buf, explicitly clearing before.
///
/// If payload, ttl or tcp_checksum is given, override view's one.
fn split_packet_0(
    view: &PktView,
    start: u32,
    end: Option<u32>,
    out_buf: &mut Vec<u8>,
    payload: Option<&[u8]>,
    ttl: Option<u8>,
    tcp_checksum: Option<u16>
) -> Result<()> {
    use etherparse::*;

    let ip = &view.ip;
    let tcp = &view.tcp;
    let payload = payload.unwrap_or(tcp.payload());

    let end = end.unwrap_or(payload.len().try_into()?);

    if start > end || payload.len() < end as usize {
        return Err(anyhow!("invalid index"));
    }

    let opts = tcp.options();
    let mut tcp_hdr = tcp.to_header();
    tcp_hdr.sequence_number += start;

    let (builder, l3_len) = match ip {
        IpSlice::Ipv4(hdr) => {
            let mut ip_hdr = hdr.header().to_header();
            if let Some(t) = ttl { ip_hdr.time_to_live = t; };

            let exts = hdr.extensions().to_header();
            let l3_len = ip_hdr.header_len() + exts.header_len();

            (PacketBuilder::ip(IpHeaders::Ipv4(
                ip_hdr,
                hdr.extensions().to_header()
            )), l3_len)
        },

        IpSlice::Ipv6(hdr) => {
            let mut ip6_hdr = hdr.header().to_header();
            if let Some(t) = ttl { ip6_hdr.hop_limit = t; };

            let l3_len = Ipv6Header::LEN;

            (PacketBuilder::ip(IpHeaders::Ipv6(
                ip6_hdr,
                Default::default()
            )), l3_len)
        }
    };

    let builder = builder.tcp_header(tcp_hdr).options_raw(opts)?;

    let payload = &payload[start as usize..end as usize];

    out_buf.clear();
    builder.write(out_buf, payload)?;

    if let Some(cs) = tcp_checksum {
        let tcp_csum_off = l3_len + 16;

        if out_buf.len() < tcp_csum_off + 2 {
            return Err(anyhow!("packet too short for tcp checksum patch"));
        }
        out_buf[tcp_csum_off..tcp_csum_off + 2].copy_from_slice(&cs.to_be_bytes());
    }

    Ok(())
}

fn split_packet(
    view: &PktView,
    start: u32,
    end: Option<u32>,
    out_buf: &mut Vec<u8>
) -> Result<()> {
    split_packet_0(view, start, end, out_buf, None, None, None)
}

fn send_segment(
    view: &PktView,
    start: u32,
    end: Option<u32>,
    buf: &mut Vec<u8>
) -> Result<()> {
    use platform::send_to_raw;

    if opt::fake() {
        fake::fake_clienthello(view, start, end, buf)?;
        send_to_raw(buf)?;
    }
    split_packet(view, start, end, buf)?;
    send_to_raw(buf)?;

    Ok(())
}

fn split_packet_1(view: &PktView, order: &[u32], buf: &mut Vec<u8>) -> Result<()> {
    let mut it = order.iter().copied();

    let Some(mut first) = it.next() else {
        return Err(anyhow!("split_packet_1: invalid order array"));
    };

    for next in it {
        send_segment(view, first, Some(next), buf)?;
        std::thread::sleep(std::time::Duration::from_millis(opt::delay_ms()));
        first = next;
    }

    send_segment(view, first, None, buf)?;

    Ok(())
}

fn is_synack_from_443(view: &PktView) -> bool {
    // sport == 443 and flags SYN+ACK
    view.tcp.source_port() == 443 && view.tcp.syn() && view.tcp.ack()
}

/// Return Ok(true) if packet is handled
pub fn handle_packet(pkt: &[u8], buf: &mut Vec::<u8>) -> Result<bool> {
    #[cfg(target_os = "linux")]
    let is_filtered = platform::IS_U32_SUPPORTED.load(Ordering::Relaxed);

    #[cfg(windows)]
    let is_filtered = true;

    let view = PktView::from_raw(pkt)?;

    if opt::fake_autottl() && is_synack_from_443(&view) {
        fake::saddr_hop_put(&view);
        return Ok(false);
    }

    if !is_filtered && !tls::is_client_hello(view.tcp.payload()) {
        return Ok(false);
    }

    // TODO: if clienthello packet has been (unlikely) fragmented,
    // we should find the second part and drop, reassemble it here.

    split_packet_1(&view, &[0, 1], buf)?;

    #[cfg(debug_assertions)]
    log_println!(LogLevel::Debug, "packet is handled, len={}", pkt.len());

    Ok(true)
}

#[macro_export]
macro_rules! handle_packet {
    ($bytes:expr, $buf:expr, handled => $on_handled:expr, rejected => $on_rejected:expr $(,)?) => {{
        match crate::pkt::handle_packet($bytes, $buf) {
            Ok(true) => { $on_handled }
            Ok(false) => { $on_rejected }
            Err(e) => {
                log_println!(LogLevel::Warning, "handle_packet: {e}");
                $on_rejected
            }
        }
    }};
}