use std::io;

use rustix::thread::{
    CapabilitySet, CapabilitySets, capabilities, capability_is_in_ambient_set,
    capability_is_in_bounding_set, configure_capability_in_ambient_set,
    remove_capability_from_bounding_set, set_capabilities,
};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct CapabilityState {
    pub sets: CapabilitySets,
    pub ambient: bool,
    pub bounding: bool,
}

/// # Errors
/// Returns an error when any kernel capability set cannot be inspected.
pub fn inspect(capability: CapabilitySet) -> io::Result<CapabilityState> {
    Ok(CapabilityState {
        sets: capabilities(None)?,
        ambient: capability_is_in_ambient_set(capability)?,
        bounding: capability_is_in_bounding_set(capability)?,
    })
}

/// # Errors
/// Refuses any incomplete set, bounding, ambient, or secure-bit transition.
pub fn irreversibly_drop(capability: CapabilitySet) -> io::Result<()> {
    require_single_threaded_process()?;
    let before = inspect(capability)?;
    if before.bounding && !before.sets.effective.contains(CapabilitySet::SETPCAP) {
        return Err(io::Error::from_raw_os_error(1));
    }
    if before.ambient {
        configure_capability_in_ambient_set(capability, false)?;
    }
    if before.bounding {
        remove_capability_from_bounding_set(capability)?;
    }
    let mut sets = before.sets;
    sets.effective.remove(capability);
    sets.permitted.remove(capability);
    sets.inheritable.remove(capability);
    if sets != before.sets {
        set_capabilities(None, sets)?;
    }
    let observed = inspect(capability)?;
    if observed.sets.effective.intersects(capability)
        || observed.sets.permitted.intersects(capability)
        || observed.sets.inheritable.intersects(capability)
        || observed.ambient
        || observed.bounding
    {
        return Err(io::Error::other(
            "capability remained available after irreversible drop",
        ));
    }
    Ok(())
}

/// # Errors
/// Refuses if an absent capability can be restored or any verification fails.
pub fn require_regain_refused(capability: CapabilitySet) -> io::Result<()> {
    let before = inspect(capability)?;
    if before.sets.effective.intersects(capability)
        || before.sets.permitted.intersects(capability)
        || before.sets.inheritable.intersects(capability)
        || before.ambient
        || before.bounding
    {
        return Err(io::Error::other(
            "capability remains available before regain verification",
        ));
    }
    let mut attempted = before.sets;
    attempted.effective.insert(capability);
    attempted.permitted.insert(capability);
    attempted.inheritable.insert(capability);
    if set_capabilities(None, attempted).is_ok()
        || configure_capability_in_ambient_set(capability, true).is_ok()
        || inspect(capability)? != before
    {
        return Err(io::Error::other("capability regain was not refused"));
    }
    Ok(())
}

fn require_single_threaded_process() -> io::Result<()> {
    let mut tasks = std::fs::read_dir("/proc/self/task")?;
    if tasks.next().transpose()?.is_some() && tasks.next().transpose()?.is_none() {
        Ok(())
    } else {
        Err(io::Error::other(
            "capability transition requires a single-threaded process",
        ))
    }
}

#[cfg(test)]
mod tests {
    use rustix::thread::CapabilitySet;

    use super::{inspect, irreversibly_drop, require_regain_refused};

    #[test]
    fn inspects_current_capability_sets() {
        let state = inspect(CapabilitySet::NET_ADMIN).unwrap();
        assert_eq!(
            state.sets.effective.contains(CapabilitySet::NET_ADMIN),
            state.sets.effective.intersects(CapabilitySet::NET_ADMIN)
        );
    }

    #[test]
    fn multithreaded_process_refuses_before_transition() {
        let before = inspect(CapabilitySet::NET_ADMIN).unwrap();
        let waiting = std::sync::Arc::new(std::sync::Barrier::new(2));
        let child_waiting = std::sync::Arc::clone(&waiting);
        let child = std::thread::spawn(move || child_waiting.wait());
        assert!(irreversibly_drop(CapabilitySet::NET_ADMIN).is_err());
        assert_eq!(inspect(CapabilitySet::NET_ADMIN).unwrap(), before);
        waiting.wait();
        child.join().unwrap();
    }

    #[test]
    fn absent_capability_cannot_be_regained() {
        let state = inspect(CapabilitySet::NET_ADMIN).unwrap();
        if !state.sets.effective.intersects(CapabilitySet::NET_ADMIN)
            && !state.sets.permitted.intersects(CapabilitySet::NET_ADMIN)
            && !state.sets.inheritable.intersects(CapabilitySet::NET_ADMIN)
            && !state.ambient
            && !state.bounding
        {
            require_regain_refused(CapabilitySet::NET_ADMIN).unwrap();
        }
    }
}