//
// pandora: syd's Dump Inspector & Profile Writer
// pandora.rs: Main entry point
//
// Copyright (c) 2021, 2024 Ali Polatel <alip@exherbo.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

#![allow(clippy::disallowed_methods)]

use std::{
    collections::{HashMap, HashSet},
    fs::{File, OpenOptions},
    hash::Hasher,
    io::{BufRead, BufReader},
    iter::FromIterator,
    net::IpAddr,
    os::fd::AsRawFd,
    path::Path,
    process::{exit, Command, ExitCode},
    thread,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use clap::{Arg, ArgAction};
use dns_lookup::lookup_addr;
use humantime::parse_duration;
use nix::{
    errno::Errno,
    libc::pid_t,
    sys::{
        signal::{kill, sigprocmask, SigmaskHow, Signal},
        signalfd::SigSet,
    },
    unistd::Pid,
};
use serde::{Deserialize, Serialize};
use time::{format_description, OffsetDateTime};

pub mod built_info {
    // The file has been placed there by the build script.
    include!(concat!(env!("OUT_DIR"), "/built.rs"));
}

bitflags::bitflags! {
    #[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash)]
    pub struct Capability: u32 {
        const CAP_STAT = 1 << 0;
        const CAP_READ = 1 << 1;
        const CAP_EXEC = 1 << 2;
        const CAP_WRITE = 1 << 3;
        const CAP_IOCTL = 1 << 4;
        const CAP_NET_CONNECT = 1 << 30;
        const CAP_NET_BIND = 1 << 31;
    }
}

#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(untagged)]
enum Access {
    InetAddr {
        ctx: String,
        cap: String,
        addr: String,
        t: u64,
    },
    UnixAddr {
        ctx: String,
        cap: String,
        unix: String,
        t: u64,
    },
    Path {
        ctx: String,
        cap: String,
        path: String,
        bins: Option<Vec<String>>,
        t: u64,
    },
    Run {
        ctx: String,
        command: String,
        args: Vec<String>,
        t: u64,
    },
    Any {
        ctx: String,
    },
}

fn command_profile<'b>(
    bin: &'b str,
    cmd: &[&'b str],
    output_path: &'b str,
    path_limit: u8,
    cmd_timeout: Option<Duration>,
    config: Option<&[&'b str]>,
) -> u8 {
    if Path::new(output_path).exists() {
        eprintln!("pandora: Output file {output_path} exists, bailing out.");
        return 1;
    }

    let (fd_rd, fd_rw) = match nix::unistd::pipe() {
        Ok((fd_rd, fd_rw)) => (fd_rd, fd_rw),
        Err(error) => {
            eprintln!("pandora: error creating pipe: {}", error);
            return 1;
        }
    };

    let log_fd = fd_rw.as_raw_fd().to_string();
    let mut syd = Command::new(bin);
    syd.env("SYD_NO_SYSLOG", "1");
    syd.env("SYD_LOG", "info");
    syd.env("SYD_LOG_FD", log_fd);
    syd.arg("-x");
    syd.arg("-ppandora");
    if let Some(config) = config {
        let args: Vec<String> = config.iter().map(|arg| format!("-m{arg}")).collect();
        syd.args(args);
    }
    syd.arg("--").args(cmd);
    let mut child = syd.spawn().expect("syd command failed to start");

    // Block SIGINT in the parent process.
    let mut mask = SigSet::empty();
    mask.add(Signal::SIGINT);
    sigprocmask(SigmaskHow::SIG_BLOCK, Some(&mask), None).expect("Failed to block signals");

    if let Some(cmd_timeout) = cmd_timeout {
        let pid = Pid::from_raw(child.id() as pid_t);
        thread::spawn(move || {
            thread::sleep(cmd_timeout);
            eprintln!("pandora: Timeout expired, terminating process...");
            let _ = kill(pid, Signal::SIGTERM);
        });
    }

    drop(fd_rw); // close the write end of the pipe.
    let input = Box::new(std::io::BufReader::new(std::fs::File::from(fd_rd)));
    let r = do_inspect(input, output_path, path_limit, config);

    child.wait().expect("failed to wait for syd");
    eprintln!("pandora: Profile has been written to {output_path}.");
    eprintln!("pandora: To use it, do: syd -P {output_path} command args...");

    r
}

fn command_inspect(input_path: &str, output_path: &str, path_limit: u8) -> u8 {
    let input = open_input(input_path);
    do_inspect(input, output_path, path_limit, None)
}

fn main() -> ExitCode {
    let matches = clap::Command::new(built_info::PKG_NAME)
        .about(built_info::PKG_DESCRIPTION)
        .author(built_info::PKG_AUTHORS)
        .version(built_info::PKG_VERSION)
        .arg_required_else_help(true)
        .help_expected(true)
        .next_line_help(false)
        .infer_long_args(true)
        .infer_subcommands(true)
        .propagate_version(true)
        .subcommand_required(true)
        .max_term_width(80)
        .help_template(
            r#"
{before-help}{name} {version}
{about}
Copyright (c) 2023, 2024 {author}
SPDX-License-Identifier: GPL-3.0-or-later

{usage-heading} {usage}

{all-args}{after-help}
"#,
        )
        .after_help(format!(
            "\
Hey you, out there beyond the wall,
Breaking bottles in the hall,
Can you help me?

Send bug reports to {}
Attaching poems encourages consideration tremendously.

License: {}
Homepage: {}
Repository: {}
",
            built_info::PKG_AUTHORS,
            built_info::PKG_LICENSE,
            built_info::PKG_HOMEPAGE,
            built_info::PKG_REPOSITORY,
        ))
        .subcommand(
            clap::Command::new("profile")
                .about("Execute a program under inspection and write a syd profile")
                .arg(
                    Arg::new("bin")
                        .default_value("syd")
                        .help("Path to syd binary")
                        .long("bin")
                        .env("SYD_BIN")
                        .num_args(1),
                )
                .arg(
                    Arg::new("magic")
                        .action(ArgAction::Append)
                        .help("Run a sandbox command during init, may be repeated")
                        .long("magic")
                        .short('m')
                        .num_args(1),
                )
                .arg(
                    Arg::new("output")
                        .default_value("./out.syd-3")
                        .help("Path to syd profile output")
                        .long("output")
                        .short('o')
                        .env("PANDORA_OUT")
                        .num_args(1),
                )
                .arg(
                    Arg::new("limit")
                        .default_value("3")
                        .required(false)
                        .help("Maximum number of path members before trim, 0 to disable")
                        .long("limit")
                        .short('l')
                        .value_parser(clap::value_parser!(u64).range(0..=u64::from(u8::MAX))),
                )
                .arg(
                    Arg::new("timeout")
                        .required(false)
                        .help("Human-formatted timeout duration")
                        .long("timeout")
                        .short('t')
                        .value_parser(|s: &str| parse_duration(s).map_err(|e| e.to_string())),
                )
                .arg(
                    Arg::new("cmd")
                        .required(true)
                        .help("Command to run under syd")
                        .num_args(1..),
                ),
        )
        .subcommand(
            clap::Command::new("inspect")
                .about("Read a syd core dump and write a syd profile")
                .arg(
                    Arg::new("input")
                        .required(true)
                        .help("Path to syd core dump")
                        .long("input")
                        .short('i'),
                )
                .arg(
                    Arg::new("output")
                        .default_value("./out.syd-3")
                        .required(true)
                        .help("Path to syd profile output")
                        .long("output")
                        .short('o')
                        .env("PANDORA_OUT"),
                )
                .arg(
                    Arg::new("limit")
                        .default_value("7")
                        .required(false)
                        .help("Maximum number of path members before trim, 0 to disable")
                        .long("limit")
                        .short('l')
                        .value_parser(clap::value_parser!(u64).range(0..=u64::from(u8::MAX))),
                ),
        )
        .get_matches();

    let (subcommand, submatches) = matches.subcommand().expect("missing subcommand");

    match subcommand {
        "profile" => {
            let bin = submatches.get_one::<String>("bin").expect("bin");
            let out = submatches.get_one::<String>("output").expect("output");
            let limit = *submatches.get_one::<u64>("limit").expect("limit") as u8;
            let timeout = submatches.get_one::<Duration>("timeout").copied();
            let cmd: Vec<&str> = submatches
                .get_many::<String>("cmd")
                .expect("cmd")
                .map(|s| s.as_str())
                .collect();
            let magic: Option<Vec<&str>> = if submatches.contains_id("magic") {
                Some(
                    submatches
                        .get_many::<String>("magic")
                        .expect("magic")
                        .map(|s| s.as_str())
                        .collect(),
                )
            } else {
                None
            };
            ExitCode::from(command_profile(
                bin,
                &cmd,
                out,
                limit,
                timeout,
                magic.as_deref(),
            ))
        }
        "inspect" => {
            let input = submatches.get_one::<String>("input").expect("input");
            let output = submatches.get_one::<String>("output").expect("output");
            let limit = *submatches.get_one::<u64>("limit").expect("limit") as u8;
            ExitCode::from(command_inspect(input, output, limit))
        }
        _ => unreachable!(),
    }
}

fn do_inspect(
    input: Box<dyn std::io::BufRead>,
    output_path: &str,
    path_limit: u8,
    config: Option<&[&str]>,
) -> u8 {
    let mut output = open_output(output_path);
    let mut magic = HashMap::<String, Capability>::new();
    let mut force = HashSet::<String>::new();
    let mut program_invocation_name = "?".to_string();
    let mut program_command_line = vec![];
    let mut program_startup_time = UNIX_EPOCH;

    for line in input.lines() {
        let serialized = match line {
            Ok(line) if line.is_empty() => {
                break; /* EOF */
            }
            Ok(line) => line,
            Err(error) => {
                eprintln!("pandora: failed to read line from input: {error}");
                return 1;
            }
        };

        // Parse JSON
        if let Some((comm, args, timestamp)) =
            parse_json_line(&serialized, &mut magic, &mut force, path_limit)
        {
            program_invocation_name = comm;
            program_command_line = args;
            program_startup_time = timestamp;
        }
    }

    let mut data = String::new();
    let config = config
        .map(|config| config.join("\n"))
        .unwrap_or("".to_string());
    if !config.is_empty() {
        data.push_str("###\n# User submitted options\n###\n");
        data.push_str(&config);
        data.push('\n');
    }

    /* Step 1: Print out the magic header. */
    writeln!(
        &mut output,
        "#
# syd profile generated by pandora-{}
# Date: {}

{}
###
# Auto-generated magic entries
# Program: {}
# Arguments: {:?}
###
",
        built_info::PKG_VERSION,
        format_system_time(program_startup_time),
        data,
        program_invocation_name,
        program_command_line,
    )
    .unwrap_or_else(|_| panic!("failed to print header to output »{}«", output_path));

    /* Step 2: Print out magic entries */
    let mut list = Vec::from_iter(magic);
    // Secondary alphabetical sort.
    list.sort_by_key(|(path, _)| path.to_string());
    // Primary: sort reverse by Capability
    list.sort_by_key(|(_, capability)| std::cmp::Reverse(*capability));

    let mut lastcap: Option<Capability> = None;
    for entry in list {
        if let Some(cap) = lastcap {
            if entry.1 != cap {
                writeln!(&mut output, "").unwrap();
                lastcap = Some(entry.1);
            }
        } else {
            lastcap = Some(entry.1);
        }

        let mut done = false;
        if entry.1.contains(Capability::CAP_NET_BIND) {
            if entry.0.starts_with('/') {
                // UNIX abstract/domain socket
                writeln!(&mut output, "allow/net/bind+{}", entry.0).unwrap();
            } else {
                // IPv{4,6} address
                let ip = entry.0.splitn(2, '!').next().unwrap();
                let ip = ip.parse::<IpAddr>().unwrap_or_else(|e| {
                    panic!("Failed to parse IP address `{}': {}", ip, e);
                });
                if let Ok(host) = lookup_addr(&ip) {
                    writeln!(&mut output, "# {host}").unwrap();
                }
                writeln!(&mut output, "allow/net/bind+{}", entry.0).unwrap();
            }
            done = true;
        }
        if entry.1.contains(Capability::CAP_NET_CONNECT) {
            if entry.0.starts_with('/') {
                // UNIX abstract/domain socket
                writeln!(&mut output, "allow/net/connect+{}", entry.0).unwrap();
            } else {
                let ip = entry.0.splitn(2, '!').next().unwrap();
                let ip = ip.parse::<IpAddr>().unwrap_or_else(|e| {
                    panic!("Failed to parse IP address `{}': {}", ip, e);
                });
                if let Ok(host) = lookup_addr(&ip) {
                    writeln!(&mut output, "# {host}").unwrap();
                }
                writeln!(&mut output, "allow/net/connect+{}", entry.0).unwrap();
            }
            done = true;
        }

        if done {
            continue;
        }

        let mut caps = vec![];
        if entry.1.contains(Capability::CAP_IOCTL) {
            caps.push("ioctl");
        }
        if entry.1.contains(Capability::CAP_WRITE) {
            caps.push("write")
        }
        if entry.1.contains(Capability::CAP_EXEC) {
            caps.push("exec")
        }
        if entry.1.contains(Capability::CAP_READ) {
            caps.push("read")
        }
        if entry.1.contains(Capability::CAP_STAT) {
            caps.push("stat")
        }
        assert!(!caps.is_empty(), "Invalid rule!");

        writeln!(&mut output, "allow/{}+{}", caps.join(","), entry.0).unwrap();
    }

    /* Step 3: Print Force entries if available. */
    if !force.is_empty() {
        write!(&mut output, "\n###\n# Auto-generated force entries\n###").unwrap();
        let mut force: Vec<_> = force.iter().collect();
        force.sort_by_cached_key(|arg| (arg.len(), arg.to_string()));
        for entry in force {
            if let Some(line) = path2force(entry) {
                write!(&mut output, "\n{line}").unwrap();
            }
        }
        writeln!(&mut output, "").unwrap();
    }

    writeln!(
        &mut output,
        "\n# Turn on binary verification.\nsandbox/force:on"
    )
    .unwrap_or_else(|_| {
        panic!(
            "failed to turn on binary verification for output »{}«",
            output_path
        )
    });
    writeln!(&mut output, "\n# Lock configuration.\nlock:on")
        .unwrap_or_else(|_| panic!("failed to lock configuration for output »{}«", output_path));

    0
}

#[allow(clippy::type_complexity)]
fn parse_json_line(
    serialized: &str,
    magic: &mut HashMap<String, Capability>,
    force: &mut HashSet<String>,
    path_limit: u8,
) -> Option<(String, Vec<String>, SystemTime)> {
    match serde_json::from_str(serialized)
        .unwrap_or_else(|e| panic!("failed to parse line: »{}«", e))
    {
        Access::Path {
            cap, path, bins, ..
        } => {
            for c in cap.chars() {
                let capability = match c {
                    'r' => Capability::CAP_READ,
                    's' => Capability::CAP_STAT,
                    'w' => Capability::CAP_WRITE,
                    'f' | 'x' => Capability::CAP_EXEC,
                    'i' => Capability::CAP_IOCTL,
                    _ => unreachable!("Undefined capability `{}`!", c),
                };
                if capability == Capability::CAP_EXEC {
                    force.insert(path.clone());
                }
                magic
                    .entry(process_path(&path, path_limit))
                    .or_insert_with(Capability::empty)
                    .insert(capability);
            }
            if let Some(bins) = bins {
                for path in bins {
                    force.insert(path.clone());
                    magic
                        .entry(process_path(&path, path_limit))
                        .or_insert_with(Capability::empty)
                        .insert(Capability::CAP_EXEC);
                }
            }
        }
        Access::InetAddr { cap, addr, .. }
        | Access::UnixAddr {
            cap, unix: addr, ..
        } => {
            let capability = match cap.as_str() {
                "b" => Capability::CAP_NET_BIND,
                "c" => Capability::CAP_NET_CONNECT,
                _ => unreachable!(),
            };
            magic
                .entry(addr)
                .or_insert_with(Capability::empty)
                .insert(capability);
        }
        Access::Run {
            command, args, t, ..
        } => {
            return Some((command, args, UNIX_EPOCH + Duration::from_secs(t)));
        }
        _ => {}
    };

    None
}

fn open_input(path_or_stdin: &str) -> Box<dyn std::io::BufRead> {
    match path_or_stdin {
        "-" => Box::new(std::io::BufReader::new(std::io::stdin())),
        path => Box::new(std::io::BufReader::new(
            match OpenOptions::new().read(true).open(path) {
                Ok(file) => file,
                Err(error) => {
                    eprintln!("pandora: Failed to open file »{path}«: {error}");
                    exit(1);
                }
            },
        )),
    }
}

fn open_output(path_or_stdout: &str) -> Box<dyn std::io::Write> {
    match path_or_stdout {
        "-" => Box::new(std::io::BufWriter::new(std::io::stdout())),
        path => Box::new(std::io::BufWriter::new(
            match OpenOptions::new().write(true).create_new(true).open(path) {
                Ok(file) => file,
                Err(error) => {
                    eprintln!("pandora: Failed to open file »{path}«: {error}");
                    exit(1);
                }
            },
        )),
    }
}

fn process_path(path: &str, limit: u8) -> String {
    if limit == 0 || path == "/" {
        path.to_string()
    } else if let Some(glob) = path2glob(path) {
        glob
    } else {
        let limit = limit as usize;
        let members: Vec<&str> = path.split('/').filter(|&x| !x.is_empty()).collect();
        if limit > 0 && limit < members.len() {
            format!("/{}/***", members[0..limit].join("/"))
        } else {
            format!("/{}", members.join("/"))
        }
    }
}

fn path2force(path: &str) -> Option<String> {
    let file = BufReader::new(File::open(path).ok()?);
    let hash = const_hex::encode(hash(file).ok()?);

    Some(format!("force+{path}:{hash}:kill"))
}

fn path2glob(path: &str) -> Option<String> {
    let components: Vec<&str> = path.split('/').collect();
    let mut new_path = String::new();
    let mut handled = false;

    if path.starts_with("/proc/") {
        if components.len() >= 3 && components[2].chars().all(char::is_numeric) {
            if components.len() > 4
                && components[4].chars().all(char::is_numeric)
                && components[3] == "task"
            {
                // Handle the /proc/$pid/task/$tid/... case
                let rest_of_path = if components.len() > 5 {
                    format!("/{}", components[5..].join("/"))
                } else {
                    String::new()
                };
                new_path = format!("/proc/[0-9]*/task/[0-9]*{}", rest_of_path);
                handled = true;

                // Specifically handle the /proc/$pid/task/$tid/{fd,ns}/... cases.
                if components.len() > 5 && components[5] == "fd" {
                    let fd_rest_of_path = if components.len() > 6 {
                        format!("/{}", components[6..].join("/"))
                    } else {
                        String::new()
                    };
                    new_path = format!("/proc/[0-9]*/task/[0-9]*/fd{}", fd_rest_of_path);
                } else if components.len() > 5 && components[5] == "ns" {
                    let ns_rest_of_path = if components.len() > 6 {
                        format!("/{}", components[6..].join("/"))
                    } else {
                        String::new()
                    };
                    new_path = format!("/proc/[0-9]*/task/[0-9]*/ns{}", ns_rest_of_path);
                }
            } else {
                // Handle the general /proc/$pid/... case
                let rest_of_path = if components.len() > 3 {
                    format!("/{}", components[3..].join("/"))
                } else {
                    String::new()
                };
                new_path = format!("/proc/[0-9]*{}", rest_of_path);
                handled = true;

                // Specifically handle the /proc/$pid/{fd,ns}/... cases.
                if components.len() > 3 && components[3] == "fd" {
                    let fd_rest_of_path = if components.len() > 4 {
                        format!("/{}", components[4..].join("/"))
                    } else {
                        String::new()
                    };
                    new_path = format!("/proc/[0-9]*/fd{}", fd_rest_of_path);
                } else if components.len() > 3 && components[3] == "ns" {
                    let ns_rest_of_path = if components.len() > 4 {
                        format!("/{}", components[4..].join("/"))
                    } else {
                        String::new()
                    };
                    new_path = format!("/proc/[0-9]*/ns{}", ns_rest_of_path);
                }
            }
        }

        // Further handle /{fd,ns}/... parts.
        if new_path.contains("/fd/") || new_path.contains("/ns/") {
            let mut final_path = String::new();
            let fd_components: Vec<&str> = new_path.split('/').collect();
            for (i, component) in fd_components.iter().enumerate() {
                if i > 0 {
                    final_path.push('/');
                }
                if i == fd_components.len() - 1 && component.chars().all(char::is_numeric) {
                    // Convert numeric fd/ns component to [0-9]*.
                    final_path.push_str("[0-9]*");
                } else if component.contains(':') {
                    // Handle foo:[number] pattern
                    let parts: Vec<&str> = component.split(':').collect();
                    if parts.len() == 2 && parts[1].starts_with('[') && parts[1].ends_with(']') {
                        let inner = &parts[1][1..parts[1].len() - 1];
                        if inner.chars().all(char::is_numeric) {
                            final_path.push_str(&format!("{}:[0-9]*", parts[0]));
                            continue;
                        }
                    }
                    final_path.push_str(component);
                } else {
                    final_path.push_str(component);
                }
            }
            return Some(final_path);
        }
    }

    if handled {
        return Some(new_path);
    }

    // Handle /dev/pts/[number] case
    if path.starts_with("/dev/pts/") {
        if path.split('/').count() == 4
            && path
                .split('/')
                .nth(3)
                .unwrap()
                .chars()
                .all(char::is_numeric)
        {
            return Some("/dev/pts/[0-9]*".to_string());
        } else {
            return None;
        }
    }

    // Handle /dev/tty case
    if path.starts_with("/dev/tty") {
        return Some("/dev/tty*".to_string());
    }

    // Return None if no cases match
    None
}

fn format_system_time(system_time: SystemTime) -> String {
    let datetime = OffsetDateTime::from(system_time);
    let format =
        format_description::parse("[year]-[month]-[day] [hour]:[minute]:[second].[subsecond]")
            .unwrap();
    datetime.format(&format).unwrap()
}

/// Calculate sha{1,256,512} of the given buffered reader.
/// Returns a byte array.
fn hash<R: BufRead>(mut reader: R) -> Result<Vec<u8>, Errno> {
    let mut hasher = rs_sha3_512::Sha3_512Hasher::default();
    loop {
        let consumed = {
            let buf = reader
                .fill_buf()
                .map_err(|e| Errno::from_raw(e.raw_os_error().unwrap_or(nix::libc::EINVAL)))?;
            if buf.is_empty() {
                break;
            }
            hasher.write(buf);
            buf.len()
        };
        reader.consume(consumed);
    }
    Ok(rs_sha3_512::HasherContext::finish(&mut hasher)
        .as_ref()
        .to_vec())
}