//
// Syd: rock-solid unikernel
// src/syd-sys.rs: Given a number, print the matching syscall name and exit.
//                 Given a regex, print case-insensitively matching syscall names and exit.
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{process::ExitCode, str::FromStr, time::Duration};

use getargs::{Opt, Options};
use libseccomp::{ScmpArch, ScmpSyscall};
use nix::{
    errno::Errno,
    sys::{
        signal::{kill, Signal},
        stat::lstat,
        wait::{waitpid, WaitPidFlag, WaitStatus},
    },
    unistd::{fork, ForkResult},
};
use syd::{
    err::SydResult,
    wildmatch::{is_literal, wildmatch},
};

// Set global allocator to mimalloc.
#[global_allocator]
static GLOBAL: mimalloc2_rust::GlobalMiMalloc = mimalloc2_rust::GlobalMiMalloc;

fn main() -> SydResult<ExitCode> {
    syd::set_sigpipe_dfl()?;

    // Configure mimalloc.
    syd::configure_mimalloc();

    let args = std::env::args().skip(1).collect::<Vec<_>>();

    // If no arguments are provided, display help.
    if args.is_empty() {
        help();
        return Ok(ExitCode::SUCCESS);
    }

    let args: &'static Vec<String> = Box::leak(Box::new(args));
    let mut opts = Options::new(args.iter().map(String::as_str));
    let mut arch = ScmpArch::Native; // -a
    let mut ghost = false; // -g
    let mut probe = false; // -p
    let mut tmout = Duration::from_secs(3); // -t

    while let Some(opt) = opts.next_opt()? {
        match opt {
            Opt::Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Opt::Short('p') => probe = true,
            Opt::Short('g') => {
                // -g doesn't make sense without -p.
                ghost = true;
                probe = true;
            }
            Opt::Short('t') => {
                tmout = match opts.value() {
                    Ok(value) => match value.parse::<f64>() {
                        Ok(value) => Duration::from_secs_f64(value),
                        Err(error) => {
                            eprintln!("Invalid timeout for -t: {error}");
                            return Ok(ExitCode::FAILURE);
                        }
                    },
                    Err(_) => {
                        eprintln!("-t requires an argument!");
                        return Ok(ExitCode::FAILURE);
                    }
                };
            }
            Opt::Short('a') => {
                let value = match opts.value() {
                    Ok(value) => value,
                    Err(_) => {
                        eprintln!("-a requires an argument!");
                        eprintln!("Do '-a list' to print the list of architectures.");
                        return Ok(ExitCode::FAILURE);
                    }
                };
                if matches!(value.to_ascii_lowercase().as_str(), "help" | "list") {
                    syd::print_seccomp_architectures();
                    return Ok(ExitCode::SUCCESS);
                }
                arch = match ScmpArch::from_str(&format!(
                    "SCMP_ARCH_{}",
                    value.to_ascii_uppercase()
                )) {
                    Ok(arch) => arch,
                    Err(_) => {
                        eprintln!("Invalid architecture `{value}'");
                        return Ok(ExitCode::FAILURE);
                    }
                };
            }
            Opt::Short(opt) => {
                eprintln!("Unknown option: -{opt}!");
                help();
                return Ok(ExitCode::FAILURE);
            }
            Opt::Long(opt) => {
                eprintln!("Unknown option: --{opt}!");
                help();
                return Ok(ExitCode::FAILURE);
            }
        }
    }

    let sysarg = if let Some(value) = opts.positionals().next() {
        value
    } else {
        eprintln!("Expected syscall number or name regex as first argument!");
        return Ok(ExitCode::FAILURE);
    };

    let syscalls = match sysarg.parse::<i32>() {
        Ok(num) => {
            let syscall = ScmpSyscall::from(num);
            if !probe {
                if let Ok(name) = syscall.get_name_by_arch(arch) {
                    println!("{num}\t{name}");
                    return Ok(ExitCode::SUCCESS);
                } else {
                    return Ok(ExitCode::FAILURE);
                }
            }
            vec![syscall]
        }
        Err(_) => {
            let glob = if !is_literal(sysarg.as_bytes()) {
                sysarg.to_string()
            } else {
                format!("*{sysarg}*")
            };
            let mut ok = false;
            let mut syscalls = vec![];
            for (call, name) in (0..1024)
                .map(|n| {
                    let call = ScmpSyscall::from(n);
                    (call, call.get_name_by_arch(arch).unwrap_or_default())
                })
                .filter(|(_, name)| !name.is_empty())
            {
                if wildmatch(glob.as_bytes(), name.as_bytes()) {
                    if probe {
                        syscalls.push(call);
                    } else {
                        let num = i32::from(call);
                        println!("{num}\t{name}");
                        ok = true;
                    }
                }
            }
            if !probe {
                return Ok(if ok {
                    ExitCode::SUCCESS
                } else {
                    ExitCode::FAILURE
                });
            }
            syscalls
        }
    };

    // Probe
    if ghost {
        if let Err(errno) = enable_ghost_mode() {
            eprintln!("Failed to enable Syd's Ghost mode: {errno}");
            if errno == Errno::ENOENT {
                eprintln!("Ensure you're running under Syd, and the sandbox lock is off.");
            }
            return Ok(ExitCode::FAILURE);
        }
    }

    // Prepare system call arguments.
    let mut args: [Option<nix::libc::c_long>; 6] = [None; 6];
    #[allow(clippy::needless_range_loop)]
    for argc in 0..6 {
        if let Some(value) = opts.positionals().next() {
            args[argc] = match value.parse::<nix::libc::c_long>() {
                Ok(value) => Some(value),
                Err(error) => {
                    eprintln!("Argument {argc} is invalid: {error}");
                    return Ok(ExitCode::FAILURE);
                }
            };
        } else {
            break;
        }
    }

    for syscall in syscalls {
        println!("{}", probe_syscall(syscall, &args, tmout));
    }

    Ok(ExitCode::SUCCESS)
}

fn help() {
    println!("Usage: syd-sys [-hgpt] [-a list|native|x86|x86_64|aarch64...] number|name-glob [<probe-args>...]");
    println!("Given a number, print the matching syscall name and exit.");
    println!("Given a glob, print case-insensitively matching syscall names and exit.");
    println!("Given -p, probe the system call and print result.");
    println!("Given -g with -p, enable Syd's Ghost mode prior to probing.");
    println!("Specify syscall probe timeout in seconds, defaults to 3 seconds.");
}

fn probe_syscall(
    syscall: ScmpSyscall,
    args: &[Option<nix::libc::c_long>; 6],
    timeout: Duration,
) -> String {
    let snum = i32::from(syscall);
    let name = syscall.get_name().unwrap_or(snum.to_string());
    let argc = args
        .iter()
        .enumerate()
        .rev()
        .find(|&(_, elem)| elem.is_some())
        .map_or(0, |(idx, _)| idx + 1);
    #[allow(clippy::disallowed_methods)]
    match unsafe { fork() }.expect("fork") {
        ForkResult::Child => unsafe {
            match argc {
                0 => nix::libc::syscall(snum.into()),
                1 => nix::libc::syscall(snum.into(), args[0].unwrap()),
                2 => nix::libc::syscall(snum.into(), args[0].unwrap(), args[1].unwrap()),
                3 => nix::libc::syscall(
                    snum.into(),
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                ),
                4 => nix::libc::syscall(
                    snum.into(),
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                    args[3].unwrap(),
                ),
                5 => nix::libc::syscall(
                    snum.into(),
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                    args[3].unwrap(),
                    args[4].unwrap(),
                ),
                6 => nix::libc::syscall(
                    snum.into(),
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                    args[3].unwrap(),
                    args[4].unwrap(),
                    args[5].unwrap(),
                ),
                _ => unreachable!(),
            };
            nix::libc::_exit(Errno::last() as i32);
        },
        ForkResult::Parent { child, .. } => {
            let start = std::time::Instant::now();

            let result = loop {
                match waitpid(child, Some(WaitPidFlag::WNOHANG)) {
                    Ok(WaitStatus::Exited(_, code)) => {
                        if code == 0 {
                            break "0".to_string();
                        } else {
                            break errstr(code).to_string();
                        }
                    }
                    Ok(WaitStatus::Signaled(_, sig, core)) => {
                        if core {
                            break format!("{sig}!");
                        } else {
                            break format!("{sig}");
                        }
                    }
                    Ok(WaitStatus::StillAlive) => {
                        if start.elapsed() >= timeout {
                            let _ = kill(child, Signal::SIGKILL);
                            break "TMOUT".to_string();
                        }
                    }
                    Err(Errno::ECHILD) => break "ECHILD".to_string(),
                    _ => {}
                }
            };

            match argc {
                0 => format!("{name}()={result}"),
                1 => format!("{name}(0x{:x})={result}", args[0].unwrap()),
                2 => format!(
                    "{name}(0x{:x}, 0x{:x})={result}",
                    args[0].unwrap(),
                    args[1].unwrap()
                ),
                3 => format!(
                    "{name}(0x{:x}, 0x{:x}, 0x{:x})={result}",
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap()
                ),
                4 => format!(
                    "{name}(0x{:x}, 0x{:x}, 0x{:x}, 0x{:x})={result}",
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                    args[3].unwrap()
                ),
                5 => format!(
                    "{name}(0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x})={result}",
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                    args[3].unwrap(),
                    args[4].unwrap()
                ),
                6 => format!(
                    "{name}(0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x}, 0x{:x})={result}",
                    args[0].unwrap(),
                    args[1].unwrap(),
                    args[2].unwrap(),
                    args[3].unwrap(),
                    args[4].unwrap(),
                    args[5].unwrap()
                ),
                _ => unreachable!(),
            }
        }
    }
}

fn enable_ghost_mode() -> Result<(), Errno> {
    match lstat("/dev/syd/ghost") {
        Err(Errno::EOWNERDEAD) => Ok(()),
        Err(errno) => Err(errno),
        Ok(_) => Err(Errno::EOWNERDEAD),
    }
}

fn errstr(errno: i32) -> String {
    if let Some((name, _)) = Errno::from_raw(errno).to_string().split_once(':') {
        name.to_string()
    } else {
        errno.to_string()
    }
}
