use nix::{
    libc::{c_long, EPERM, ORIG_RAX},
    sys::stat::Mode,
    unistd::{getpid, getppid, Pid},
};
use std::io::Write;
use std::{
    ffi::{c_void, CString},
    mem::size_of,
};

use clap::*;

use libseccomp::*;

#[derive(clap::Parser, Debug)]
struct Args {
    #[cfg(feature = "tracing-mode")]
    #[clap(short, long)]
    log_failed_to: Option<String>,

    #[clap(required = true)]
    command: Vec<String>,
}

fn main() -> anyhow::Result<()> {
    env_logger::init();

    let args = Args::parse();

    log::trace!("args parsed: {args:?}");

    #[cfg(feature = "tracing-mode")]
    let tracing = args.log_failed_to != None;

    let mut default_action = ScmpAction::Errno(EPERM);

    #[cfg(feature = "tracing-mode")]
    if let Some(log_fail_to) = args.log_failed_to {
        use nix::sys::{
            ptrace::Options,
            wait::{waitpid, WaitStatus},
        };
        log::trace!("tracing mode. forking into a tracer...");

        let fork_result = match unsafe { nix::unistd::fork() } {
            Err(err) => panic!("failed to fork a child: {}", err),
            Ok(pid) => pid,
        };

        log::trace!("fork result: {fork_result:?}");

        if let nix::unistd::ForkResult::Parent { child } = fork_result {
            // i'm parent, setting up ptrace

            log::trace!("waiting for child to be ready...");
            waitpid(child, None)?;

            let mut output = match std::fs::OpenOptions::new()
                .append(true)
                .write(true)
                .create(true)
                .open(log_fail_to)
            {
                Ok(file) => file,
                Err(err) => {
                    nix::sys::ptrace::kill(child)?;
                    return Err(err.into());
                }
            };
            nix::sys::ptrace::setoptions(
                child,
                Options::PTRACE_O_TRACESECCOMP
                    | Options::PTRACE_O_TRACECLONE
                    | Options::PTRACE_O_TRACEFORK
                    | Options::PTRACE_O_TRACEVFORK,
            )?;
            nix::sys::ptrace::cont(child, None)?;
            log::trace!("child is ready");

            loop {
                log::debug!("parent: waitpid...");
                let waitstatus = waitpid(Pid::from_raw(-1), None)?;

                match waitstatus {
                    WaitStatus::Exited(pid, ret) => {
                        log::debug!("child {pid} exited with return code {ret}");
                        if pid == child {
                            break;
                        }
                    }
                    WaitStatus::PtraceEvent(pid, sig, _) => {
                        let syscall_nr = nix::sys::ptrace::read_user(
                            pid,
                            (size_of::<c_long>() * ORIG_RAX as usize) as *mut c_void,
                        )? as i32;
                        let syscall = ScmpSyscall::from(syscall_nr);
                        let syscall_name = syscall
                            .get_name()
                            .unwrap_or(format!("syscall({syscall_nr})"));
                        log::debug!("parent: child {pid} received signal {sig:?} syscall: {syscall_name}({syscall_nr})");

                        if let Err(err) = writeln!(output, "{} {}", pid.as_raw(), syscall_name) {
                            log::warn!("failed to write to output file: {err}")
                        }
                        nix::sys::ptrace::cont(pid, None).ok();
                    }
                    _ => {
                        if let Some(pid) = waitstatus.pid() {
                            nix::sys::ptrace::cont(pid, None).ok();
                        }
                    }
                }
            }
            return Ok(());
        }
    }

    let command: Vec<_> = args
        .command
        .iter()
        .map(|s| CString::new(s.as_str()).unwrap())
        .collect();
    log::trace!("command: {command:?}");
    log::trace!("executable: {:?}", args.command[0]);
    let env: Vec<CString> = std::env::vars()
        .into_iter()
        .map(|(k, v)| format!("{k}={v}"))
        .map(|s| CString::new(s).unwrap())
        .collect();

    let exe_fd = nix::fcntl::open(
        format!("{}", args.command[0]).as_str(),
        nix::fcntl::OFlag::O_RDONLY | nix::fcntl::OFlag::O_PATH,
        Mode::empty(),
    )?;

    #[cfg(feature = "tracing-mode")]
    if tracing {
        // i'm child. trace me
        log::debug!("child: traceme");
        nix::sys::ptrace::traceme()?;

        let me = getpid();
        log::info!("waiting for parent to be ready...");
        nix::sys::signal::kill(me, nix::sys::signal::Signal::SIGSTOP)?;
        log::info!("parent is ready");

        let parent = getppid();
        default_action = ScmpAction::Trace(parent.as_raw() as u16);
        log::info!("set default action to {default_action:?}");
    }

    log::trace!("create filter context...");
    let mut filter = ScmpFilterContext::new_filter(default_action)?;

    filter.add_rule_conditional(
        ScmpAction::Allow,
        ScmpSyscall::from(nix::libc::SYS_execveat as i32),
        &[ScmpArgCompare::new(
            0,
            libseccomp::ScmpCompareOp::Equal,
            exe_fd as u64,
        )],
    )?;

    x2t_sandbox_rulegen::generate! {
        log::trace!("accepting {}({})", syscall_name, syscall_nr);
    };

    #[cfg(feature = "tracing-mode")]
    if tracing {
        log::debug!("no need to restrict myself by set_no_new_privs");
    } else {
        log::debug!("restrict myself by set_no_new_privs...");
        nix::sys::prctl::set_no_new_privs()?;
    }

    log::info!("loading filter into kernel...");
    if let Err(err) = filter.load() {
        log::error!("failed to load filter into kernel: {err}");
        return Err(err.into());
    }
    log::trace!("loaded");

    log::debug!("executing {:?}", args.command);
    if let Err(err) = nix::unistd::fexecve(exe_fd, command.as_slice(), env.as_slice()) {
        panic!("failed to execve for {err}");
    }

    panic!("unreachable");
}