use crate::{Rule, SecurityError, SecurityResult, ViolationAction};

#[cfg(target_os = "linux")]
use crate::bpf::{FilterRule, RuleAction};
#[cfg(target_os = "linux")]
use crate::{ArgCmp, RawOsErrorNum, Syscall};
#[cfg(target_os = "linux")]
use seccompiler::BpfProgram;
#[cfg(target_os = "linux")]
use std::collections::BTreeMap;
#[cfg(target_os = "linux")]
use std::io;

#[cfg(target_os = "linux")]
#[repr(C)]
struct SockFilterProgram {
    len: libc::c_ushort,
    filter: *const seccompiler::sock_filter,
}

#[cfg(target_os = "linux")]
type RuleKey = (i64, Vec<ArgCmp>);

#[cfg(target_os = "linux")]
#[derive(Default)]
struct Policy {
    rules: Vec<FilterRule>,
    actions: BTreeMap<RuleKey, RuleAction>,
}

#[cfg(target_os = "linux")]
struct CompiledPolicy {
    program: BpfProgram,
}

#[cfg(target_os = "linux")]
impl Policy {
    fn from_rules(rules: &[Rule]) -> SecurityResult<Self> {
        let mut policy = Self::default();
        for rule in rules {
            policy.insert(rule)?;
        }
        Ok(policy)
    }

    fn insert(&mut self, rule: &Rule) -> SecurityResult<()> {
        match rule {
            Rule::Allow(syscall, args) => self.insert_allowed(*syscall, args),
            Rule::AllowAndLog(syscall, args) => self.insert_logged(*syscall, args),
            Rule::ReturnError(syscall, errno, args) => self.insert_error(*syscall, *errno, args),
        }
    }

    fn insert_allowed(&mut self, syscall: Syscall, args: &[ArgCmp]) -> SecurityResult<()> {
        self.note_action(syscall, args, RuleAction::Allow)?;
        self.rules
            .push(FilterRule::new(syscall, args.to_vec(), RuleAction::Allow));
        Ok(())
    }

    fn insert_logged(&mut self, syscall: Syscall, args: &[ArgCmp]) -> SecurityResult<()> {
        self.note_action(syscall, args, RuleAction::Log)?;
        self.rules
            .push(FilterRule::new(syscall, args.to_vec(), RuleAction::Log));
        Ok(())
    }

    fn insert_error(
        &mut self,
        syscall: Syscall,
        errno: RawOsErrorNum,
        args: &[ArgCmp],
    ) -> SecurityResult<()> {
        self.note_action(syscall, args, RuleAction::Error(errno))?;
        self.rules.push(FilterRule::new(
            syscall,
            args.to_vec(),
            RuleAction::Error(errno),
        ));
        Ok(())
    }

    fn compile(self, violation_action: ViolationAction) -> SecurityResult<CompiledPolicy> {
        Ok(CompiledPolicy {
            program: crate::bpf::compile(self.rules, violation_action)?,
        })
    }

    fn note_action(
        &mut self,
        syscall: Syscall,
        args: &[ArgCmp],
        action: RuleAction,
    ) -> SecurityResult<()> {
        let key = action_key(syscall, args);
        let Some(previous) = self.actions.insert(key, action) else {
            return Ok(());
        };
        if previous == action {
            return Ok(());
        }
        Err(SecurityError::new(
            format!(
                "conflicting seccomp actions for syscall {}",
                syscall.number()
            ),
            "keep one action for each syscall and argument comparison set",
        ))
    }
}

#[cfg(target_os = "linux")]
impl CompiledPolicy {
    fn apply(self) -> SecurityResult<()> {
        set_no_new_privs()?;
        apply_filter(&self.program)?;
        Ok(())
    }
}

#[cfg(target_os = "linux")]
/// Install a seccomp filter for the current thread
pub fn enable_syscall_sandboxing(
    violation_action: ViolationAction,
    exception_rules: &[Rule],
) -> SecurityResult<()> {
    Policy::from_rules(exception_rules)?
        .compile(violation_action)?
        .apply()
}

#[cfg(not(target_os = "linux"))]
/// Return an unsupported-target error outside Linux
pub fn enable_syscall_sandboxing(
    violation_action: ViolationAction,
    exception_rules: &[Rule],
) -> SecurityResult<()> {
    let _ = (violation_action, exception_rules);
    Err(SecurityError::unsupported_target())
}

#[cfg(target_os = "linux")]
/// Check whether the current thread already has seccomp enabled
pub fn is_syscall_sandboxing_enabled_for_current_thread() -> SecurityResult<bool> {
    // SAFETY: PR_GET_SECCOMP reads kernel state and passes no pointers
    let rc = unsafe { libc::prctl(libc::PR_GET_SECCOMP) };
    if rc < 0 {
        return Err(prctl_error("read the current seccomp mode"));
    }
    match rc {
        0 => Ok(false),
        mode if mode == libc::SECCOMP_MODE_FILTER as i32 => Ok(true),
        mode => Err(SecurityError::new(
            format!("unsupported seccomp mode {mode}"),
            "run with no seccomp filter or Linux SECCOMP_MODE_FILTER",
        )),
    }
}

#[cfg(not(target_os = "linux"))]
/// Return an unsupported-target error outside Linux
pub fn is_syscall_sandboxing_enabled_for_current_thread() -> SecurityResult<bool> {
    Err(SecurityError::unsupported_target())
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
/// Block direct x86_64 time-stamp-counter access in the process
pub fn forbid_x86_64_cpu_cycle_counter() -> SecurityResult<()> {
    // SAFETY: PR_SET_TSC passes plain integer flags and touches no memory
    let rc = unsafe { libc::prctl(libc::PR_SET_TSC, libc::PR_TSC_SIGSEGV, 0, 0, 0) };
    if rc != 0 {
        return Err(prctl_error("disable the x86_64 cycle counter"));
    }
    Ok(())
}

#[cfg(all(target_os = "linux", not(target_arch = "x86_64")))]
/// No-op on Linux targets without an x86_64 time-stamp counter
pub fn forbid_x86_64_cpu_cycle_counter() -> SecurityResult<()> {
    Ok(())
}

#[cfg(not(target_os = "linux"))]
/// Return an unsupported-target error outside Linux
pub fn forbid_x86_64_cpu_cycle_counter() -> SecurityResult<()> {
    Err(SecurityError::unsupported_target())
}

#[cfg(target_os = "linux")]
fn action_key(syscall: Syscall, args: &[ArgCmp]) -> RuleKey {
    let mut args = args.to_vec();
    args.sort();
    (syscall.number(), args)
}

#[cfg(target_os = "linux")]
fn apply_filter(program: &BpfProgram) -> SecurityResult<()> {
    let len = filter_len(program)?;
    let filter = SockFilterProgram {
        len,
        filter: program.as_ptr(),
    };
    // SAFETY: filter points to a valid SockFilterProgram that outlives the call
    let rc = unsafe { libc::syscall(libc::SYS_seccomp, libc::SECCOMP_SET_MODE_FILTER, 0, &filter) };
    if rc == 0 {
        return Ok(());
    }
    Err(seccomp_error())
}

#[cfg(target_os = "linux")]
fn set_no_new_privs() -> SecurityResult<()> {
    // SAFETY: PR_SET_NO_NEW_PRIVS passes plain integers and touches no memory
    let rc = unsafe { libc::prctl(libc::PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) };
    if rc == 0 {
        return Ok(());
    }
    Err(prctl_error("set no_new_privs before installing seccomp"))
}

#[cfg(target_os = "linux")]
fn filter_len(program: &BpfProgram) -> SecurityResult<libc::c_ushort> {
    libc::c_ushort::try_from(program.len()).map_err(|_| {
        SecurityError::new(
            format!("seccomp filter has {} BPF instructions", program.len()),
            "reduce the allow list or split the service boundary",
        )
    })
}

#[cfg(target_os = "linux")]
fn seccomp_error() -> SecurityError {
    SecurityError::new(
        format!(
            "failed to install seccomp filter: {}",
            io::Error::last_os_error()
        ),
        "apply the sandbox after startup opens files and sockets but before spawning workers",
    )
}

#[cfg(target_os = "linux")]
fn prctl_error(action: &'static str) -> SecurityError {
    SecurityError::new(
        format!("failed to {action}: {}", io::Error::last_os_error()),
        "run on Linux with seccomp and prctl support",
    )
}
