Ratbin/src/desktop.rs

use std::{fmt, path::Path};

const MAIN: &str = "Desktop Entry";

#[derive(Clone, Debug)]
pub struct DesktopEntry {
    groups: Vec<Group>,
}

#[derive(Clone, Debug)]
struct Group {
    name: Option<String>,
    lines: Vec<Line>,
}

#[derive(Clone, Debug)]
enum Line {
    Pair(String, String),
    Raw(String),
}

impl DesktopEntry {
    pub fn new() -> Self {
        Self {
            groups: vec![Group {
                name: Some(MAIN.into()),
                lines: Vec::new(),
            }],
        }
    }

    pub fn parse(text: &str) -> Self {
        let mut groups = vec![Group {
            name: None,
            lines: Vec::new(),
        }];
        for line in text.lines() {
            let trimmed = line.trim();
            let current = groups.last_mut().expect("groups is never empty");
            if trimmed.starts_with('[') && trimmed.ends_with(']') {
                groups.push(Group {
                    name: Some(trimmed[1..trimmed.len() - 1].to_string()),
                    lines: Vec::new(),
                });
            } else if !trimmed.starts_with('#')
                && let Some((key, value)) = line.split_once('=')
            {
                current
                    .lines
                    .push(Line::Pair(key.trim().into(), value.trim_start().into()));
            } else {
                current.lines.push(Line::Raw(line.into()));
            }
        }
        if groups[0].lines.is_empty() {
            groups.remove(0);
        }
        Self { groups }
    }

    pub fn get(&self, key: &str) -> Option<String> {
        self.group(MAIN)?.lines.iter().find_map(|line| match line {
            Line::Pair(k, v) if k == key => Some(unescape(v)),
            _ => None,
        })
    }

    pub fn set(&mut self, key: &str, value: &str) {
        let value = escape(value);
        let group = self.main_mut();
        for line in &mut group.lines {
            if let Line::Pair(k, v) = line
                && k == key
            {
                *v = value;
                return;
            }
        }
        group.lines.push(Line::Pair(key.into(), value));
    }

    pub fn remove(&mut self, key: &str) {
        self.main_mut()
            .lines
            .retain(|line| !matches!(line, Line::Pair(k, _) if k == key));
    }

    pub fn remove_localized(&mut self, key: &str) {
        self.main_mut().lines.retain(|line| match line {
            Line::Pair(k, _) => !(k.starts_with(key) && k[key.len()..].starts_with('[')),
            Line::Raw(_) => true,
        });
    }

    pub fn map_values(&mut self, key: &str, f: impl Fn(&str) -> String) {
        for group in &mut self.groups {
            for line in &mut group.lines {
                if let Line::Pair(k, v) = line
                    && k == key
                {
                    *v = escape(&f(&unescape(v)));
                }
            }
        }
    }

    fn group(&self, name: &str) -> Option<&Group> {
        self.groups.iter().find(|g| g.name.as_deref() == Some(name))
    }

    fn main_mut(&mut self) -> &mut Group {
        let index = match self
            .groups
            .iter()
            .position(|g| g.name.as_deref() == Some(MAIN))
        {
            Some(index) => index,
            None => {
                let at = usize::from(self.groups.first().is_some_and(|g| g.name.is_none()));
                self.groups.insert(
                    at,
                    Group {
                        name: Some(MAIN.into()),
                        lines: Vec::new(),
                    },
                );
                at
            }
        };
        &mut self.groups[index]
    }
}

impl fmt::Display for DesktopEntry {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for group in &self.groups {
            if let Some(name) = &group.name {
                writeln!(f, "[{name}]")?;
            }
            for line in &group.lines {
                match line {
                    Line::Pair(k, v) => writeln!(f, "{k}={v}")?,
                    Line::Raw(raw) => writeln!(f, "{raw}")?,
                }
            }
        }
        Ok(())
    }
}

fn unescape(value: &str) -> String {
    let mut out = String::with_capacity(value.len());
    let mut chars = value.chars();
    while let Some(c) = chars.next() {
        if c != '\\' {
            out.push(c);
            continue;
        }
        match chars.next() {
            Some('s') => out.push(' '),
            Some('n') => out.push('\n'),
            Some('t') => out.push('\t'),
            Some('r') => out.push('\r'),
            Some('\\') => out.push('\\'),
            Some(other) => {
                out.push('\\');
                out.push(other);
            }
            None => out.push('\\'),
        }
    }
    out
}

fn escape(value: &str) -> String {
    value
        .replace('\\', "\\\\")
        .replace('\n', "\\n")
        .replace('\t', "\\t")
        .replace('\r', "\\r")
}

pub fn quote_exec_arg(arg: &str) -> String {
    const RESERVED: &[char] = &[
        ' ', '\t', '\n', '"', '\'', '\\', '>', '<', '~', '|', '&', ';', '$', '*', '?', '#', '(',
        ')', '`',
    ];
    let arg = arg.replace('%', "%%");
    if !arg.is_empty() && !arg.contains(RESERVED) {
        return arg;
    }
    let mut out = String::from('"');
    for c in arg.chars() {
        if matches!(c, '"' | '`' | '$' | '\\') {
            out.push('\\');
        }
        out.push(c);
    }
    out.push('"');
    out
}

pub fn exec_args(exec: &str) -> &str {
    let exec = exec.trim_start();
    let end = match exec.strip_prefix('"') {
        Some(rest) => {
            let mut escaped = false;
            rest.char_indices()
                .find_map(|(i, c)| {
                    match c {
                        _ if escaped => escaped = false,
                        '\\' => escaped = true,
                        '"' => return Some(i + 2),
                        _ => {}
                    }
                    None
                })
                .unwrap_or(exec.len())
        }
        None => exec.find(char::is_whitespace).unwrap_or(exec.len()),
    };
    exec[end..].trim()
}

pub fn build_exec(program: &Path, args: &str) -> String {
    let program = quote_exec_arg(&program.to_string_lossy());
    if args.is_empty() {
        program
    } else {
        format!("{program} {args}")
    }
}