proto-reg 0.1.0

Manage custom URL protocols in the Windows registry
//! Protocol model and management operations.

use std::path::Path;

use serde::{Deserialize, Serialize};

use crate::backup::{self, RestoreReport};
use crate::error::{Error, Result};
use crate::registry;

/// A custom URL protocol registered in the Windows registry.
///
/// A protocol is stored under `HKEY_CLASSES_ROOT\<name>`:
///
/// ```text
/// HKEY_CLASSES_ROOT\<name>
///     (Default)          = description
///     "URL Protocol"     = ""    (marks the key as a URL protocol)
///     DefaultIcon
///         (Default)      = icon
///     shell\open\command
///         (Default)      = command ("%1" is replaced by the URL)
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Protocol {
    /// Scheme name, e.g. `"myapp"`. Must be a valid URL scheme (RFC 3986).
    pub name: String,
    /// Friendly description shown in the registry; the key's default value.
    #[serde(default)]
    pub description: String,
    /// Whether the `"URL Protocol"` marker is set. URL protocols only.
    #[serde(default = "default_url_protocol")]
    pub url_protocol: bool,
    /// Optional path to the icon shown by the OS; the `DefaultIcon` value.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub icon: Option<String>,
    /// Command line executed when the protocol is invoked; `%1` receives the
    /// URL. Example: `"C:\app.exe" "%1"`.
    pub command: String,
}

fn default_url_protocol() -> bool {
    true
}

impl Protocol {
    /// Creates a new protocol, validating the scheme name.
    ///
    /// `description` defaults to the empty string, `icon` to `None` and
    /// `url_protocol` to `true`; set the public fields to customize.
    pub fn new(name: impl Into<String>, command: impl Into<String>) -> Result<Self> {
        let name = name.into();
        Self::validate_name(&name)?;
        Ok(Self {
            name,
            description: String::new(),
            url_protocol: true,
            icon: None,
            command: command.into(),
        })
    }

    /// Returns `true` when `name` is a valid URL scheme: starts with a letter
    /// and contains only letters, digits, `+`, `-` or `.` (RFC 3986).
    pub fn is_valid_name(name: &str) -> bool {
        let mut chars = name.chars();
        match chars.next() {
            Some(first) if first.is_ascii_alphabetic() => {}
            _ => return false,
        }
        chars.all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.'))
    }

    /// Validates the scheme name, returning [`Error::InvalidProtocolName`] on
    /// failure.
    pub fn validate_name(name: &str) -> Result<()> {
        if Self::is_valid_name(name) {
            Ok(())
        } else {
            Err(Error::InvalidProtocolName {
                name: name.to_string(),
            })
        }
    }
}

/// Entry point for protocol management operations.
///
/// Read operations ([`Self::query`], [`Self::list`], backups) only need a
/// normal user account. Write operations ([`Self::add`], [`Self::remove`],
/// [`Self::restore_from_file`]) modify `HKLM\Software\Classes` and therefore
/// require an elevated (Administrator) process; use [`Self::is_elevated`] to
/// check beforehand.
pub struct ProtocolManager;

impl ProtocolManager {
    /// Returns `true` when the current process is running with administrator
    /// privileges.
    ///
    /// Write operations ([`Self::add`], [`Self::remove`],
    /// [`Self::restore_from_file`]) fail fast with
    /// [`Error::ElevationRequired`] when this returns `false`.
    pub fn is_elevated() -> bool {
        registry::is_elevated()
    }

    /// Returns the registered protocol `name`, or `None` if it is not
    /// registered.
    pub fn query(name: &str) -> Result<Option<Protocol>> {
        registry::read_protocol(name)
    }

    /// Returns `true` when the protocol is registered.
    pub fn exists(name: &str) -> Result<bool> {
        Ok(registry::read_protocol(name)?.is_some())
    }

    /// Registers a new protocol or updates an existing one.
    ///
    /// Requires administrator privileges.
    pub fn add(protocol: &Protocol) -> Result<()> {
        Protocol::validate_name(&protocol.name)?;
        registry::write_protocol(protocol)
    }

    /// Deletes a protocol from both the machine-wide (`HKLM\Software\Classes`)
    /// and the per-user (`HKCU\Software\Classes`) hives.
    ///
    /// Returns `true` when anything was removed. Requires administrator
    /// privileges when the protocol lives in the machine-wide hive.
    pub fn remove(name: &str) -> Result<bool> {
        Protocol::validate_name(name)?;
        registry::delete_protocol(name)
    }

    /// Lists every registered URL protocol (keys carrying the
    /// `"URL Protocol"` marker), sorted by name.
    pub fn list() -> Result<Vec<Protocol>> {
        let mut protocols = registry::list_url_protocols()?;
        protocols.sort_by(|a, b| a.name.cmp(&b.name));
        Ok(protocols)
    }

    /// Exports all registered URL protocols as a JSON array.
    pub fn backup_to_json() -> Result<Vec<Protocol>> {
        Self::list()
    }

    /// Exports all registered URL protocols to `path` as pretty JSON and
    /// returns them.
    pub fn backup_to_file(path: impl AsRef<Path>) -> Result<Vec<Protocol>> {
        let protocols = Self::backup_to_json()?;
        backup::write_to_file(&protocols, path)?;
        Ok(protocols)
    }

    /// Exports a single registered protocol as JSON.
    ///
    /// Returns `Ok(None)` when `name` is not registered.
    pub fn backup_protocol_to_json(name: &str) -> Result<Option<Protocol>> {
        Protocol::validate_name(name)?;
        registry::read_protocol(name)
    }

    /// Exports a single registered protocol to `path` as pretty JSON.
    ///
    /// Returns `Ok(None)` when `name` is not registered.
    pub fn backup_protocol_to_file(name: &str, path: impl AsRef<Path>) -> Result<Option<Protocol>> {
        let Some(protocol) = Self::backup_protocol_to_json(name)? else {
            return Ok(None);
        };
        backup::write_to_file(std::slice::from_ref(&protocol), path)?;
        Ok(Some(protocol))
    }

    /// Restores protocols from a JSON string.
    ///
    /// Invalid entries are skipped and reported in the returned
    /// [`RestoreReport`]. Requires administrator privileges.
    pub fn restore_from_json(json: &str) -> Result<RestoreReport> {
        let protocols = backup::parse_backup(json)?;
        restore_backup(&protocols)
    }

    /// Restores protocols from a JSON backup file created by
    /// [`Self::backup_to_file`]. Requires administrator privileges.
    pub fn restore_from_file(path: impl AsRef<Path>) -> Result<RestoreReport> {
        let protocols = backup::read_from_file(path)?;
        restore_backup(&protocols)
    }

    /// Restores only the protocol `name` from a JSON backup file.
    ///
    /// Returns `false` when the backup does not contain `name`. Requires
    /// administrator privileges.
    pub fn restore_protocol_from_file(name: &str, path: impl AsRef<Path>) -> Result<bool> {
        Protocol::validate_name(name)?;
        let protocols = backup::read_from_file(path)?;
        let Some(protocol) = protocols.iter().find(|protocol| protocol.name == name) else {
            return Ok(false);
        };
        registry::write_protocol(protocol)?;
        Ok(true)
    }

    /// Shows or hides the "Always open these links" checkbox in the
    /// Chrome/Edge external-protocol confirmation dialog.
    ///
    /// Modern Chrome/Edge hide this checkbox by default, so users must
    /// confirm the prompt on every click. `enabled = true` writes the
    /// `ExternalProtocolDialogShowAlwaysOpenCheckbox` group-policy value
    /// (DWORD 1) under `HKLM\Software\Policies\Google\Chrome` and
    /// `HKLM\Software\Policies\Microsoft\Edge` to show it again;
    /// `enabled = false` removes the value (default browser behavior).
    ///
    /// Requires administrator privileges.
    pub fn set_browser_policy(enabled: bool) -> Result<()> {
        registry::set_browser_policy(enabled)
    }
}

fn restore_backup(protocols: &[Protocol]) -> Result<RestoreReport> {
    // Fail fast instead of collecting one elevation error per protocol.
    if !protocols.is_empty() {
        registry::ensure_elevated()?;
    }

    let mut report = RestoreReport::default();
    for protocol in protocols {
        let outcome = (|| -> Result<()> {
            if protocol.command.trim().is_empty() {
                return Err(Error::InvalidBackup(format!(
                    "protocol `{}` has an empty command",
                    protocol.name
                )));
            }
            Protocol::validate_name(&protocol.name)?;
            registry::write_protocol(protocol)
        })();
        match outcome {
            Ok(()) => report.restored += 1,
            Err(err) => report.failed.push((protocol.name.clone(), err.to_string())),
        }
    }
    Ok(report)
}