use std::{collections::HashMap, path::PathBuf, str::FromStr};

use anyhow::{Context, Result, anyhow};
use gumdrop::Options;
use landlock::{
    ABI, Access, AccessFs, AccessNet, Compatible, NetPort, Ruleset, RulesetAttr,
    RulesetCreatedAttr, Scope, path_beneath_rules,
};

#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
enum Direction {
    In,
    Out,
}

#[derive(Debug, Options)]
struct Chas {
    #[options(no_multi, parse(try_from_str = "fs_parse"))]
    fs: HashMap<Permissions, Vec<PathBuf>>,
    #[options(no_multi, parse(try_from_str = "tcp_parse"))]
    tcp: HashMap<Direction, Vec<u16>>,
    #[options(no_multi, parse(try_from_str = "env_parse"))]
    env: HashMap<String, String>,
    clear_env: bool,
    #[options(free)]
    exec: Vec<String>,
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
enum Permissions {
    Read,
    Write,
}

fn env_parse(s: &str) -> Result<HashMap<String, String>> {
    let pairs = s.split(',');
    let mut envs = HashMap::new();
    for pair in pairs {
        let (k, v) = pair.split_once('=').ok_or(anyhow!("invalid env var"))?;
        envs.entry(k.to_string())
            .and_modify(|iv: &mut String| {
                *iv = v.to_string();
            })
            .or_insert(v.to_string());
    }
    Ok(envs)
}

// TODO handle ioctl perms
fn fs_parse(s: &str) -> Result<HashMap<Permissions, Vec<PathBuf>>> {
    let pairs = s.split(',');
    let mut rules = HashMap::new();
    for pair in pairs {
        let (s_perm, s_path) = pair
            .split_once('=')
            .ok_or(anyhow!("invalid filesystem pair"))?;
        if s_perm.len() > 1 {
            return Err(anyhow!("permission specifier too long"));
        }
        let perm = unsafe {
            match s_perm.get_unchecked(0..1) {
                "r" | "R" => Permissions::Read,
                "w" | "W" => Permissions::Write,
                _ => {
                    return Err(anyhow!("invalid permission specifier"));
                }
            }
        };
        for sub in s_path.split(':') {
            let path = PathBuf::from_str(sub).context("invalid path specifier")?;
            rules
                .entry(perm)
                .and_modify(|v: &mut Vec<PathBuf>| {
                    v.push(path.clone());
                })
                .or_insert(vec![path]);
        }
    }
    Ok(rules)
}

fn tcp_parse(s: &str) -> Result<HashMap<Direction, Vec<u16>>> {
    let pairs = s.split(',');
    let mut rules = HashMap::new();
    for pair in pairs {
        let (s_io, s_port) = pair.split_once('=').context("invalid tcp pair")?;
        let dir = match s_io {
            "i" => Direction::In,
            "o" => Direction::Out,
            _ => {
                return Err(anyhow!("invalid tcp specifier"));
            }
        };
        for sub in s_port.split(':') {
            let port = sub.parse().context("invalid port")?;
            rules
                .entry(dir)
                .and_modify(|v: &mut Vec<u16>| {
                    v.push(port);
                })
                .or_insert(vec![port]);
        }
    }
    Ok(rules)
}

fn main() -> Result<()> {
    let opts = Chas::parse_args_or_exit(gumdrop::ParsingStyle::StopAtFirstFree);
    println!("{opts:?}");
    let mut preempt = Ruleset::default();
    preempt = preempt.set_compatibility(landlock::CompatLevel::HardRequirement);
    preempt = preempt.scope(Scope::Signal).context("scoping signals")?;
    preempt = preempt
        .scope(Scope::AbstractUnixSocket)
        .context("scoping sockets")?;
    // let is_fs = !opts.fs.is_empty();
    let is_fs = true;
    let is_tcp = !opts.tcp.is_empty();
    preempt = if is_fs {
        preempt
            .handle_access(AccessFs::from_all(ABI::V6))
            .context("handling fs access")?
    } else {
        preempt
    };
    preempt = if is_tcp {
        preempt
            .handle_access(AccessNet::BindTcp)
            .context("handling tcp bind access")?
            .handle_access(AccessNet::ConnectTcp)
            .context("handling tcp conn access")?
    } else {
        preempt
    };
    let mut ruleset = preempt.create().context("creating ruleset")?;
    for (perms, paths) in opts.fs {
        let access = match perms {
            Permissions::Read => AccessFs::from_read(ABI::V6),
            Permissions::Write => AccessFs::from_write(ABI::V6),
        };
        ruleset = ruleset
            .add_rules(path_beneath_rules(paths, access))
            .context("adding fs rule")?;
    }
    for (dir, ports) in opts.tcp {
        let access = match dir {
            Direction::In => AccessNet::BindTcp,
            Direction::Out => AccessNet::ConnectTcp,
        };
        for port in ports {
            ruleset = ruleset
                .add_rule(NetPort::new(port, access))
                .context("adding tcp rule")?;
        }
    }
    ruleset.restrict_self().context("enforcing ruleset")?;
    if !opts.exec.is_empty() {
        let mut cmd = exec::Command::new(&opts.exec[0]);
        if opts.exec.len() > 1 {
            cmd.args(&opts.exec[1..]);
        }
        if opts.clear_env {
            for (k, _) in std::env::vars() {
                unsafe {
                    std::env::remove_var(k);
                }
            }
        }
        if !opts.env.is_empty() {
            for (k, v) in opts.env {
                unsafe {
                    std::env::set_var(k, v);
                }
            }
        }
        let err = cmd.exec();
        eprintln!("failed to run process: {}", err);
    } else {
        eprintln!(
            "
            yoke -- simple command sandboxer

            use: yoke [rules] [command]

            rules:
                --fs [access]=/path:/another/path,[access]=/more/path
                --tcp [access]=1234:5678,[access]=91011

            access specifiers (only one may be used per entry):
                fs:
                    r - read (implies execute)
                    w - write (implies read+execute)
                tcp:
                    i - in/bind
                    o - out/connect
            example:
                yoke --fs r=$(which ls):/etc ls /etc
                yoke --tcp i=80:443,o=80:443 --fs r=/srv/web myhttpserver
            "
        );
    }
    Ok(())
}