Ratbin/src/appimage.rs

use std::{
    collections::HashMap,
    ffi::OsStr,
    fs::File,
    io::{BufReader, Read, Seek, SeekFrom},
    path::{Component, Path, PathBuf},
};

use anyhow::{Context, Result};
use backhand::{FilesystemReader, InnerNode, SquashfsFileReader};

use crate::desktop::DesktopEntry;

pub struct Icon {
    pub data: Vec<u8>,
    pub ext: &'static str,
}

pub struct Embedded {
    pub desktop: Option<DesktopEntry>,
    pub icon: Option<Icon>,
}

pub fn read(path: &Path, offset: u64) -> Result<Embedded> {
    let fs = open(path, offset)?;
    let image = Image::new(&fs);
    let desktop = image.root_desktop();
    let icon_name = desktop.as_ref().and_then(|d| d.get("Icon"));
    let icon = image.icon(icon_name.as_deref());
    Ok(Embedded { desktop, icon })
}

pub fn desktop_id(path: &Path, offset: u64) -> Option<String> {
    let fs = open(path, offset).ok()?;
    let image = Image::new(&fs);
    let name = image.root_desktop_path()?.file_name()?;
    Some(name.to_string_lossy().into_owned())
}

pub fn update_info(path: &Path) -> Option<String> {
    let mut file = File::open(path).ok()?;
    let mut header = [0u8; 64];
    file.read_exact(&mut header).ok()?;
    if !header.starts_with(b"\x7fELF") {
        return None;
    }
    let little_endian = header[5] == 1;
    let wide = match header[4] {
        1 => false,
        2 => true,
        _ => return None,
    };
    let uint = |bytes: &[u8], offset: usize, len: usize| -> Option<u64> {
        let bytes = bytes.get(offset..offset + len)?;
        let fold = |acc: u64, &b: &u8| (acc << 8) | u64::from(b);
        Some(if little_endian {
            bytes.iter().rev().fold(0, fold)
        } else {
            bytes.iter().fold(0, fold)
        })
    };
    let (shoff, shentsize, shnum, shstrndx) = if wide {
        (
            uint(&header, 0x28, 8)?,
            uint(&header, 0x3A, 2)?,
            uint(&header, 0x3C, 2)?,
            uint(&header, 0x3E, 2)?,
        )
    } else {
        (
            uint(&header, 0x20, 4)?,
            uint(&header, 0x2E, 2)?,
            uint(&header, 0x30, 2)?,
            uint(&header, 0x32, 2)?,
        )
    };
    let mut table = vec![0u8; usize::try_from(shentsize * shnum).ok()?];
    file.seek(SeekFrom::Start(shoff)).ok()?;
    file.read_exact(&mut table).ok()?;
    let section = |index: u64| -> Option<(u64, u64, u64)> {
        let base = usize::try_from(index * shentsize).ok()?;
        let name = uint(&table, base, 4)?;
        Some(if wide {
            (
                name,
                uint(&table, base + 0x18, 8)?,
                uint(&table, base + 0x20, 8)?,
            )
        } else {
            (
                name,
                uint(&table, base + 0x10, 4)?,
                uint(&table, base + 0x14, 4)?,
            )
        })
    };
    let mut read_at = |offset: u64, size: u64| -> Option<Vec<u8>> {
        let mut data = vec![0u8; usize::try_from(size.min(1 << 16)).ok()?];
        file.seek(SeekFrom::Start(offset)).ok()?;
        file.read_exact(&mut data).ok()?;
        Some(data)
    };
    let (_, names_offset, names_size) = section(shstrndx)?;
    let names = read_at(names_offset, names_size)?;
    let (_, offset, size) = (0..shnum).filter_map(section).find(|(name, _, _)| {
        usize::try_from(*name)
            .ok()
            .and_then(|start| names.get(start..))
            .is_some_and(|rest| rest.starts_with(b".upd_info\0"))
    })?;
    let data = read_at(offset, size)?;
    let end = data.iter().position(|&b| b == 0).unwrap_or(data.len());
    let info = String::from_utf8_lossy(&data[..end]).trim().to_string();
    (!info.is_empty()).then_some(info)
}

fn open(path: &Path, offset: u64) -> Result<FilesystemReader<'static>> {
    let file = BufReader::new(File::open(path)?);
    FilesystemReader::from_reader_with_offset(file, offset).with_context(|| {
        format!(
            "cannot read the squashfs image embedded in {}",
            path.display()
        )
    })
}

struct Image<'a, 'b> {
    fs: &'a FilesystemReader<'b>,
    nodes: HashMap<PathBuf, &'a InnerNode<SquashfsFileReader>>,
}

impl<'a, 'b> Image<'a, 'b> {
    fn new(fs: &'a FilesystemReader<'b>) -> Self {
        let nodes = fs
            .files()
            .map(|node| (normalize(&node.fullpath), &node.inner))
            .collect();
        Self { fs, nodes }
    }

    fn read(&self, path: &Path) -> Option<Vec<u8>> {
        let mut path = normalize(path);
        for _ in 0..16 {
            match self.nodes.get(&path)? {
                InnerNode::File(file) => {
                    let mut data = Vec::new();
                    self.fs.file(file).reader().read_to_end(&mut data).ok()?;
                    return Some(data);
                }
                InnerNode::Symlink(symlink) => {
                    let parent = path.parent().unwrap_or(Path::new(""));
                    path = normalize(&parent.join(&symlink.link));
                }
                _ => return None,
            }
        }
        None
    }

    fn root_desktop_path(&self) -> Option<&PathBuf> {
        self.nodes
            .keys()
            .filter(|p| p.components().count() == 1 && p.extension() == Some(OsStr::new("desktop")))
            .min()
    }

    fn root_desktop(&self) -> Option<DesktopEntry> {
        let data = self.read(self.root_desktop_path()?)?;
        Some(DesktopEntry::parse(&String::from_utf8_lossy(&data)))
    }

    fn icon(&self, name: Option<&str>) -> Option<Icon> {
        let mut candidates = Vec::new();
        if let Some(name) = name.filter(|n| !n.is_empty()) {
            let name = Path::new(name)
                .file_name()
                .and_then(OsStr::to_str)
                .unwrap_or(name);
            candidates.push(PathBuf::from(name));
            for ext in ["svg", "png", "xpm"] {
                candidates.push(PathBuf::from(format!("{name}.{ext}")));
            }
            candidates.extend(self.hicolor(name));
        }
        candidates.push(PathBuf::from(".DirIcon"));
        candidates
            .iter()
            .find_map(|path| self.read(path).and_then(icon_from))
    }

    fn hicolor(&self, name: &str) -> Vec<PathBuf> {
        let mut found: Vec<(u32, PathBuf)> = self
            .nodes
            .keys()
            .filter_map(|path| {
                let rest = path.strip_prefix("usr/share/icons/hicolor").ok()?;
                if path.file_stem()? != OsStr::new(name) {
                    return None;
                }
                let size = rest.components().next()?.as_os_str().to_str()?;
                let rank = if size == "scalable" {
                    u32::MAX
                } else {
                    size.split('x').next()?.parse().ok()?
                };
                Some((rank, path.clone()))
            })
            .collect();
        found.sort_by_key(|(rank, _)| std::cmp::Reverse(*rank));
        found.into_iter().map(|(_, path)| path).collect()
    }
}

fn icon_from(data: Vec<u8>) -> Option<Icon> {
    let ext = if data.starts_with(b"\x89PNG") {
        "png"
    } else if data.starts_with(b"/* XPM */") {
        "xpm"
    } else if String::from_utf8_lossy(&data[..data.len().min(4096)]).contains("<svg") {
        "svg"
    } else {
        return None;
    };
    Some(Icon { data, ext })
}

fn normalize(path: &Path) -> PathBuf {
    let mut out = PathBuf::new();
    for component in path.components() {
        match component {
            Component::Normal(part) => out.push(part),
            Component::ParentDir => {
                out.pop();
            }
            _ => {}
        }
    }
    out
}