swift_secure_enclave_tool_rs/

lib.rs

use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use rust_util::{XResult, debugging, opt_result, simple_error, util_env};
use serde::{Deserialize, de};
use serde_json::Value;
use std::process::{Command, Output};

const ENV_SWIFT_SECURE_ENCLAVE_TOOL_CMD_KEY: &str = "SWIFT_SECURE_ENCLAVE_TOOL_CMD";
const DEFAULT_SWIFT_SECURE_ENCLAVE_TOOL_CMD: &str = "swift-secure-enclave-tool-v2";

#[derive(Debug, Clone, Copy)]
pub enum KeyPurpose {
    Signing,
    KeyAgreement,
}

#[derive(Debug, Clone, Copy)]
pub enum KeyMlKem {
    MlKem768,
    MlKem1024,
}

#[derive(Debug, Clone, Copy)]
pub enum ControlFlag {
    None,
    UserPresence,
    DevicePasscode,
    BiometryAny,
    BiometryCurrentSet,
}

impl ControlFlag {
    fn to_str(&self) -> &'static str {
        match self {
            ControlFlag::None => "none",
            ControlFlag::UserPresence => "userPresence",
            ControlFlag::DevicePasscode => "devicePasscode",
            ControlFlag::BiometryAny => "biometryAny",
            ControlFlag::BiometryCurrentSet => "biometryCurrentSet",
        }
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum DigestType {
    Raw,
    Sha256,
    Sha384,
    Sha512,
}

impl DigestType {
    pub fn to_str(&self) -> &'static str {
        match self {
            DigestType::Raw => "raw",
            DigestType::Sha256 => "sha256",
            DigestType::Sha384 => "sha384",
            DigestType::Sha512 => "sha512",
        }
    }

    pub fn bytes(&self) -> Option<u32> {
        match self {
            DigestType::Raw => None,
            DigestType::Sha256 => Some(256 / 8),
            DigestType::Sha384 => Some(384 / 8),
            DigestType::Sha512 => Some(512 / 8),
        }
    }

    pub fn parse(t: Option<&str>) -> XResult<DigestType> {
        Ok(match t {
            None | Some("") | Some("raw") => DigestType::Raw,
            Some("sha256") => DigestType::Sha256,
            Some("sha384") => DigestType::Sha384,
            Some("sha512") => DigestType::Sha512,
            Some(other_digest_type) => {
                return simple_error!("Invalid digest type: {}", other_digest_type);
            }
        })
    }
}

#[derive(Debug)]
pub struct KeyMaterial {
    pub public_key_point: Vec<u8>,
    pub public_key_der: Vec<u8>,
    pub private_key_representation: Vec<u8>,
}

#[derive(Debug, Deserialize)]
struct ErrorResult {
    #[allow(dead_code)]
    pub success: bool,
    pub error: String,
}

#[derive(Debug, Deserialize)]
struct IsSupportSecureEnclaveResult {
    #[allow(dead_code)]
    pub success: bool,
    pub supported: bool,
}

#[derive(Debug, Deserialize)]
struct KeyPairResult {
    #[allow(dead_code)]
    pub success: bool,
    pub public_key_point_base64: Option<String>,
    pub public_key_base64: String,
    pub data_representation_base64: String,
}

#[derive(Debug, Deserialize)]
struct SignResult {
    #[allow(dead_code)]
    pub success: bool,
    pub signature_base64: String,
}

#[derive(Debug, Deserialize)]
struct DhResult {
    #[allow(dead_code)]
    pub success: bool,
    pub shared_secret_hex: String,
}

#[derive(Debug, Deserialize)]
pub struct ExternalSpec {
    #[allow(dead_code)]
    pub success: bool,
    pub agent: String,
    pub specification: String,
}

#[derive(Debug, Deserialize)]
struct ExternalPublicKey {
    #[allow(dead_code)]
    pub success: bool,
    pub public_key_base64: String,
}

pub fn is_secure_enclave_supported() -> XResult<bool> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg("is_support_secure_enclave");

    let cmd_stdout = run_command_stdout(cmd)?;
    if is_success(&cmd_stdout)? {
        let is_support_se_result: IsSupportSecureEnclaveResult = from_str(&cmd_stdout)?;
        Ok(is_support_se_result.supported)
    } else {
        let error_result: ErrorResult = from_str(&cmd_stdout)?;
        simple_error!("{}", error_result.error)
    }
}

pub fn generate_keypair(
    key_purpose: KeyPurpose,
    control_flag: ControlFlag,
) -> XResult<KeyMaterial> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg(match key_purpose {
        KeyPurpose::Signing => "generate_p256_ecsign_keypair",
        KeyPurpose::KeyAgreement => "generate_p256_ecdh_keypair",
    });
    cmd.arg("--control-flag");
    cmd.arg(control_flag.to_str());

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_keypair_result(&cmd_stdout)
}

pub fn generate_mlkem_keypair(
    key_ml_kem: KeyMlKem,
    control_flag: ControlFlag,
) -> XResult<KeyMaterial> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg(match key_ml_kem {
        KeyMlKem::MlKem768 => "generate_mlkem768_ecdh_keypair",
        KeyMlKem::MlKem1024 => "generate_mlkem1024_ecdh_keypair",
    });
    cmd.arg("--control-flag");
    cmd.arg(control_flag.to_str());

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_keypair_result(&cmd_stdout)
}

pub fn recover_keypair(
    key_purpose: KeyPurpose,
    private_key_representation: &[u8],
) -> XResult<KeyMaterial> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg(match key_purpose {
        KeyPurpose::Signing => "recover_p256_ecsign_public_key",
        KeyPurpose::KeyAgreement => "recover_p256_ecdh_public_key",
    });
    cmd.arg("--private-key");
    cmd.arg(STANDARD.encode(private_key_representation));

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_keypair_result(&cmd_stdout)
}

pub fn recover_mlkem_keypair(
    key_ml_kem: KeyMlKem,
    private_key_representation: &[u8],
) -> XResult<KeyMaterial> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg(match key_ml_kem {
        KeyMlKem::MlKem768 => "recover_mlkem768_public_key",
        KeyMlKem::MlKem1024 => "recover_mlkem1024_public_key",
    });
    cmd.arg("--private-key");
    cmd.arg(STANDARD.encode(private_key_representation));

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_keypair_result(&cmd_stdout)
}

pub fn private_key_sign(private_key_representation: &[u8], content: &[u8]) -> XResult<Vec<u8>> {
    private_key_sign_digested(private_key_representation, content, DigestType::Raw)
}

pub fn private_key_sign_digested(
    private_key_representation: &[u8],
    content: &[u8],
    digest_type: DigestType,
) -> XResult<Vec<u8>> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg("compute_p256_ecsign");
    cmd.arg("--private-key");
    cmd.arg(STANDARD.encode(private_key_representation));
    cmd.arg("--message-base64");
    cmd.arg(STANDARD.encode(content));
    if digest_type != DigestType::Raw {
        cmd.arg("--message-type");
        cmd.arg(digest_type.to_str());
    }

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_sign_result(&cmd_stdout)
}

// ephemera_public_key MUST be DER format public key
pub fn private_key_ecdh(
    private_key_representation: &[u8],
    ephemera_public_key: &[u8],
) -> XResult<Vec<u8>> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg("compute_p256_ecdh");
    cmd.arg("--private-key");
    cmd.arg(STANDARD.encode(private_key_representation));
    cmd.arg("--ephemera-public-key");
    cmd.arg(STANDARD.encode(ephemera_public_key));

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_ecdh_result(&cmd_stdout)
}

pub fn private_key_mlkem_ecdh(
    key_ml_kem: KeyMlKem,
    private_key_representation: &[u8],
    ephemera_public_key: &[u8],
) -> XResult<Vec<u8>> {
    let mut cmd = Command::new(get_swift_secure_enclave_tool_cmd());
    cmd.arg(match key_ml_kem {
        KeyMlKem::MlKem768 => "compute_mlkem768_ecdh",
        KeyMlKem::MlKem1024 => "compute_mlkem1024_ecdh",
    });
    cmd.arg("--private-key");
    cmd.arg(STANDARD.encode(private_key_representation));
    cmd.arg("--ephemera-public-key");
    cmd.arg(STANDARD.encode(ephemera_public_key));

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_ecdh_result(&cmd_stdout)
}

pub fn external_sign(
    external_command: &str,
    parameter: &str,
    alg: &str,
    content: &[u8],
) -> XResult<Vec<u8>> {
    let mut cmd = Command::new(external_command);
    cmd.arg("external_sign");
    cmd.arg("--parameter");
    cmd.arg(parameter);
    cmd.arg("--alg");
    cmd.arg(alg);
    cmd.arg("--message-base64");
    cmd.arg(STANDARD.encode(content));

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_sign_result(&cmd_stdout)
}

pub fn external_ecdh(
    external_command: &str,
    parameter: &str,
    ephemera_public_key: &[u8],
) -> XResult<Vec<u8>> {
    let mut cmd = Command::new(external_command);
    cmd.arg("external_ecdh");
    cmd.arg("--parameter");
    cmd.arg(parameter);
    cmd.arg("--epk");
    cmd.arg(STANDARD.encode(ephemera_public_key));

    let cmd_stdout = run_command_stdout(cmd)?;
    parse_ecdh_result(&cmd_stdout)
}

pub fn external_spec(external_command: &str) -> XResult<ExternalSpec> {
    let mut cmd = Command::new(external_command);
    cmd.arg("external_spec");

    let cmd_stdout = run_command_stdout(cmd)?;
    if is_success(&cmd_stdout)? {
        let external_spec: ExternalSpec = from_str(&cmd_stdout)?;
        Ok(external_spec)
    } else {
        let error_result: ErrorResult = from_str(&cmd_stdout)?;
        simple_error!("{}", error_result.error)
    }
}

pub fn external_public_key(external_command: &str, parameter: &str) -> XResult<Vec<u8>> {
    let mut cmd = Command::new(external_command);
    cmd.arg("external_public_key");
    cmd.arg("--parameter");
    cmd.arg(parameter);

    let cmd_stdout = run_command_stdout(cmd)?;
    if is_success(&cmd_stdout)? {
        let external_public_key: ExternalPublicKey = from_str(&cmd_stdout)?;
        Ok(STANDARD.decode(&external_public_key.public_key_base64)?)
    } else {
        let error_result: ErrorResult = from_str(&cmd_stdout)?;
        simple_error!("{}", error_result.error)
    }
}

fn run_command_stdout(cmd: Command) -> XResult<String> {
    let output = run_command(cmd)?;
    let stdout_text = opt_result!(String::from_utf8(output.stdout), "Parse stdout failed:{}");
    Ok(stdout_text.trim().to_string())
}

fn run_command(mut cmd: Command) -> XResult<Output> {
    debugging!("Run command: {:?}", cmd);
    let output = cmd.output();
    match output {
        Err(e) => simple_error!("Run command failed: {:?}", e),
        Ok(output) => {
            debugging!("Output: {:?}", output);
            if !output.status.success() {
                let stderr = String::from_utf8_lossy(&output.stderr);
                let stdout = String::from_utf8_lossy(&output.stdout);
                simple_error!(
                    "Run command not success: {:?}\n - stdout: {}\n - stderr: {}",
                    output.status.code(),
                    stdout,
                    stderr
                )
            } else {
                Ok(output)
            }
        }
    }
}

fn parse_keypair_result(cmd_stdout: &str) -> XResult<KeyMaterial> {
    if is_success(&cmd_stdout)? {
        let key_pair_result: KeyPairResult = from_str(&cmd_stdout)?;
        let public_key_point = match &key_pair_result.public_key_point_base64 {
            Some(public_key_point_base64) => STANDARD.decode(public_key_point_base64)?,
            None => vec![],
        };
        let public_key_der = STANDARD.decode(&key_pair_result.public_key_base64)?;
        let private_key_representation =
            STANDARD.decode(&key_pair_result.data_representation_base64)?;
        Ok(KeyMaterial {
            public_key_point,
            public_key_der,
            private_key_representation,
        })
    } else {
        let error_result: ErrorResult = from_str(&cmd_stdout)?;
        simple_error!("{}", error_result.error)
    }
}

fn parse_sign_result(stdout: &str) -> XResult<Vec<u8>> {
    if is_success(stdout)? {
        let sign_result: SignResult = from_str(stdout)?;
        Ok(STANDARD.decode(&sign_result.signature_base64)?)
    } else {
        let error_result: ErrorResult = from_str(stdout)?;
        simple_error!("{}", error_result.error)
    }
}

fn parse_ecdh_result(stdout: &str) -> XResult<Vec<u8>> {
    if is_success(stdout)? {
        let dh_result: DhResult = from_str(stdout)?;
        Ok(hex::decode(&dh_result.shared_secret_hex)?)
    } else {
        let error_result: ErrorResult = from_str(stdout)?;
        simple_error!("{}", error_result.error)
    }
}

pub fn from_str<'a, T>(s: &'a str) -> XResult<T>
where
    T: de::Deserialize<'a>,
{
    match serde_json::from_str(s) {
        Ok(result) => Ok(result),
        Err(e) => simple_error!("Parse JSON: {}, error: {}", s, e),
    }
}

fn is_success(cmd_stdout: &str) -> XResult<bool> {
    let val = opt_result!(
        serde_json::from_str::<Value>(cmd_stdout),
        "Parse result: {}, failed: {}",
        cmd_stdout
    );
    if let Value::Object(map) = val {
        if let Some(success_value) = map.get("success") {
            if let Value::Bool(result) = success_value {
                return Ok(*result);
            }
        }
    }
    simple_error!("Bad result: {}", cmd_stdout)
}

fn get_swift_secure_enclave_tool_cmd() -> String {
    util_env::env_var(ENV_SWIFT_SECURE_ENCLAVE_TOOL_CMD_KEY)
        .unwrap_or_else(|| DEFAULT_SWIFT_SECURE_ENCLAVE_TOOL_CMD.to_string())
}