//
// Syd: rock-solid application kernel
// src/utils/syd-lock.rs: Run a command under Landlock
//
// Copyright (c) 2024, 2025, 2026 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0

use std::{
    os::unix::ffi::OsStrExt,
    process::{Command, ExitCode},
};

use memchr::arch::all::is_equal;
use nix::errno::Errno;
use syd::{
    confine::{lock_enabled, run_cmd},
    err::err2no,
    landlock::{
        AddRuleError, AddRulesError, CompatLevel, CreateRulesetError, Errata, RestrictSelfError,
        RestrictSelfFlags, RulesetError, RulesetStatus, ABI,
    },
    landlock_policy::{LandlockPolicy, LANDLOCK_ACCESS_FS, LANDLOCK_ACCESS_NET},
    parsers::sandbox::parse_landlock_cmd,
    path::XPathBuf,
};

// Set global allocator to GrapheneOS allocator.
#[cfg(all(
    not(coverage),
    not(feature = "prof"),
    not(target_os = "android"),
    not(target_arch = "riscv64"),
    target_page_size_4k,
    target_pointer_width = "64"
))]
#[global_allocator]
static GLOBAL: hardened_malloc::HardenedMalloc = hardened_malloc::HardenedMalloc;

// Set global allocator to tcmalloc if profiling is enabled.
#[cfg(feature = "prof")]
#[global_allocator]
static GLOBAL: tcmalloc::TCMalloc = tcmalloc::TCMalloc;

const COMPAT_LEVEL_HARD: &[&[u8]] = &[b"h", b"hard", b"hard-requirement"];
const COMPAT_LEVEL_SOFT: &[&[u8]] = &[b"s", b"soft", b"soft-requirement"];
const COMPAT_LEVEL_BEST: &[&[u8]] = &[b"b", b"best", b"best-effort"];

syd::main! {
    use lexopt::prelude::*;

    syd::set_sigpipe_dfl()?;

    // Parse CLI options.
    //
    // Note, option parsing is POSIXly correct:
    // POSIX recommends that no more options are parsed after the first
    // positional argument. The other arguments are then all treated as
    // positional arguments.
    // See: https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap12.html#tag_12_02
    let mut opt_abick = false;
    let mut opt_check = false;
    let mut opt_verbose = false;
    let mut opt_cmd = None;
    let mut opt_arg = Vec::new();
    let mut policy = LandlockPolicy {
        compat_level: Some(CompatLevel::HardRequirement),
        ..Default::default()
    };

    let mut parser = lexopt::Parser::from_env();
    while let Some(arg) = parser.next()? {
        match arg {
            Short('h') => {
                help();
                return Ok(ExitCode::SUCCESS);
            }
            Short('A') => opt_abick = true,
            Short('V') => opt_check = true,
            Short('v') => opt_verbose = true,
            // Interface to Landlock erratas.
            Short('E') => {
                let errata = parser.value()?;
                let errata = errata.as_bytes();

                if is_equal(errata, b"list") {
                    let errata = Errata::query();
                    for fix in errata {
                        match fix {
                            Errata::TCP_SOCKET_IDENTIFICATION => println!("tcp_socket_identification"),
                            Errata::SCOPED_SIGNAL_SAME_TGID => println!("scoped_signal_same_tgid"),
                            unknown => { // bundle of unknown bits.
                                let mut unknown = unknown.bits();
                                while unknown != 0 {
                                    let lsb = unknown & unknown.wrapping_neg();
                                    println!("{lsb:#x}");
                                    unknown &= unknown - 1;
                                }
                            }
                        }
                    }
                    return Ok(ExitCode::SUCCESS);
                }

                let requested = LandlockPolicy::parse_errata(errata)?;
                let supported = Errata::query();

                let code = if supported.contains(requested) {
                    0 // all errata are available.
                } else if supported.intersects(requested) {
                    1 // some errata is missing.
                } else {
                    2 // all errata is missing.
                };

                return Ok(ExitCode::from(code));
            }
            // Scoped signals and UNIX abstract sockets.
            Short('S') => policy.scoped_sig = true,
            Short('U') => policy.scoped_abs = true,
            // Interface to Landlock compatibility levels.
            Short('C') => {
                let level = parser.value()?;
                let level = level.as_bytes();

                if is_equal(level, b"list") {
                    println!("hard-requirement");
                    println!("soft-requirement");
                    println!("best-effort");
                    return Ok(ExitCode::SUCCESS);
                }

                if COMPAT_LEVEL_HARD.iter().any(|&l| is_equal(level, l)) {
                    policy.compat_level = Some(CompatLevel::HardRequirement);
                } else if COMPAT_LEVEL_SOFT.iter().any(|&l| is_equal(level, l)) {
                    policy.compat_level = Some(CompatLevel::SoftRequirement);
                } else if COMPAT_LEVEL_BEST.iter().any(|&l| is_equal(level, l)) {
                    policy.compat_level = Some(CompatLevel::BestEffort);
                } else {
                    return Err(Errno::EINVAL.into());
                }
            }
            // Interface to landlock_restrict_self(2) flags.
            Short('F') => {
                let flags = parser.value()?;
                let flags = flags.as_bytes();

                if is_equal(flags, b"list") {
                    for flag in RestrictSelfFlags::all().iter() {
                        println!("{flag}\t{}", flag.bits());
                    }
                    return Ok(ExitCode::SUCCESS);
                }

                let flags = LandlockPolicy::parse_restrict_self_flags(flags, true)?;
                policy.restrict_self_flags.insert(flags);
            }
            // New interface with refined categories.
            Short('l') => {
                let cmd = parser.value().map(XPathBuf::from)?;

                if cmd.is_equal(b"list") {
                    for (name, access) in LANDLOCK_ACCESS_FS.iter() {
                        println!("{name}\t{}", access.bits());
                    }
                    for (name, access) in LANDLOCK_ACCESS_NET.iter() {
                        println!("{name}\t{}", access.bits());
                    }
                    return Ok(ExitCode::SUCCESS);
                }

                let cmd = parse_landlock_cmd(&format!("allow/lock/{cmd}"))?;
                policy.edit(cmd, None)?;
            }
            // Old interface with practical read/write generalization.
            Short('r') => {
                let path = parser.value().map(XPathBuf::from)?;
                let cmd = parse_landlock_cmd(&format!("allow/lock/read,readdir,exec,ioctl+{path}"))?;
                policy.edit(cmd, None)?;
            }
            Short('w') => {
                let path = parser.value().map(XPathBuf::from)?;
                let cmd = parse_landlock_cmd(&format!("allow/lock/all+{path}"))?;
                policy.edit(cmd, None)?;
            }
            Short('b') => {
                let port = parser.value().map(XPathBuf::from)?;
                let cmd = parse_landlock_cmd(&format!("allow/lock/bind+{port}"))?;
                policy.edit(cmd, None)?;
            }
            Short('c') => {
                let port = parser.value().map(XPathBuf::from)?;
                let cmd = parse_landlock_cmd(&format!("allow/lock/connect+{port}"))?;
                policy.edit(cmd, None)?;
            }
            Value(prog) => {
                opt_cmd = Some(prog);
                opt_arg.extend(parser.raw_args()?);
            }
            _ => return Err(arg.unexpected().into()),
        }
    }

    if opt_abick && opt_check {
        eprintln!("-A and -V are mutually exclusive!");
        return Err(Errno::EINVAL.into());
    }

    let abi = ABI::new_current();
    if opt_abick {
        let abi = abi as i32 as u8;
        print!("{abi}");
        return Ok(ExitCode::from(abi));
    } else if opt_check {
        if abi == ABI::Unsupported {
            println!("Landlock is not supported.");
            return Ok(ExitCode::from(127));
        }

        let state = lock_enabled(abi);
        let state_verb = match state {
            0 => "fully enforced",
            1 => "partially enforced",
            2 => "not enforced",
            _ => "unsupported",
        };
        println!("Landlock ABI {} is {state_verb}.", abi as i32);
        return Ok(ExitCode::from(state));
    }

    // Prepare command or bail if not passed.
    let cmd = if let Some(cmd) = opt_cmd {
        cmd
    } else {
        help();
        return Ok(ExitCode::SUCCESS);
    };

    // Set up Landlock sandbox.
    macro_rules! vprintln {
        ($($arg:tt)*) => {
            if opt_verbose {
                eprintln!($($arg)*);
            }
        };
    }

    match policy.restrict_self(abi) {
        Ok(status) => match status.ruleset {
            RulesetStatus::FullyEnforced => {
                vprintln!("syd-lock: Landlock ABI {} is fully enforced.", abi as i32)
            }
            RulesetStatus::PartiallyEnforced => {
                vprintln!(
                    "syd-lock: Landlock ABI {} is partially enforced.",
                    abi as i32
                )
            }
            RulesetStatus::NotEnforced => {
                eprintln!("syd-lock: Landlock ABI {} is not enforced!", abi as i32);
                return Ok(ExitCode::FAILURE);
            }
        },
        Err(error) if policy.compat_level == Some(CompatLevel::BestEffort) => {
            eprintln!("syd-lock: Landlock ABI {} is unsupported: {error}!", abi as i32);
            eprintln!("syd-lock: Compatibility level is best effort, resuming...");
            // fall-through.
        }
        Err(
            RulesetError::AddRules(AddRulesError::Fs(AddRuleError::AddRuleCall { source, .. }))
            | RulesetError::AddRules(AddRulesError::Net(AddRuleError::AddRuleCall {
                source, ..
            }))
            | RulesetError::AddRules(AddRulesError::Scope(AddRuleError::AddRuleCall {
                source, ..
            })),
        ) => {
            let errno = err2no(&source);
            eprintln!("syd-lock: Landlock add rules error: {source}!");
            return Err(errno.into());
        }
        Err(RulesetError::CreateRuleset(CreateRulesetError::CreateRulesetCall {
            source, ..
        })) => {
            let errno = err2no(&source);
            eprintln!("syd-lock: Landlock create ruleset error: {source}!");
            return Err(errno.into());
        }
        Err(RulesetError::RestrictSelf(RestrictSelfError::SetNoNewPrivsCall {
            source, ..
        })) => {
            let errno = err2no(&source);
            eprintln!("syd-lock: Set no new privs error: {source}!");
            return Err(errno.into());
        }
        Err(RulesetError::RestrictSelf(RestrictSelfError::RestrictSelfCall { source, .. })) => {
            let errno = err2no(&source);
            eprintln!("syd-lock: Landlock restrict self error: {source}!");
            return Err(errno.into());
        }
        Err(source) => {
            eprintln!("syd-lock: Landlock handle accesses error: {source}!");
            return Err(Errno::ENOTSUP.into());
        }

    };

    // Execute command, /bin/sh by default.
    let mut cmd = Command::new(cmd);
    let cmd = cmd.args(opt_arg);
    Ok(ExitCode::from(run_cmd(cmd)))
}

fn help() {
    println!("Usage: syd-lock [-bchrvwASUV] [-C level] [-E errata] [-F flag]... [-l cat[,cat...]{{+|-}}path|port[-port]]... {{command [args...]}}");
    println!("Run a program under landlock(7).");
    println!("Use -v to increase verbosity.");
    println!("Use -A to exit with Landlock ABI version.");
    println!("Use -V to check for Landlock support.");
    println!("Use -l cat[,cat...]{{+|-}}path|port[-port] to specify categories with path or closed port range.");
    println!("Use -C level to set ABI compatibility level, one of hard-requirement, soft-requirement, best-effort.");
    println!("Use -E errata to check for fixes in current ABI. Argument may be a name or number.");
    println!("Use -F flags to set landlock_restrict_self(2) flags.");
    println!("Use -S to enabled scoped signals.");
    println!("Use -U to enabled scoped UNIX abstract sockets.");
    println!(
        "Use `list' with -l, -C, -E, -F to list categories, compat-levels, erratas and flags."
    );
    println!("Use -r path as a shorthand for -l read,readdir,exec,ioctl+path.");
    println!("Use -w path as a shorthand for -l all+path.");
    println!("Use -b port as a shorthand for -l bind+port.");
    println!("Use -c port as a shorthand for -l connect+port.");
    println!("Categories:");
    println!("\tall   = *");
    println!("\trpath = read + readdir");
    println!("\twpath = write + truncate");
    println!("\tcpath = create + delete + rename");
    println!("\tdpath = mkbdev + mkcdev");
    println!("\tspath = mkfifo + symlink");
    println!("\ttpath = mkdir + rmdir");
    println!("\tinet = bind + connect");
    println!("\t        ioctl");
    println!(
        "Refer to the \"Sandboxing\" and \"Lock Sandboxing\" sections of the syd(7) manual page."
    );
}
