proto-reg 0.1.0

//! Thin wrappers around the Windows registry (`winreg`).
//!
//! Protocols are stored under `HKEY_CLASSES_ROOT\<scheme>`. On disk this is
//! either `HKLM\Software\Classes\<scheme>` (machine-wide, requires an
//! elevated process) or `HKCU\Software\Classes\<scheme>` (per-user).
//!
//! * Reads use the merged `HKEY_CLASSES_ROOT` view.
//! * Writes target `HKLM\Software\Classes`, mirroring what importing the
//!   `.reg` files generated by the original helpers did.

use winreg::RegKey;
use winreg::enums::*;

use crate::error::{Error, Result};
use crate::protocol::Protocol;

/// `Software\Classes` under both `HKLM` and `HKCU`.
const CLASSES_PATH: &str = r"Software\Classes";

/// Windows error code `ERROR_FILE_NOT_FOUND`.
const ERROR_FILE_NOT_FOUND: i32 = 2;
/// Windows error code `ERROR_PATH_NOT_FOUND`.
const ERROR_PATH_NOT_FOUND: i32 = 3;
/// Windows error code `ERROR_ACCESS_DENIED`.
const ERROR_ACCESS_DENIED: i32 = 5;

/// Returns `true` when the current process is running with administrator
/// privileges.
pub(crate) fn is_elevated() -> bool {
    is_admin::is_admin()
}

/// Checks that the current process is running with administrator privileges,
/// returning [`Error::ElevationRequired`] otherwise.
///
/// Write operations target `HKLM\Software\Classes`, so this is a fast,
/// user-friendly pre-flight check before touching the registry.
pub(crate) fn ensure_elevated() -> Result<()> {
    if is_elevated() {
        Ok(())
    } else {
        Err(Error::ElevationRequired(
            "the current process is not running with administrator privileges".to_string(),
        ))
    }
}

/// Opens `HKLM\Software\Classes` for writing, mapping access-denied errors to
/// the friendlier [`Error::ElevationRequired`]. Kept as a safety net on top of
/// [`ensure_elevated`].
fn hklm_classes() -> Result<RegKey> {
    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
    match hklm.open_subkey_with_flags(CLASSES_PATH, KEY_READ | KEY_WRITE) {
        Ok(key) => Ok(key),
        Err(err) if err.raw_os_error() == Some(ERROR_ACCESS_DENIED) => {
            Err(Error::ElevationRequired(err.to_string()))
        }
        Err(err) => Err(err.into()),
    }
}

/// Opens `path` under `root` read-only, returning `Ok(None)` when the key
/// does not exist.
fn open_optional(root: &RegKey, path: &str) -> Result<Option<RegKey>> {
    match root.open_subkey_with_flags(path, KEY_READ) {
        Ok(key) => Ok(Some(key)),
        Err(err)
            if matches!(
                err.raw_os_error(),
                Some(ERROR_FILE_NOT_FOUND | ERROR_PATH_NOT_FOUND)
            ) =>
        {
            Ok(None)
        }
        Err(err) => Err(err.into()),
    }
}

/// Reads the protocol `name` from the merged `HKEY_CLASSES_ROOT` view.
pub(crate) fn read_protocol(name: &str) -> Result<Option<Protocol>> {
    let hkcr = RegKey::predef(HKEY_CLASSES_ROOT);
    let Some(key) = open_optional(&hkcr, name)? else {
        return Ok(None);
    };
    Ok(Some(protocol_from_key(&key, name)))
}

/// Reads a protocol from an already-open `HKEY_CLASSES_ROOT` key.
fn protocol_from_key(key: &RegKey, name: &str) -> Protocol {
    let description = key.get_value("").unwrap_or_default();
    let url_protocol = key.get_value::<String, _>("URL Protocol").is_ok();
    let icon = key
        .open_subkey("DefaultIcon")
        .ok()
        .and_then(|icon_key| icon_key.get_value("").ok());
    let command = key
        .open_subkey(r"shell\open\command")
        .ok()
        .and_then(|command_key| command_key.get_value("").ok())
        .unwrap_or_default();

    Protocol {
        name: name.to_string(),
        description,
        url_protocol,
        icon,
        command,
    }
}

/// Writes (creates or updates) a protocol under `HKLM\Software\Classes`.
///
/// Requires administrator privileges; fails fast with
/// [`Error::ElevationRequired`] when the process is not elevated.
pub(crate) fn write_protocol(protocol: &Protocol) -> Result<()> {
    ensure_elevated()?;
    let classes = hklm_classes()?;

    // Root key: description + URL protocol marker.
    let (root, _disposition) = classes.create_subkey(&protocol.name)?;
    root.set_value("", &protocol.description)?;
    if protocol.url_protocol {
        root.set_value("URL Protocol", &"")?;
    } else {
        let _ = root.delete_value("URL Protocol");
    }

    // DefaultIcon (drop a stale icon when the new protocol has none).
    match &protocol.icon {
        Some(icon) => {
            let (icon_key, _) = root.create_subkey("DefaultIcon")?;
            icon_key.set_value("", icon)?;
        }
        None => {
            if root.open_subkey("DefaultIcon").is_ok() {
                let _ = root.delete_subkey_all("DefaultIcon");
            }
        }
    }

    // shell\open\command
    let (shell, _) = root.create_subkey(r"shell\open\command")?;
    shell.set_value("", &protocol.command)?;

    Ok(())
}

/// Deletes the protocol from both `HKLM\Software\Classes` and
/// `HKCU\Software\Classes`.
///
/// Returns `true` when at least one hive contained the protocol. Requires
/// administrator privileges **only** when the protocol lives in the
/// machine-wide hive; per-user deletions work without elevation.
pub(crate) fn delete_protocol(name: &str) -> Result<bool> {
    let mut removed = false;

    // Machine-wide hive (requires elevation).
    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
    let in_hklm = hklm
        .open_subkey(CLASSES_PATH)
        .map(|classes| classes.open_subkey(name).is_ok())
        .unwrap_or(false);
    if in_hklm {
        ensure_elevated()?;
        hklm_classes()?.delete_subkey_all(name)?;
        removed = true;
    }

    // Per-user hive (no elevation needed).
    let hkcu = RegKey::predef(HKEY_CURRENT_USER);
    if let Ok(classes) = hkcu.open_subkey_with_flags(CLASSES_PATH, KEY_READ | KEY_WRITE)
        && classes.open_subkey(name).is_ok()
    {
        classes.delete_subkey_all(name)?;
        removed = true;
    }

    Ok(removed)
}

/// Lists every protocol under `HKEY_CLASSES_ROOT` that carries the
/// `"URL Protocol"` marker.
pub(crate) fn list_url_protocols() -> Result<Vec<Protocol>> {
    let hkcr = RegKey::predef(HKEY_CLASSES_ROOT);
    let mut protocols = Vec::new();

    for entry in hkcr.enum_keys() {
        let name = entry?;
        // Fast filter: only keys that look like URL schemes.
        if !Protocol::is_valid_name(&name) {
            continue;
        }
        // Open once and skip keys without the "URL Protocol" marker before
        // reading the (more expensive) subkey values.
        let Some(key) = open_optional(&hkcr, &name)? else {
            continue;
        };
        if key.get_value::<String, _>("URL Protocol").is_err() {
            continue;
        }
        protocols.push(protocol_from_key(&key, &name));
    }

    Ok(protocols)
}

/// Chrome/Edge group-policy keys under `HKLM\Software\Policies` that control
/// the external-protocol confirmation dialog.
const BROWSER_POLICY_KEYS: &[&str] = &[
    r"Software\Policies\Google\Chrome",
    r"Software\Policies\Microsoft\Edge",
];

/// Group-policy value that shows the "Always open these links" checkbox.
const ALWAYS_OPEN_CHECKBOX_VALUE: &str = "ExternalProtocolDialogShowAlwaysOpenCheckbox";

/// Shows or hides the "Always open these links" checkbox in the Chrome/Edge
/// external-protocol dialog.
///
/// Modern Chrome/Edge hide this checkbox by default (to deter drive-by
/// downloads), so users must confirm the prompt on every click.
///
/// * `enabled = true` writes the `ExternalProtocolDialogShowAlwaysOpenCheckbox`
///   value (DWORD 1) under `HKLM\Software\Policies\...` for Chrome and Edge,
///   making them show the checkbox again.
/// * `enabled = false` removes the value, restoring the default behavior.
///
/// Requires administrator privileges.
pub(crate) fn set_browser_policy(enabled: bool) -> Result<()> {
    ensure_elevated()?;
    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
    for path in BROWSER_POLICY_KEYS {
        let (key, _) = hklm.create_subkey(path)?;
        if enabled {
            key.set_value(ALWAYS_OPEN_CHECKBOX_VALUE, &1u32)?;
        } else {
            match key.delete_value(ALWAYS_OPEN_CHECKBOX_VALUE) {
                // Not set yet — already in the requested state.
                Err(err) if err.raw_os_error() == Some(ERROR_FILE_NOT_FOUND) => {}
                other => other?,
            }
        }
    }
    Ok(())
}