//
// Syd: rock-solid unikernel
// src/syd-stat.rs: Print process status of the given PID or the current process.
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

#![recursion_limit = "256"]

use std::process::ExitCode;

use nix::{libc::pid_t, sys::signal::Signal, unistd::Pid};
use procfs::process::Process;
use serde::{ser::SerializeSeq, Serialize, Serializer};
use serde_json::json;
use syd::{
    caps::Capability::{self, *},
    err::SydResult,
};

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

struct SigMap(u64);

impl SigMap {
    /// Converts the bitmap to a vector of signal names.
    fn names(&self) -> Vec<&'static str> {
        (0..32) // Focus on the defined signals, ignoring real-time signals.
            .filter_map(|bit| {
                if (self.0 >> bit) & 1 == 1 {
                    // Attempt to cast bit position to Signal enum
                    Signal::try_from(bit + 1).ok()
                } else {
                    None
                }
            })
            .map(|signal| signal.as_str())
            .collect()
    }
}

impl Serialize for SigMap {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let signal_names = self.names();
        let mut seq = serializer.serialize_seq(Some(signal_names.len()))?;
        for signal_name in signal_names {
            seq.serialize_element(signal_name)?;
        }
        seq.end()
    }
}

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

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

    let pid = match std::env::args().nth(1).map(|arg| arg.parse::<pid_t>()) {
        Some(Ok(pid)) => pid,
        None => Pid::this().as_raw(),
        Some(Err(_)) => {
            help();
            return Ok(ExitCode::SUCCESS);
        }
    };

    let proc = match Process::new(pid) {
        Ok(proc) => proc,
        Err(error) => {
            eprintln!("syd-stat: {error}");
            return Ok(ExitCode::FAILURE);
        }
    };

    let stat = match proc.status() {
        Ok(stat) => stat,
        Err(error) => {
            eprintln!("syd-stat: {error}");
            return Ok(ExitCode::FAILURE);
        }
    };

    let seccomp = match stat.seccomp {
        Some(0) => "disabled",
        Some(1) => "strict",
        Some(2) => "filter",
        _ => "unknown",
    };

    #[allow(clippy::disallowed_methods)]
    let status = json!({
        "pid": pid,
        "ppid": stat.ppid,
        "tgid": stat.tgid,
        "ngid": stat.ngid.unwrap_or(0),
        "tracerpid": stat.tracerpid,
        "threads": stat.threads,
        "name": stat.name,
        "state": stat.state,
        "umask": stat.umask.unwrap_or(0),
        "ruid": stat.ruid,
        "euid": stat.euid,
        "suid": stat.suid,
        "fuid": stat.fuid,
        "rgid": stat.rgid,
        "egid": stat.egid,
        "sgid": stat.sgid,
        "fgid": stat.fgid,
        "fdsize": stat.fdsize,
        "vmpeak": stat.vmpeak.unwrap_or(0),
        "vmsize": stat.vmsize.unwrap_or(0),
        "vmlck": stat.vmlck.unwrap_or(0),
        "vmpin": stat.vmpin.unwrap_or(0),
        "vmhwm": stat.vmhwm.unwrap_or(0),
        "vmrss": stat.vmrss.unwrap_or(0),
        "rssanon": stat.rssanon.unwrap_or(0),
        "rssfile": stat.rssfile.unwrap_or(0),
        "rssshmem": stat.rssshmem.unwrap_or(0),
        "vmdata": stat.vmdata.unwrap_or(0),
        "vmstk": stat.vmstk.unwrap_or(0),
        "vmexe": stat.vmexe.unwrap_or(0),
        "vmlib": stat.vmlib.unwrap_or(0),
        "vmpte": stat.vmpte.unwrap_or(0),
        "vmswap": stat.vmswap.unwrap_or(0),
        "hugetlbpages": stat.hugetlbpages.unwrap_or(0),
        "nonewprivs": stat.nonewprivs.map_or(false, |val| val != 0),
        "seccomp": seccomp,
        "speculation_store_bypass": stat.speculation_store_bypass,
        "caps": {
            "ambient": mask2cap(stat.capamb.unwrap_or(0)),
            "bounding": mask2cap(stat.capbnd.unwrap_or(0)),
            "effective": mask2cap(stat.capeff),
            "inheritable": mask2cap(stat.capinh),
            "permitted": mask2cap(stat.capprm),
        },
        "sig": {
            "blk": SigMap(stat.sigblk),
            "cgt": SigMap(stat.sigcgt),
            "ign": SigMap(stat.sigign),
            "pnd": SigMap(stat.sigpnd),
        }
    });

    #[allow(clippy::disallowed_methods)]
    let status = serde_json::to_string_pretty(&status).unwrap();
    println!("{status}");

    Ok(ExitCode::SUCCESS)
}

fn help() {
    println!("Usage: syd-stat [PID]");
    println!("Print detailed information about the given process process or current process.");
}

fn mask2cap(bitmask: u64) -> Vec<String> {
    CAPS.iter()
        .filter(|&&cap| {
            let cap_bit = cap.index() as u64;
            bitmask & (1 << cap_bit) != 0
        })
        .map(|cap| cap.to_string())
        .collect()
}

const CAPS: &[Capability] = &[
    CAP_CHOWN,
    CAP_DAC_OVERRIDE,
    CAP_DAC_READ_SEARCH,
    CAP_FOWNER,
    CAP_FSETID,
    CAP_KILL,
    CAP_SETGID,
    CAP_SETUID,
    CAP_SETPCAP,
    CAP_LINUX_IMMUTABLE,
    CAP_NET_BIND_SERVICE,
    CAP_NET_BROADCAST,
    CAP_NET_ADMIN,
    CAP_NET_RAW,
    CAP_IPC_LOCK,
    CAP_IPC_OWNER,
    CAP_SYS_MODULE,
    CAP_SYS_RAWIO,
    CAP_SYS_CHROOT,
    CAP_SYS_PTRACE,
    CAP_SYS_PACCT,
    CAP_SYS_ADMIN,
    CAP_SYS_BOOT,
    CAP_SYS_NICE,
    CAP_SYS_RESOURCE,
    CAP_SYS_TIME,
    CAP_SYS_TTY_CONFIG,
    CAP_MKNOD,
    CAP_LEASE,
    CAP_AUDIT_WRITE,
    CAP_AUDIT_CONTROL,
    CAP_SETFCAP,
    CAP_MAC_OVERRIDE,
    CAP_MAC_ADMIN,
    CAP_SYSLOG,
    CAP_WAKE_ALARM,
    CAP_BLOCK_SUSPEND,
    CAP_AUDIT_READ,
    CAP_PERFMON,
    CAP_BPF,
    CAP_CHECKPOINT_RESTORE,
];
