use crate::{ArgCmp, RawOsErrorNum, SecurityError, SecurityResult, Syscall, ViolationAction};
use seccompiler::{BpfProgram, SeccompAction, sock_filter};

const BPF_LD: u16 = 0x00;
const BPF_ALU: u16 = 0x04;
const BPF_JMP: u16 = 0x05;
const BPF_RET: u16 = 0x06;
const BPF_W: u16 = 0x00;
const BPF_ABS: u16 = 0x20;
const BPF_AND: u16 = 0x50;
const BPF_JA: u16 = 0x00;
const BPF_JEQ: u16 = 0x10;
const BPF_JGT: u16 = 0x20;
const BPF_JGE: u16 = 0x30;
const BPF_K: u16 = 0x00;
const BPF_MAX_LEN: usize = 4096;
const SECCOMP_DATA_NR_OFFSET: u32 = 0;
const SECCOMP_DATA_ARCH_OFFSET: u32 = 4;
const SECCOMP_DATA_ARGS_OFFSET: u32 = 16;
const SECCOMP_DATA_ARG_SIZE: u32 = 8;

const AUDIT_ARCH_X86_64: u32 = 62 | 0x8000_0000 | 0x4000_0000;
const AUDIT_ARCH_AARCH64: u32 = 183 | 0x8000_0000 | 0x4000_0000;
const AUDIT_ARCH_RISCV64: u32 = 243 | 0x8000_0000 | 0x4000_0000;

pub struct FilterRule {
    syscall: Syscall,
    args: Vec<ArgCmp>,
    action: RuleAction,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RuleAction {
    Allow,
    Log,
    Error(RawOsErrorNum),
}

impl FilterRule {
    pub fn new(syscall: Syscall, args: Vec<ArgCmp>, action: RuleAction) -> Self {
        Self {
            syscall,
            args,
            action,
        }
    }
}

impl RuleAction {
    fn rank(self) -> u8 {
        match self {
            Self::Error(_) => 0,
            Self::Log => 1,
            Self::Allow => 2,
        }
    }

    fn seccomp(self) -> SeccompAction {
        match self {
            Self::Allow => SeccompAction::Allow,
            Self::Log => SeccompAction::Log,
            Self::Error(errno) => SeccompAction::Errno(u32::from(errno)),
        }
    }
}

pub fn compile(
    mut rules: Vec<FilterRule>,
    violation_action: ViolationAction,
) -> SecurityResult<BpfProgram> {
    rules.sort_by_key(|rule| rule.action.rank());
    let mut program = arch_check()?;
    for rule in rules {
        program.extend(compile_rule(rule)?);
    }
    program.push(ret(violation(violation_action)));
    validate_len(&program)?;
    Ok(program)
}

fn compile_rule(rule: FilterRule) -> SecurityResult<BpfProgram> {
    validate_syscall_number(rule.syscall.number())?;
    let mut body = compile_rule_body(rule.args, rule.action)?;
    let mut program = syscall_gate(rule.syscall.number(), body.len())?;
    program.append(&mut body);
    Ok(program)
}

fn compile_rule_body(args: Vec<ArgCmp>, action: RuleAction) -> SecurityResult<BpfProgram> {
    if args.is_empty() {
        return Ok(vec![ret(action.seccomp())]);
    }
    let mut program = compile_conditions(&args)?;
    program.push(ret(action.seccomp()));
    Ok(program)
}

fn syscall_gate(number: i64, body_len: usize) -> SecurityResult<BpfProgram> {
    let skip = u32::try_from(body_len).map_err(|_| filter_too_large(body_len))?;
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, SECCOMP_DATA_NR_OFFSET),
        jump(BPF_JMP | BPF_JEQ | BPF_K, number as u32, 1, 0),
        stmt(BPF_JMP | BPF_JA, skip),
    ])
}

fn arch_check() -> SecurityResult<BpfProgram> {
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, SECCOMP_DATA_ARCH_OFFSET),
        jump(BPF_JMP | BPF_JEQ | BPF_K, audit_arch()?, 1, 0),
        ret(SeccompAction::KillProcess),
    ])
}

fn compile_conditions(args: &[ArgCmp]) -> SecurityResult<BpfProgram> {
    let mut remaining = condition_program_len(args)?;
    let mut program = Vec::new();
    for cmp in args {
        remaining -= condition_len(*cmp);
        program.extend(condition_bpf(*cmp, jump_offset(remaining + 1)?)?);
    }
    Ok(program)
}

fn condition_bpf(cmp: ArgCmp, offset: u8) -> SecurityResult<BpfProgram> {
    let arg_idx = arg_index(arg_cmp_index(cmp))?;
    match cmp {
        ArgCmp::NotEqual { value, .. } => ne_bpf(arg_idx, value.get(), offset),
        ArgCmp::LessThan { value, .. } => lt_bpf(arg_idx, value.get(), offset),
        ArgCmp::LessThanOrEqual { value, .. } => le_bpf(arg_idx, value.get(), offset),
        ArgCmp::Equal { value, .. } => eq_bpf(arg_idx, value.get(), offset),
        ArgCmp::GreaterThanOrEqual { value, .. } => ge_bpf(arg_idx, value.get(), offset),
        ArgCmp::GreaterThan { value, .. } => gt_bpf(arg_idx, value.get(), offset),
        ArgCmp::EqualMasked { mask, value, .. } => {
            masked_eq_bpf(arg_idx, mask.get(), value.get(), offset)
        }
    }
}

fn condition_program_len(args: &[ArgCmp]) -> SecurityResult<usize> {
    let len = args.iter().map(|cmp| condition_len(*cmp)).sum();
    jump_offset(len + 1)?;
    Ok(len)
}

fn condition_len(cmp: ArgCmp) -> usize {
    match cmp {
        ArgCmp::EqualMasked { .. } => 6,
        ArgCmp::LessThan { .. }
        | ArgCmp::LessThanOrEqual { .. }
        | ArgCmp::GreaterThanOrEqual { .. }
        | ArgCmp::GreaterThan { .. } => 5,
        ArgCmp::NotEqual { .. } | ArgCmp::Equal { .. } => 4,
    }
}

fn arg_cmp_index(cmp: ArgCmp) -> u32 {
    match cmp {
        ArgCmp::NotEqual { arg_idx, .. }
        | ArgCmp::LessThan { arg_idx, .. }
        | ArgCmp::LessThanOrEqual { arg_idx, .. }
        | ArgCmp::Equal { arg_idx, .. }
        | ArgCmp::GreaterThanOrEqual { arg_idx, .. }
        | ArgCmp::GreaterThan { arg_idx, .. }
        | ArgCmp::EqualMasked { arg_idx, .. } => arg_idx,
    }
}

fn arg_index(value: u32) -> SecurityResult<u8> {
    let index = u8::try_from(value).map_err(|_| invalid_arg_index(value))?;
    if index > 5 {
        return Err(invalid_arg_index(value));
    }
    Ok(index)
}

fn eq_bpf(arg_idx: u8, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb, lsb) = split_value(value);
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        jump(
            BPF_JMP | BPF_JEQ | BPF_K,
            msb,
            0,
            checked_offset(offset, 2)?,
        ),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        jump(BPF_JMP | BPF_JEQ | BPF_K, lsb, 0, offset),
    ])
}

fn ne_bpf(arg_idx: u8, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb, lsb) = split_value(value);
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        jump(BPF_JMP | BPF_JEQ | BPF_K, msb, 0, 2),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        jump(BPF_JMP | BPF_JEQ | BPF_K, lsb, offset, 0),
    ])
}

fn ge_bpf(arg_idx: u8, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb, lsb) = split_value(value);
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        jump(BPF_JMP | BPF_JGT | BPF_K, msb, 3, 0),
        jump(
            BPF_JMP | BPF_JEQ | BPF_K,
            msb,
            0,
            checked_offset(offset, 2)?,
        ),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        jump(BPF_JMP | BPF_JGE | BPF_K, lsb, 0, offset),
    ])
}

fn gt_bpf(arg_idx: u8, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb, lsb) = split_value(value);
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        jump(BPF_JMP | BPF_JGT | BPF_K, msb, 3, 0),
        jump(
            BPF_JMP | BPF_JEQ | BPF_K,
            msb,
            0,
            checked_offset(offset, 2)?,
        ),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        jump(BPF_JMP | BPF_JGT | BPF_K, lsb, 0, offset),
    ])
}

fn le_bpf(arg_idx: u8, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb, lsb) = split_value(value);
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        jump(
            BPF_JMP | BPF_JGT | BPF_K,
            msb,
            checked_offset(offset, 3)?,
            0,
        ),
        jump(BPF_JMP | BPF_JEQ | BPF_K, msb, 0, 2),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        jump(BPF_JMP | BPF_JGT | BPF_K, lsb, offset, 0),
    ])
}

fn lt_bpf(arg_idx: u8, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb, lsb) = split_value(value);
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        jump(
            BPF_JMP | BPF_JGT | BPF_K,
            msb,
            checked_offset(offset, 3)?,
            0,
        ),
        jump(BPF_JMP | BPF_JEQ | BPF_K, msb, 0, 2),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        jump(BPF_JMP | BPF_JGE | BPF_K, lsb, offset, 0),
    ])
}

fn masked_eq_bpf(arg_idx: u8, mask: u64, value: u64, offset: u8) -> SecurityResult<BpfProgram> {
    let (msb_offset, lsb_offset) = data_offsets(arg_idx);
    let (msb, lsb) = split_value(value & mask);
    let (mask_msb, mask_lsb) = split_value(mask);
    Ok(vec![
        stmt(BPF_LD | BPF_W | BPF_ABS, msb_offset),
        stmt(BPF_ALU | BPF_AND | BPF_K, mask_msb),
        jump(
            BPF_JMP | BPF_JEQ | BPF_K,
            msb,
            0,
            checked_offset(offset, 3)?,
        ),
        stmt(BPF_LD | BPF_W | BPF_ABS, lsb_offset),
        stmt(BPF_ALU | BPF_AND | BPF_K, mask_lsb),
        jump(BPF_JMP | BPF_JEQ | BPF_K, lsb, 0, offset),
    ])
}

fn data_offsets(arg_idx: u8) -> (u32, u32) {
    let offset = SECCOMP_DATA_ARGS_OFFSET + u32::from(arg_idx) * SECCOMP_DATA_ARG_SIZE;
    (offset + SECCOMP_DATA_ARG_SIZE / 2, offset)
}

fn split_value(value: u64) -> (u32, u32) {
    ((value >> 32) as u32, value as u32)
}

fn checked_offset(offset: u8, extra: u8) -> SecurityResult<u8> {
    offset
        .checked_add(extra)
        .ok_or_else(|| filter_too_large(usize::from(offset) + usize::from(extra)))
}

fn jump_offset(value: usize) -> SecurityResult<u8> {
    u8::try_from(value).map_err(|_| filter_too_large(value))
}

fn validate_syscall_number(number: i64) -> SecurityResult<()> {
    if (0..=i64::from(u32::MAX)).contains(&number) {
        return Ok(());
    }
    Err(SecurityError::new(
        format!("invalid Linux syscall number {number}"),
        "use a named Syscall constant or a raw number between 0 and u32::MAX",
    ))
}

fn validate_len(program: &BpfProgram) -> SecurityResult<()> {
    if program.len() < BPF_MAX_LEN {
        return Ok(());
    }
    Err(filter_too_large(program.len()))
}

fn audit_arch() -> SecurityResult<u32> {
    match std::env::consts::ARCH {
        "x86_64" => Ok(AUDIT_ARCH_X86_64),
        "aarch64" => Ok(AUDIT_ARCH_AARCH64),
        "riscv64" => Ok(AUDIT_ARCH_RISCV64),
        arch => Err(SecurityError::new(
            format!("unsupported seccomp target architecture: {arch}"),
            "run the security feature on x86_64, aarch64, or riscv64 Linux",
        )),
    }
}

fn violation(action: ViolationAction) -> SeccompAction {
    match action {
        ViolationAction::KillProcess => SeccompAction::KillProcess,
        ViolationAction::AllowAndLog => SeccompAction::Log,
    }
}

fn invalid_arg_index(value: u32) -> SecurityError {
    SecurityError::new(
        format!("invalid syscall argument index {value}"),
        "use arg_idx between 0 and 5",
    )
}

fn filter_too_large(len: usize) -> SecurityError {
    SecurityError::new(
        format!("seccomp filter has {len} BPF instructions"),
        "reduce the allow list or split the service boundary",
    )
}

fn ret(action: SeccompAction) -> sock_filter {
    stmt(BPF_RET | BPF_K, action.into())
}

fn stmt(code: u16, k: u32) -> sock_filter {
    sock_filter {
        code,
        jt: 0,
        jf: 0,
        k,
    }
}

fn jump(code: u16, k: u32, jt: u8, jf: u8) -> sock_filter {
    sock_filter { code, jt, jf, k }
}

#[cfg(test)]
mod tests;
