//! Deterministic integrity checks for programs executed directly from EROFS.
use crate::{MAX_ENTRIES, MAX_UNPACKED, Software, relative};
use fds_common::{Error, Result};
use sha2::{Digest, Sha256};
use std::{
    fs::{self, File, OpenOptions},
    io::Read,
    os::unix::fs::{OpenOptionsExt, PermissionsExt, symlink},
    path::{Component, Path, PathBuf},
};

#[derive(Debug, PartialEq, Eq)]
pub struct Inventory {
    pub bytes: u64,
    pub entries: u32,
    pub sha256: String,
}

fn paths(root: &Path, directory: &Path, output: &mut Vec<PathBuf>) -> Result<()> {
    for entry in fs::read_dir(directory)? {
        let entry = entry?;
        let path = entry.path();
        let name = path.strip_prefix(root).unwrap();
        if !name.to_str().is_some_and(relative) || output.len() >= MAX_ENTRIES as usize {
            return Err(Error(
                "Invalid software path or too many installed files".into(),
            ));
        }
        output.push(name.to_owned());
        if entry.file_type()?.is_dir() {
            paths(root, &path, output)?;
        }
    }
    Ok(())
}

fn link_inside(root: &Path, path: &Path, target: &Path) -> Result<()> {
    if target.is_absolute() || target.as_os_str().is_empty() {
        return Err(Error(
            "Installed software links must be relative to their tree".into(),
        ));
    }
    let mut depth = path
        .parent()
        .unwrap()
        .strip_prefix(root)
        .unwrap()
        .components()
        .count();
    for part in target.components() {
        match part {
            Component::Normal(_) => depth += 1,
            Component::CurDir => (),
            Component::ParentDir if depth > 0 => depth -= 1,
            _ => return Err(Error("Installed software symlink escapes its tree".into())),
        }
    }
    match path.canonicalize() {
        Ok(destination) if !destination.starts_with(root) => {
            return Err(Error(
                "Installed software symlink resolves outside its tree".into(),
            ));
        }
        Ok(_) => (),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => (),
        Err(e) => return Err(e.into()),
    }
    Ok(())
}

pub fn inspect(root: &Path, architecture: &str) -> Result<Inventory> {
    if !fs::symlink_metadata(root)?.is_dir() {
        return Err(Error(
            "Installed software root must be a real directory".into(),
        ));
    }
    let root = root.canonicalize()?;
    let mut entries = Vec::new();
    paths(&root, &root, &mut entries)?;
    entries.sort();
    let mut hash = Sha256::new();
    let mut bytes = 0u64;
    for relative in &entries {
        let path = root.join(relative);
        let metadata = path.symlink_metadata()?;
        let name = relative.as_os_str().as_encoded_bytes();
        hash.update((name.len() as u64).to_le_bytes());
        hash.update(name);
        let mode = metadata.permissions().mode();
        if !metadata.is_symlink() && mode & 0o7022 != 0 {
            return Err(Error(format!(
                "Privileged or group/world-writable software file: {}",
                relative.display()
            )));
        }
        hash.update((mode & 0o777).to_le_bytes());
        if metadata.is_dir() {
            hash.update(b"d");
        } else if metadata.is_symlink() {
            hash.update(b"l");
            let target = fs::read_link(&path)?;
            link_inside(&root, &path, &target)?;
            let value = target.as_os_str().as_encoded_bytes();
            hash.update((value.len() as u64).to_le_bytes());
            hash.update(value);
        } else if metadata.is_file() {
            hash.update(b"f");
            bytes = bytes
                .checked_add(metadata.len())
                .ok_or_else(|| Error("Software size overflow".into()))?;
            if bytes > MAX_UNPACKED {
                return Err(Error("Installed software tree exceeds 1 GiB".into()));
            }
            hash.update(metadata.len().to_le_bytes());
            let mut file = OpenOptions::new()
                .read(true)
                .custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK)
                .open(&path)?;
            let mut prefix = [0u8; 20];
            let n = file.read(&mut prefix)?;
            if prefix.starts_with(b"\x7fELF")
                && (architecture != "aarch64"
                    || n < 20
                    || prefix[4] != 2
                    || prefix[5] != 1
                    || prefix[18..20] != [183, 0])
            {
                return Err(Error(
                    "Installed software contains an ELF file for the wrong architecture".into(),
                ));
            }
            hash.update(&prefix[..n]);
            let copied = std::io::copy(&mut file, &mut HashWriter(&mut hash))?;
            if copied + n as u64 != metadata.len() {
                return Err(Error("Installed software changed during inspection".into()));
            }
        } else {
            return Err(Error(
                "Installed software permits only files, directories and internal symlinks".into(),
            ));
        }
    }
    Ok(Inventory {
        bytes,
        entries: entries.len() as u32,
        sha256: format!("{:x}", hash.finalize()),
    })
}

struct HashWriter<'a>(&'a mut Sha256);
impl std::io::Write for HashWriter<'_> {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        self.0.update(bytes);
        Ok(bytes.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

pub fn verify(root: &Path, software: &Software) -> Result<()> {
    software.validate()?;
    if !software.installed {
        return Err(Error("Expected an installed software tree".into()));
    }
    let actual = inspect(root, &software.architecture)?;
    if actual.sha256 != software.sha256
        || actual.bytes != software.unpacked_bytes
        || actual.entries != software.entries
    {
        return Err(Error(
            "Installed software tree digest, size or entry count disagrees with catalogue".into(),
        ));
    }
    let canonical = root.canonicalize()?;
    for command in software.commands.values() {
        let executable = root.join(command).canonicalize()?;
        if !executable.starts_with(&canonical)
            || !executable.is_file()
            || executable.metadata()?.permissions().mode() & 0o111 == 0
        {
            return Err(Error(
                "Software command must resolve to an executable within its installed tree".into(),
            ));
        }
    }
    Ok(())
}

/// XBPS trees contain root-relative symlinks. Relocate those on the workstation
/// so the exact tree can be mounted under /run without pointing into SYSTEM.
pub fn relocate(root: &Path) -> Result<()> {
    let mut entries = Vec::new();
    paths(root, root, &mut entries)?;
    for name in entries {
        let path = root.join(&name);
        if path.is_symlink() {
            let target = fs::read_link(&path)?;
            if target.is_absolute() {
                let mut relative = PathBuf::new();
                for _ in name.parent().unwrap().components() {
                    relative.push("..");
                }
                relative.push(target.strip_prefix("/").unwrap());
                fs::remove_file(&path)?;
                symlink(relative, &path)?;
            }
        }
    }
    Ok(())
}

pub fn copy(root: &Path, destination: &Path) -> Result<()> {
    fs::create_dir(destination)?;
    fs::set_permissions(destination, fs::Permissions::from_mode(0o755))?;
    let mut entries = Vec::new();
    paths(root, root, &mut entries)?;
    entries.sort();
    for name in entries {
        let from = root.join(&name);
        let to = destination.join(name);
        let metadata = from.symlink_metadata()?;
        if metadata.is_symlink() {
            symlink(fs::read_link(from)?, to)?;
        } else if metadata.is_dir() {
            fs::create_dir(&to)?;
            fs::set_permissions(to, metadata.permissions())?;
        } else if metadata.is_file() {
            fs::copy(from, to)?;
        } else {
            return Err(Error("Unsupported installed software file type".into()));
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};
    struct Fixture(PathBuf);
    impl Fixture {
        fn new() -> Self {
            static NEXT: AtomicU64 = AtomicU64::new(0);
            let root = std::env::temp_dir().join(format!(
                "fds-tree-{}-{}",
                std::process::id(),
                NEXT.fetch_add(1, Ordering::Relaxed)
            ));
            fs::create_dir(&root).unwrap();
            fs::create_dir_all(root.join("usr/bin")).unwrap();
            fs::write(root.join("usr/bin/tool"), b"#!/bin/sh\necho installed\n").unwrap();
            fs::set_permissions(root.join("usr/bin/tool"), fs::Permissions::from_mode(0o755))
                .unwrap();
            symlink("usr/bin", root.join("bin")).unwrap();
            Self(root)
        }
        fn software(&self) -> Software {
            let found = inspect(&self.0, "aarch64").unwrap();
            Software {
                id: "demo.tool".into(),
                name: "Tool".into(),
                version: "1".into(),
                architecture: "aarch64".into(),
                partition: 2,
                installed: true,
                packages: vec!["WindowMaker-0.96.0_1".into()],
                archive_bytes: 0,
                unpacked_bytes: found.bytes,
                entries: found.entries,
                sha256: found.sha256,
                commands: [("tool".into(), "bin/tool".into())].into(),
            }
        }
    }
    impl Drop for Fixture {
        fn drop(&mut self) {
            let _ = fs::remove_dir_all(&self.0);
        }
    }
    #[test]
    fn installed_tree_roundtrip_and_tampering() {
        let root = Fixture::new();
        let software = root.software();
        verify(&root.0, &software).unwrap();
        let copy_path = root.0.with_extension("copy");
        copy(&root.0, &copy_path).unwrap();
        verify(&copy_path, &software).unwrap();
        fs::remove_dir_all(copy_path).unwrap();
        fs::write(root.0.join("usr/bin/tool"), b"changed").unwrap();
        assert!(verify(&root.0, &software).is_err());
    }
    #[test]
    fn escaping_links_and_privileged_files_are_rejected() {
        let root = Fixture::new();
        symlink("../../../etc/passwd", root.0.join("usr/bin/escape")).unwrap();
        assert!(inspect(&root.0, "aarch64").is_err());
        fs::remove_file(root.0.join("usr/bin/escape")).unwrap();
        fs::set_permissions(
            root.0.join("usr/bin/tool"),
            fs::Permissions::from_mode(0o4755),
        )
        .unwrap();
        assert!(inspect(&root.0, "aarch64").is_err());
    }
    #[test]
    fn root_relative_xbps_links_are_relocated_and_wrong_elf_is_rejected() {
        let root = Fixture::new();
        symlink("/usr/bin/tool", root.0.join("usr/bin/alias")).unwrap();
        assert!(inspect(&root.0, "aarch64").is_err());
        relocate(&root.0).unwrap();
        verify(&root.0, &root.software()).unwrap();
        fs::write(
            root.0.join("usr/bin/tool"),
            b"\x7fELF\x02\x01\0\0\0\0\0\0\0\0\0\0\0\0\x3e\0",
        )
        .unwrap();
        assert!(inspect(&root.0, "aarch64").is_err());
    }
}

/// Use the package's own glibc loader when present. Static ELFs and scripts
/// execute normally; dynamically linked binaries retain their packaged ABI.
pub fn executable(root: &Path, relative: &str) -> Result<Vec<String>> {
    let path = root.join(relative);
    let mut file = File::open(&path)?;
    let mut header = [0u8; 64];
    let n = file.read(&mut header)?;
    if n == header.len() && header.starts_with(b"\x7fELF") && header[4] == 2 && header[5] == 1 {
        use std::io::{Seek, SeekFrom};
        let offset = u64::from_le_bytes(header[32..40].try_into().unwrap());
        let size = u16::from_le_bytes(header[54..56].try_into().unwrap()) as u64;
        let count = u16::from_le_bytes(header[56..58].try_into().unwrap()) as u64;
        if size >= 56 && count <= 1024 {
            for index in 0..count {
                file.seek(SeekFrom::Start(
                    offset
                        .checked_add(index * size)
                        .ok_or_else(|| Error("ELF header overflow".into()))?,
                ))?;
                let mut program = [0u8; 56];
                file.read_exact(&mut program)?;
                if u32::from_le_bytes(program[..4].try_into().unwrap()) == 3 {
                    for name in ["lib/ld-linux-aarch64.so.1", "usr/lib/ld-linux-aarch64.so.1"] {
                        let loader = root.join(name);
                        if loader.is_file() {
                            return Ok(vec![
                                loader.display().to_string(),
                                "--library-path".into(),
                                format!("{0}/usr/lib:{0}/lib", root.display()),
                                path.display().to_string(),
                            ]);
                        }
                    }
                    break;
                }
            }
        }
    }
    Ok(vec![path.display().to_string()])
}