//
// Syd: rock-solid unikernel
// src/syd-elf.rs: Syd's ELF information utility
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{fs::File, os::fd::AsRawFd, process::ExitCode};

use getargs::{Opt, Options};
use libseccomp::{scmp_cmp, ScmpAction, ScmpFilterContext, ScmpSyscall};
use nix::errno::Errno;
use syd::{
    elf::{ElfError, ElfType, ExecutableFile, LinkingType},
    err::SydResult,
    proc::proc_mmap_min_addr,
};

// 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 etyp = false;
    let mut is_32bit = false;
    let mut is_64bit = false;
    let mut is_dynamic = false;
    let mut is_static = false;
    let mut is_pie = false;
    let mut is_script = false;
    let mut is_xstack = false;

    while let Some(opt) = opts.next_opt()? {
        match opt {
            Opt::Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Opt::Short('3') => {
                is_32bit = true;
            }
            Opt::Short('6') => {
                is_64bit = true;
            }
            Opt::Short('d') => {
                is_dynamic = true;
            }
            Opt::Short('s') => {
                is_static = true;
            }
            Opt::Short('p') => {
                is_pie = true;
            }
            Opt::Short('x') => {
                is_script = true;
            }
            Opt::Short('X') => {
                is_xstack = true;
            }
            Opt::Short('t') => {
                etyp = true;
            }
            Opt::Short(opt) => {
                eprintln!("Unknown option: -{opt}!");
                return Err(Errno::EINVAL.into());
            }
            Opt::Long(opt) => {
                eprintln!("Unknown option: --{opt}!");
                return Err(Errno::EINVAL.into());
            }
        }
    }

    let flags = [
        is_32bit, is_64bit, is_dynamic, is_static, is_pie, etyp, is_script, is_xstack,
    ];
    let info = match flags.iter().filter(|&&flag| flag).count() {
        0 => true,
        1 => false,
        _ => {
            eprintln!("At most one of -3, -6, -d, -s, -p, -t, -x and -X must be given!");
            return Err(Errno::EINVAL.into());
        }
    };

    let path = match opts.positionals().next() {
        None => {
            eprintln!("Expected an ELF file as the first argument!");
            return Err(Errno::EINVAL.into());
        }
        Some(value) => value,
    };

    let check_linking = info || is_dynamic || is_static || is_pie || is_xstack;

    // Step 1: Open file.
    let file = File::open(path)?;

    // Step 2: Confine.
    confine(&file)?;

    // Step 3: Parse ELF.
    let exe = match ExecutableFile::parse(file, check_linking) {
        Ok(exe) => Some(exe),
        Err(ElfError::BadMagic) => None,
        Err(error) => return Err(error.into()),
    };

    // Step 4: Report result.
    if is_script {
        return Ok(match exe {
            Some(ExecutableFile::Script) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if is_32bit {
        return Ok(match exe {
            Some(ExecutableFile::Elf {
                elf_type: ElfType::Elf32,
                ..
            }) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if is_64bit {
        return Ok(match exe {
            Some(ExecutableFile::Elf {
                elf_type: ElfType::Elf64,
                ..
            }) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if is_dynamic {
        return Ok(match exe {
            Some(ExecutableFile::Elf {
                linking_type: Some(LinkingType::Dynamic),
                ..
            }) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if is_static {
        return Ok(match exe {
            Some(ExecutableFile::Elf {
                linking_type: Some(LinkingType::Static),
                ..
            }) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if is_pie {
        return Ok(match exe {
            Some(ExecutableFile::Elf { pie: true, .. }) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if is_xstack {
        return Ok(match exe {
            Some(ExecutableFile::Elf { xs: true, .. }) => ExitCode::SUCCESS,
            _ => ExitCode::FAILURE,
        });
    } else if etyp {
        let name = match exe {
            Some(ExecutableFile::Elf { file_type, .. }) => file_type.to_string(),
            Some(ExecutableFile::Script) => "script".to_string(),
            None => "unknown".to_string(),
        };
        println!("{name}");
        return Ok(ExitCode::SUCCESS);
    } else if let Some(exe) = exe {
        // Print ELF information or SCRIPT.
        println!("{path}:{exe}");
    } else {
        println!("{path}:UNKNOWN");
    }

    Ok(ExitCode::SUCCESS)
}

fn help() {
    println!("Usage: syd-elf [-36dhpstxX] binary|script");
    println!("Given a binary, file name and print ELF information.");
    println!("Given a script, print file name and \"SCRIPT\".");
    println!("The information line is a list of fields delimited by colons.");
    println!("Given -3, exit with success if the given binary is 32-bit.");
    println!("Given -6, exit with success if the given binary is 64-bit.");
    println!("Given -d, exit with success if the given binary is dynamically linked.");
    println!("Given -s, exit with success if the given binary is statically linked.");
    println!("Given -p, exit with success if the given binary is PIE.");
    println!("Given -t, print the type of the file.");
    println!("Given -x, exit with success if the given executable is a script.");
    println!("Given -X, exit with success if the given binary has executable stack.");
}

fn confine<Fd: AsRawFd>(fd: &Fd) -> SydResult<()> {
    // Step 1: Set up a Landlock sandbox to disallow all access.
    // Ignore errors as Landlock may not be supported.
    let p = vec![];
    let _ = syd::landlock_operation(&p, &p, syd::config::LANDLOCK_ABI);

    // Step 2: Ensure W^X via MDWE (if available) and seccomp-bpf.
    const PR_SET_MDWE: nix::libc::c_int = 65;
    const PR_MDWE_REFUSE_EXEC_GAIN: nix::libc::c_ulong = 1;

    // Ignore errors as
    // 1. MDWE may be unsupported -> EINVAL.
    // 2. MDWE may already be applied -> EPERM.
    // SAFETY: In libc, we trust.
    let _ = unsafe { nix::libc::prctl(PR_SET_MDWE, PR_MDWE_REFUSE_EXEC_GAIN, 0, 0, 0) };

    // W^X filter allows by default and kills offending memory access.
    let mut ctx = ScmpFilterContext::new_filter(ScmpAction::Allow)?;

    // Enforce the NO_NEW_PRIVS functionality before
    // loading the seccomp filter into the kernel.
    ctx.set_ctl_nnp(true)?;

    // Kill process for bad arch.
    ctx.set_act_badarch(ScmpAction::KillProcess)?;

    // Use a binary tree sorted by syscall number, if possible.
    let _ = ctx.set_ctl_optimize(2);

    // Restriction -1: Prevent mmap(addr<${mmap_min_addr}, MAP_FIXED).
    const MAP_FIXED: u64 = nix::libc::MAP_FIXED as u64;
    const MAP_FIXED_NOREPLACE: u64 = nix::libc::MAP_FIXED_NOREPLACE as u64;
    let mmap_min_addr = proc_mmap_min_addr().unwrap_or(4096);
    for sysname in ["mmap", "mmap2"] {
        #[allow(clippy::disallowed_methods)]
        let syscall = ScmpSyscall::from_name(sysname).unwrap();
        ctx.add_rule_conditional(
            ScmpAction::KillProcess,
            syscall,
            &[
                scmp_cmp!($arg0 < mmap_min_addr),
                scmp_cmp!($arg3 & MAP_FIXED == MAP_FIXED),
            ],
        )?;
        ctx.add_rule_conditional(
            ScmpAction::KillProcess,
            syscall,
            &[
                scmp_cmp!($arg0 < mmap_min_addr),
                scmp_cmp!($arg3 & MAP_FIXED_NOREPLACE == MAP_FIXED_NOREPLACE),
            ],
        )?;
    }

    // Restriction 0: Prohibit attempts to create memory mappings
    // that are writable and executable at the same time, or to
    // change existing memory mappings to become executable, or
    // mapping shared memory segments as executable.
    const W: u64 = nix::libc::PROT_WRITE as u64;
    const X: u64 = nix::libc::PROT_EXEC as u64;
    const WX: u64 = W | X;
    const MAP_S: u64 = nix::libc::MAP_SHARED as u64;
    for sysname in ["mmap", "mmap2"] {
        // Prevent writable and executable memory.
        #[allow(clippy::disallowed_methods)]
        let syscall = ScmpSyscall::from_name(sysname).unwrap();
        ctx.add_rule_conditional(
            ScmpAction::KillProcess,
            syscall,
            &[scmp_cmp!($arg2 & WX == WX)],
        )?;

        // Prevent executable shared memory.
        ctx.add_rule_conditional(
            ScmpAction::KillProcess,
            syscall,
            &[scmp_cmp!($arg2 & X == X), scmp_cmp!($arg3 & MAP_S == MAP_S)],
        )?;
    }

    ctx.add_rule_conditional(
        ScmpAction::KillProcess,
        ScmpSyscall::from_name("mprotect")?,
        &[scmp_cmp!($arg2 & X == X)],
    )?;

    // All set, load the seccomp filter.
    ctx.load()?;

    // Step 3: Confine system calls to least privilege using seccomp-bpf.
    let mut ctx = ScmpFilterContext::new_filter(ScmpAction::KillProcess)?;

    // Enforce the NO_NEW_PRIVS functionality before
    // loading the seccomp filter into the kernel.
    ctx.set_ctl_nnp(true)?;

    // Kill process for bad arch.
    ctx.set_act_badarch(ScmpAction::KillProcess)?;

    // Use a binary tree sorted by syscall number, if possible.
    let _ = ctx.set_ctl_optimize(2);

    // Allow base set.
    for sysname in [
        "brk",
        "exit",
        "exit_group",
        "madvise",
        "mmap",
        "mmap2",
        "mprotect",
        "mremap",
        "munmap",
        "rt_sigprocmask",
        "sigaltstack",
        "sigprocmask",
    ] {
        ctx.add_rule(ScmpAction::Allow, ScmpSyscall::from_name(sysname)?)?;
    }

    // Allow read, seek, close of file.
    let fd = fd.as_raw_fd() as u64;
    for sysname in ["close", "read", "readv", "_llseek", "lseek"] {
        ctx.add_rule_conditional(
            ScmpAction::Allow,
            ScmpSyscall::from_name(sysname)?,
            &[scmp_cmp!($arg0 == fd)],
        )?;
    }

    // Allow {g,s}etting file descriptor flags.
    const F_GETFD: u64 = nix::libc::F_GETFD as u64;
    const F_SETFD: u64 = nix::libc::F_SETFD as u64;
    for sysname in ["fcntl", "fcntl64"] {
        if let Ok(syscall) = ScmpSyscall::from_name(sysname) {
            for op in [F_GETFD, F_SETFD] {
                ctx.add_rule_conditional(
                    ScmpAction::Allow,
                    syscall,
                    &[scmp_cmp!($arg0 == fd), scmp_cmp!($arg1 == op)],
                )?;
            }
        }
    }

    // Allow writes to standard output and error.
    const FD_1: u64 = nix::libc::STDOUT_FILENO as u64;
    const FD_2: u64 = nix::libc::STDERR_FILENO as u64;
    for sysname in ["write", "writev"] {
        for fd in [FD_1, FD_2] {
            ctx.add_rule_conditional(
                ScmpAction::Allow,
                ScmpSyscall::from_name(sysname)?,
                &[scmp_cmp!($arg0 == fd)],
            )?;
        }
    }

    // All set, load the seccomp filter.
    ctx.load()?;

    Ok(())
}
