tiny-encrypt-rs/src/cmd_simple_encrypt_decrypt.rs

use crate::cmd_decrypt::try_decrypt_key;
use crate::config::TinyEncryptConfig;
use crate::consts::TINY_ENC_CONFIG_FILE;
use crate::spec::TinyEncryptEnvelop;
use crate::{cmd_encrypt, crypto_cryptor, util, util_env};
use base64::engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD};
use base64::Engine;
use clap::Args;
use rust_util::{debugging, opt_result, simple_error, XResult};
use serde::Serialize;
use std::process::exit;

// Reference: https://git.hatter.ink/hatter/tiny-encrypt-rs/issues/3
const SIMPLE_ENCRYPTION_HEADER: &str = "tinyencrypt-dir";
const SIMPLE_ENCRYPTION_DOT: &str = ".";

#[derive(Debug, Args)]
pub struct CmdSimpleEncrypt {
    /// Encryption profile (use default when --key-filter is assigned)
    #[arg(long, short = 'p')]
    pub profile: Option<String>,

    /// Encryption key filter (key_id or type:TYPE(e.g. ecdh, pgp, ecdh-p384, pgp-ed25519), multiple joined by ',', ALL for all)
    #[arg(long, short = 'k')]
    pub key_filter: Option<String>,

    /// Encrypt value from stdin
    #[arg(long)]
    pub value_stdin: bool,

    /// Encrypt value
    #[arg(long, short = 'v')]
    pub value: Option<String>,

    /// Encrypt value in bse64
    #[arg(long)]
    pub value_base64: Option<String>,

    /// Encrypt value in hex
    #[arg(long)]
    pub value_hex: Option<String>,
}

#[derive(Debug, Args)]
pub struct CmdSimpleDecrypt {
    /// PGP or PIV PIN
    #[arg(long, short = 'p')]
    pub pin: Option<String>,

    /// Decrypt key ID
    #[arg(long, short = 'k')]
    pub key_id: Option<String>,

    /// PIV slot
    #[arg(long, short = 's')]
    pub slot: Option<String>,

    /// Decrypt value from stdin
    #[arg(long)]
    pub value_stdin: bool,

    /// Decrypt value
    #[arg(long, short = 'v')]
    pub value: Option<String>,

    /// Decrypt result output type (plain, hex, bse64)
    #[arg(long, short = 'o')]
    pub output_type: Option<String>,
}

impl CmdSimpleEncrypt {
    pub fn get_value(&self) -> XResult<Option<Vec<u8>>> {
        if self.value_stdin {
            return Ok(Some(util::read_stdin()?));
        }
        if let Some(value) = &self.value {
            return Ok(Some(value.as_bytes().to_vec()));
        }
        if let Some(value_base64) = &self.value_base64 {
            return Ok(Some(opt_result!(STANDARD.decode(value_base64), "Parse value base64 failed: {}")));
        }
        if let Some(value_hex) = &self.value_hex {
            return Ok(Some(opt_result!(hex::decode(value_hex), "Parse value hex failed: {}")));
        }
        Ok(None)
    }
}

impl CmdSimpleDecrypt {
    pub fn get_value(&self) -> XResult<Option<String>> {
        if self.value_stdin {
            return Ok(Some(opt_result!(String::from_utf8(util::read_stdin()?), "Read stdin value failed: {}")));
        }
        Ok(self.value.clone())
    }
}

#[derive(Serialize)]
pub struct CmdResult {
    pub code: i32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<String>,
}

impl CmdResult {
    pub fn fail(code: i32, message: &str) -> Self {
        Self {
            code,
            message: Some(message.to_string()),
            result: None,
        }
    }

    pub fn success(result: &str) -> Self {
        Self {
            code: 0,
            message: None,
            result: Some(result.to_string()),
        }
    }

    pub fn print_exit(&self) -> ! {
        let result = serde_json::to_string_pretty(self).unwrap();
        println!("{}", result);
        exit(self.code)
    }
}

pub fn simple_encrypt(cmd_simple_encrypt: CmdSimpleEncrypt) -> XResult<()> {
    if let Err(inner_result_error) = inner_simple_encrypt(cmd_simple_encrypt) {
        CmdResult::fail(-1, &format!("{}", inner_result_error)).print_exit();
    }
    Ok(())
}

pub fn simple_decrypt(cmd_simple_decrypt: CmdSimpleDecrypt) -> XResult<()> {
    if let Err(inner_result_error) = inner_simple_decrypt(cmd_simple_decrypt) {
        CmdResult::fail(-1, &format!("{}", inner_result_error)).print_exit();
    }
    Ok(())
}

pub fn inner_simple_encrypt(cmd_simple_encrypt: CmdSimpleEncrypt) -> XResult<()> {
    let config = TinyEncryptConfig::load(TINY_ENC_CONFIG_FILE)?;
    debugging!("Found tiny encrypt config: {:?}", config);
    let envelops = config.find_envelops(&cmd_simple_encrypt.profile, &cmd_simple_encrypt.key_filter)?;
    if envelops.is_empty() { return simple_error!("Cannot find any valid envelops"); }
    debugging!("Found envelops: {:?}", envelops);
    let envelop_tkids: Vec<_> = envelops.iter()
        .map(|e| format!("{}:{}", e.r#type.get_name(), e.kid))
        .collect();
    debugging!("Matched {} envelop(s): \n- {}", envelops.len(), envelop_tkids.join("\n- "));

    if envelop_tkids.is_empty() {
        return simple_error!("no matched envelops found");
    }

    let value = match cmd_simple_encrypt.get_value()? {
        None => return simple_error!("--value-stdin/value/value-base64/value-hex must assign one"),
        Some(value) => value,
    };

    let cryptor = crypto_cryptor::get_cryptor_by_encryption_algorithm(&None)?;
    let envelops = cmd_encrypt::encrypt_envelops(cryptor, &value, &envelops)?;

    let envelops_json = serde_json::to_string(&envelops)?;
    let simple_encrypt_result = format!("{}.{}",
                                        SIMPLE_ENCRYPTION_HEADER,
                                        URL_SAFE_NO_PAD.encode(envelops_json.as_bytes())
    );

    CmdResult::success(&simple_encrypt_result).print_exit();
}

pub fn inner_simple_decrypt(cmd_simple_decrypt: CmdSimpleDecrypt) -> XResult<()> {
    let config = TinyEncryptConfig::load(TINY_ENC_CONFIG_FILE).ok();

    let pin = cmd_simple_decrypt.pin.clone().or_else(util_env::get_pin);
    let slot = cmd_simple_decrypt.slot.clone();

    let output_type = cmd_simple_decrypt.output_type.as_deref().unwrap_or("plain");
    match output_type {
        "plain" | "hex" | "base64" => (),
        _ => return simple_error!("not supported output type: {}", output_type),
    };

    let value = match cmd_simple_decrypt.get_value()? {
        None => return simple_error!("--value-stdin/value must assign one"),
        Some(value) => value,
    };
    let value_parts = value.trim().split(SIMPLE_ENCRYPTION_DOT).collect::<Vec<_>>();
    if value_parts.len() != 2 {
        return simple_error!("bad value format: {}", value);
    }
    if value_parts[0] != SIMPLE_ENCRYPTION_HEADER {
        return simple_error!("bad value format: {}", value);
    }
    let envelopes_json = opt_result!(URL_SAFE_NO_PAD.decode(&value_parts[1]), "bad value format: {}");
    let envelops: Vec<TinyEncryptEnvelop> = match serde_json::from_slice(&envelopes_json) {
        Err(_) => return simple_error!("bad value format: {}", value),
        Ok(value) => value,
    };

    let filter_envelops = envelops.iter().filter(|e| {
        match &cmd_simple_decrypt.key_id {
            None => true,
            Some(key_id) => &e.kid == key_id,
        }
    }).collect::<Vec<_>>();
    if filter_envelops.is_empty() {
        return simple_error!("no envelops found: {:?}", cmd_simple_decrypt.key_id);
    }
    if filter_envelops.len() > 1 {
        return simple_error!("too many envelops: {:?}, len: {}", cmd_simple_decrypt.key_id, filter_envelops.len());
    }
    let value = try_decrypt_key(&config, filter_envelops[0], &pin, &slot, false)?;
    let value = match output_type {
        "plain" => opt_result!(String::from_utf8(value), "bad value encoding: {}"),
        "hex" => hex::encode(&value),
        "base64" => STANDARD.encode(&value),
        _ => return simple_error!("not supported output type: {}", output_type),
    };
    CmdResult::success(&value).print_exit();
}