proto-reg 0.1.0

Manage custom URL protocols in the Windows registry
//! `proto-reg` — command-line tool for managing custom URL protocols on
//! Windows.

use std::path::PathBuf;

use clap::{Parser, Subcommand};
use proto_reg::{Error, Protocol, ProtocolManager};

#[derive(Debug, Parser)]
#[command(
    name = "proto-reg",
    version,
    about = "Manage custom URL protocols in the Windows registry",
    long_about = concat!(
        "Manage custom URL protocols (e.g. `myapp://`) in the Windows registry.\n\n",
        "Protocols live under `HKEY_CLASSES_ROOT\\<scheme>` and are handled by\n",
        "`shell\\open\\command` entries where `%1` is replaced with the URL.\n\n",
        "NOTE: `add`, `remove`, `restore` and `browser-policy` write to the\n",
        "registry and require an elevated (Administrator) terminal. `query`,\n",
        "`list` and `backup` are read-only and work without elevation; `backup`\n",
        "prints JSON to stdout unless `-o <file>` is given."
    )
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    /// Show details of a registered protocol
    Query {
        /// Protocol scheme name, e.g. "myapp"
        name: String,
    },

    /// List all registered URL protocols
    List {
        /// Print the protocols as JSON
        #[arg(long)]
        json: bool,
    },

    /// Register a new protocol or update an existing one
    Add {
        /// Protocol scheme name, e.g. "myapp"
        name: String,
        /// Command line executed for the protocol; "%1" is replaced with the URL
        #[arg(short, long)]
        command: String,
        /// Friendly description shown in the registry
        #[arg(short, long)]
        description: Option<String>,
        /// Icon path shown for the protocol
        #[arg(short, long)]
        icon: Option<String>,
        /// Do not mark the protocol as a URL protocol
        #[arg(long)]
        no_url_protocol: bool,
        /// Do not also show the "Always open these links" checkbox in
        /// Chrome/Edge (enabled by default)
        #[arg(long)]
        no_browser_policy: bool,
    },

    /// Delete a protocol from the registry
    Remove {
        /// Protocol scheme name, e.g. "myapp"
        name: String,
    },

    /// Export URL protocols as JSON (prints to stdout, or to a file with -o)
    Backup {
        /// Protocol scheme name, e.g. "mpv-easy"; all protocols when omitted
        name: Option<String>,
        /// Write the JSON to this file instead of stdout
        #[arg(short, long, value_name = "FILE")]
        output: Option<PathBuf>,
    },

    /// Restore protocols from a JSON backup file
    Restore {
        /// Input JSON file created by `proto-reg backup`
        input: PathBuf,
        /// Restore only this protocol from the backup file
        #[arg(long, value_name = "NAME")]
        name: Option<String>,
    },

    /// Show the "Always open these links" checkbox in the Chrome/Edge dialog
    BrowserPolicy {
        /// Remove the policy instead of setting it (browser default behavior)
        #[arg(long)]
        remove: bool,
    },
}

fn main() {
    if let Err(error) = run() {
        eprintln!("error: {error}");
        if matches!(error, Error::ElevationRequired(_)) {
            eprintln!("hint: re-run from an elevated (Administrator) terminal");
        }
        std::process::exit(1);
    }
}

fn run() -> Result<(), Error> {
    let cli = Cli::parse();
    match cli.command {
        Command::Query { name } => match ProtocolManager::query(&name)? {
            Some(protocol) => print_protocol(&protocol),
            None => {
                eprintln!("error: protocol `{name}` is not registered");
                std::process::exit(1);
            }
        },

        Command::List { json } => {
            let protocols = ProtocolManager::list()?;
            if json {
                println!("{}", serde_json::to_string_pretty(&protocols)?);
            } else if protocols.is_empty() {
                println!("no URL protocols registered");
            } else {
                for protocol in &protocols {
                    println!("{}", protocol.name);
                }
            }
        }

        Command::Add {
            name,
            command,
            description,
            icon,
            no_url_protocol,
            no_browser_policy,
        } => {
            let mut protocol = Protocol::new(name, command)?;
            protocol.description = description.unwrap_or_default();
            protocol.icon = icon;
            protocol.url_protocol = !no_url_protocol;
            ProtocolManager::add(&protocol)?;
            println!("registered protocol `{}`", protocol.name);
            if !no_browser_policy {
                ProtocolManager::set_browser_policy(true)?;
                println!(
                    "enabled the \"Always open these links\" checkbox in Chrome/Edge"
                );
            }
        }

        Command::Remove { name } => {
            if ProtocolManager::remove(&name)? {
                println!("removed protocol `{name}`");
            } else {
                eprintln!("error: protocol `{name}` is not registered");
                std::process::exit(1);
            }
        }

        Command::Backup { name, output } => {
            let protocols = match name {
                Some(name) => match ProtocolManager::backup_protocol_to_json(&name)? {
                    Some(protocol) => vec![protocol],
                    None => {
                        eprintln!("error: protocol `{name}` is not registered");
                        std::process::exit(1);
                    }
                },
                None => ProtocolManager::backup_to_json()?,
            };

            let json = serde_json::to_string_pretty(&protocols)?;
            match output {
                Some(path) => {
                    std::fs::write(&path, json)?;
                    println!(
                        "backed up {} protocol(s) to {}",
                        protocols.len(),
                        path.display()
                    );
                }
                None => println!("{json}"),
            }
        }

        Command::Restore { input, name } => match name {
            Some(name) => match ProtocolManager::restore_protocol_from_file(&name, &input)? {
                true => println!("restored protocol `{name}` from {}", input.display()),
                false => {
                    eprintln!("error: protocol `{name}` is not in {}", input.display());
                    std::process::exit(1);
                }
            },
            None => {
                let report = ProtocolManager::restore_from_file(&input)?;
                println!(
                    "restored {} protocol(s) from {}",
                    report.restored,
                    input.display()
                );
                for (name, message) in &report.failed {
                    eprintln!("warning: could not restore `{name}`: {message}");
                }
                if !report.is_success() {
                    std::process::exit(1);
                }
            }
        },
        Command::BrowserPolicy { remove } => {
            ProtocolManager::set_browser_policy(!remove)?;
            if remove {
                println!("removed the Chrome/Edge dialog policy (browser default restored)");
            } else {
                println!(
                    "Chrome/Edge will show the \"Always open these links\" checkbox"
                );
            }
        }
    }
    Ok(())
}

fn print_protocol(protocol: &Protocol) {
    println!("name:         {}", protocol.name);
    println!("description:  {}", protocol.description);
    println!(
        "url_protocol: {}",
        if protocol.url_protocol { "yes" } else { "no" }
    );
    println!(
        "icon:         {}",
        protocol.icon.as_deref().unwrap_or("(none)")
    );
    println!("command:      {}", protocol.command);
}