keyfork/crates/keyfork/src/cli/shard.rs

use super::Keyfork;
use clap::{builder::PossibleValue, Parser, Subcommand, ValueEnum};
use keyfork_shard::Format as _;
use std::{
    io::{stdin, stdout, Read, Write},
    path::{Path, PathBuf},
};

const COULD_NOT_DETERMINE_FORMAT: &str = "could not determine format, try including --format";

#[derive(Debug, Clone)]
enum Format {
    OpenPGP(OpenPGP),
    P256(P256),
}

impl ValueEnum for Format {
    fn value_variants<'a>() -> &'a [Self] {
        &[Self::OpenPGP(OpenPGP), Self::P256(P256)]
    }

    fn to_possible_value(&self) -> Option<PossibleValue> {
        Some(match self {
            Format::OpenPGP(_) => PossibleValue::new("openpgp"),
            Format::P256(_) => PossibleValue::new("p256"),
        })
    }
}

trait ShardExec {
    fn split(
        &self,
        threshold: u8,
        max: u8,
        key_discovery: &Path,
        secret: &[u8],
        output: &mut (impl Write + Send + Sync),
    ) -> Result<(), Box<dyn std::error::Error>>;

    fn combine(
        &self,
        key_discovery: Option<&Path>,
        input: impl Read + Send + Sync,
        output: &mut impl Write,
    ) -> Result<(), Box<dyn std::error::Error>>;

    fn decrypt(
        &self,
        key_discovery: Option<&Path>,
        input: impl Read + Send + Sync,
    ) -> Result<(), Box<dyn std::error::Error>>;
}

#[derive(Clone, Debug)]
struct OpenPGP;

impl ShardExec for OpenPGP {
    fn split(
        &self,
        threshold: u8,
        max: u8,
        key_discovery: &Path,
        secret: &[u8],
        output: &mut (impl Write + Send + Sync),
    ) -> Result<(), Box<dyn std::error::Error>> {
        let opgp = keyfork_shard::openpgp::OpenPGP;
        opgp.shard_and_encrypt(threshold, max, secret, key_discovery, output)
    }

    fn combine(
        &self,
        key_discovery: Option<&Path>,
        input: impl Read + Send + Sync,
        output: &mut impl Write,
    ) -> Result<(), Box<dyn std::error::Error>>
    {
        let openpgp = keyfork_shard::openpgp::OpenPGP;
        let bytes = openpgp.decrypt_all_shards_to_secret(key_discovery, input)?;
        write!(output, "{}", smex::encode(bytes))?;

        Ok(())
    }

    fn decrypt(
        &self,
        key_discovery: Option<&Path>,
        input: impl Read + Send + Sync,
    ) -> Result<(), Box<dyn std::error::Error>>
    {
        let openpgp = keyfork_shard::openpgp::OpenPGP;
        openpgp.decrypt_one_shard_for_transport(key_discovery, input)?;
        Ok(())
    }
}

#[derive(Clone, Debug)]
struct P256;

#[derive(Subcommand, Clone, Debug)]
pub enum ShardSubcommands {
    /// Split a hex-encoded secret from input into multiple shares, using Shamir's Secret Sharing.
    ///
    /// The shares are encrypted once per key, with keys discovered either on-system or by
    /// prompting for hardware interactions. Metadata about decrypting keys is then stored and
    /// encrypted to all keys, to ensure any key that holds a share can then be used to begin the
    /// process of combining keys.
    Split {
        /// The amount of shares required to recombine a secret.
        #[arg(long)]
        threshold: u8,

        /// The total amount of shares to generate.
        #[arg(long)]
        max: u8,

        /// The path to discover public keys from.
        key_discovery: PathBuf,
    },

    /// Decrypt a single share and re-encrypt it to an ephemeral symmetric key using mnemonic-based
    /// prompts. The mnemonics can be sent over insecure channels.
    Transport {
        /// The path to load the shard from.
        shard: PathBuf,

        /// The path to discover private keys from.
        key_discovery: Option<PathBuf>,
    },

    /// Combine multiple encrypted shares into a hex-encoded secret, printed to stdout.
    ///
    /// This command only accepts input from `keyfork shard split`, and is dependent on the format
    /// used when splitting. Metadata is encrypted to all keys that may hold a share, so when using
    /// hardware metadata discovery, any hardware key used to split may be used to decrypt metadata
    /// used to combine.
    Combine {
        /// The path to load the shards from.
        shard: PathBuf,

        /// The path to discover private keys from.
        key_discovery: Option<PathBuf>,
    },
}

impl ShardSubcommands {
    pub fn handle(
        &self,
        shard: &Shard,
        _keyfork: &Keyfork,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let stdin = stdin();
        let mut stdout = stdout();
        let mut format = shard.format.clone();
        match self {
            ShardSubcommands::Split {
                threshold,
                max,
                key_discovery,
            } => {
                let input = std::io::read_to_string(stdin)?;
                assert!(threshold <= max, "threshold {threshold} <= max {max}");
                let secret = smex::decode(input.trim())?;
                match format {
                    Some(Format::OpenPGP(o)) => {
                        o.split(*threshold, *max, key_discovery, &secret, &mut stdout)
                    }
                    Some(Format::P256(_p)) => {
                        todo!()
                    }
                    None => panic!("{COULD_NOT_DETERMINE_FORMAT}"),
                }
            }
            ShardSubcommands::Transport {
                shard,
                key_discovery,
            } => {
                let shard_content = std::fs::read_to_string(shard)?;
                if shard_content.contains("BEGIN PGP MESSAGE") {
                    let _ = format.insert(Format::OpenPGP(OpenPGP));
                }

                match format {
                    Some(Format::OpenPGP(o)) => {
                        o.decrypt(key_discovery.as_deref(), shard_content.as_bytes())
                    }
                    Some(Format::P256(_p)) => todo!(),
                    None => panic!("{COULD_NOT_DETERMINE_FORMAT}"),
                }
            }
            ShardSubcommands::Combine {
                shard,
                key_discovery,
            } => {
                let shard_content = std::fs::read_to_string(shard)?;
                if shard_content.contains("BEGIN PGP MESSAGE") {
                    let _ = format.insert(Format::OpenPGP(OpenPGP));
                }

                match format {
                    Some(Format::OpenPGP(o)) => o.combine(
                        key_discovery.as_deref(),
                        shard_content.as_bytes(),
                        &mut stdout,
                    ),
                    Some(Format::P256(_p)) => {
                        todo!()
                    }
                    None => panic!("{COULD_NOT_DETERMINE_FORMAT}"),
                }
            }
        }
    }
}

#[derive(Parser, Debug, Clone)]
pub struct Shard {
    /// Which format to use for encoding/encrypting and decoding/decrypting shares.
    #[arg(long, value_enum, global = true)]
    format: Option<Format>,

    #[command(subcommand)]
    pub command: ShardSubcommands,
}