script-sign/src/lib.rs

mod keymap;
mod sign;
mod signature;
mod util;

pub use crate::keymap::KeyMap;
use crate::sign::{ecdsaverify, EcdsaAlgorithm};
use crate::signature::{
    CardEcSignResult, ScriptSignature, ScriptSignatureAlgorithm, ScriptSignatureVersion,
    SIGNATURE_PREFIX,
};
use crate::util::current_time;
use digest::Digest;
use rust_util::{debugging, opt_result, simple_error, util_cmd, XResult};
use sha2::Sha256;
use std::fs;

#[derive(Debug)]
pub struct Script {
    pub content_lines: Vec<String>,
    pub signature: Option<ScriptSignature>,
}

impl Script {
    pub fn verify_script_file_with_system_key_map(script_file: &str) -> XResult<bool> {
        Self::verify_script_file(script_file, &KeyMap::system())
    }

    pub fn verify_script_file(script_file: &str, key_map: &KeyMap) -> XResult<bool> {
        let script_content = opt_result!(
            fs::read_to_string(script_file),
            "Read script file: {script_file} failed: {}"
        );
        let script = opt_result!(
            Script::parse(&script_content),
            "Parse script file: {script_file} failed: {}"
        );
        match &script.signature {
            None => Ok(false),
            Some(_) => script.verify(key_map),
        }
    }

    pub fn parse(script: &str) -> XResult<Script> {
        let lines = script.lines().collect::<Vec<_>>();

        let mut in_signature_section = false;
        let mut signature_lines = vec![];
        let mut content_lines = vec![];
        let mut signature_line = String::new();

        let mut push_signature_line = |signature_line: &mut String| {
            if !signature_line.is_empty() {
                signature_lines.push(signature_line.clone());
                signature_line.clear();
            }
        };

        for line in &lines {
            if in_signature_section {
                if line.starts_with(SIGNATURE_PREFIX) {
                    push_signature_line(&mut signature_line);
                    signature_line.push_str(line);
                } else if line.starts_with("//") {
                    signature_line.push_str(line.chars().skip(2).collect::<String>().trim());
                } else if !line.trim().is_empty() {
                    return simple_error!("Bad signature section line, find: '{line}'");
                }
            } else {
                if line.starts_with(SIGNATURE_PREFIX) {
                    in_signature_section = true;
                    push_signature_line(&mut signature_line);
                    signature_line.push_str(line);
                } else {
                    content_lines.push(line.to_string());
                }
            }
        }
        push_signature_line(&mut signature_line);

        if signature_lines.len() > 1 {
            return simple_error!(
                "Found {} signatures, only supports one signature.",
                signature_lines.len()
            );
        }

        if signature_lines.is_empty() {
            Ok(Script {
                content_lines,
                signature: None,
            })
        } else {
            let script_signature = ScriptSignature::parse(&signature_lines[0])?;
            Ok(Script {
                content_lines,
                signature: Some(script_signature),
            })
        }
    }

    pub fn as_string(&self) -> String {
        let mut joined_content_liens = self.content_lines.join("\n");
        match &self.signature {
            None => joined_content_liens,
            Some(signature) => {
                if joined_content_liens.ends_with("\n\n") {
                    // SKIP add \n
                } else if joined_content_liens.ends_with("\n") {
                    joined_content_liens.push('\n');
                } else {
                    joined_content_liens.push_str("\n\n");
                }
                let signature_lines = signature.as_string_lines_default_width();
                for signature_line in &signature_lines {
                    joined_content_liens.push_str(&signature_line);
                    joined_content_liens.push('\n');
                }
                // joined_content_liens.push_str(&signature.as_string());
                // joined_content_liens.push('\n');
                joined_content_liens
            }
        }
    }

    pub fn has_signature(&self) -> bool {
        self.signature.is_some()
    }

    pub fn verify(&self, key_map: &KeyMap) -> XResult<bool> {
        let signature = match &self.signature {
            None => return simple_error!("Script is not signed."),
            Some(signature) => signature,
        };
        let key = match key_map.find(&signature.key_id) {
            None => return simple_error!("Sign key id: {} not found", &signature.key_id),
            Some(key) => key,
        };

        let mut verify_public_key = key.public_key_point_hex.clone();
        if ScriptSignatureVersion::V2 == signature.ver {
            match &signature.embed_signing_key {
                Some(embed_signing_key) => {
                    let mut hasher = Sha256::new();
                    hasher.update(embed_signing_key.time.as_bytes());
                    hasher.update(&embed_signing_key.public_key);
                    let embed_digest_sha256 = hasher.finalize().to_vec();
                    let key_bytes = hex::decode(&key.public_key_point_hex)?;
                    match embed_signing_key.algorithm {
                        ScriptSignatureAlgorithm::ES256 => {
                            match ecdsaverify(
                                EcdsaAlgorithm::P256,
                                &key_bytes,
                                &embed_digest_sha256,
                                &embed_signing_key.signature,
                            ) {
                                Ok(_) => {
                                    verify_public_key = hex::encode(&embed_signing_key.public_key);
                                }
                                Err(e) => {
                                    debugging!("Verify embed ecdsa signature failed: {}", e);
                                    return Ok(false);
                                }
                            }
                        }
                        _ => {
                            return simple_error!(
                                "Not supported algorithm: {:?}",
                                signature.algorithm
                            )
                        }
                    }
                }
                None => {
                    return simple_error!("Embed signing key not found");
                }
            }
        }

        let key_bytes = hex::decode(&verify_public_key)?;
        let digest_sha256 = self.normalize_content_lines_and_sha256(&signature.time);
        match signature.algorithm {
            ScriptSignatureAlgorithm::ES256 => {
                match ecdsaverify(
                    EcdsaAlgorithm::P256,
                    &key_bytes,
                    &digest_sha256,
                    &signature.signature,
                ) {
                    Ok(_) => Ok(true),
                    Err(e) => {
                        debugging!("Verify ecdsa signature failed: {}", e);
                        Ok(false)
                    }
                }
            }
            _ => simple_error!("Not supported algorithm: {:?}", signature.algorithm),
        }
    }

    pub fn sign(&mut self) -> XResult<()> {
        let (time, digest_sha256) = self.normalize_content_lines_and_sha256_with_current_time();
        let digest_sha256_hex = hex::encode(&digest_sha256);
        let output = util_cmd::run_command_or_exit(
            "card-cli",
            &["piv-ecsign", "--json", "-s", "r1", "-x", &digest_sha256_hex],
        );
        let ecsign_result: CardEcSignResult = opt_result!(
            serde_json::from_slice(&output.stdout),
            "Parse card piv-ecsign failed: {}"
        );
        if ecsign_result.algorithm == "ecdsa_p256_with_sha256" {
            self.signature = Some(ScriptSignature {
                ver: ScriptSignatureVersion::V1,
                key_id: "yk-r1".to_string(),
                embed_signing_key: None,
                algorithm: ScriptSignatureAlgorithm::ES256,
                time,
                signature: hex::decode(&ecsign_result.signed_data_hex)?,
            });
        } else {
            return simple_error!("Not supported algorithm: {}", ecsign_result.algorithm);
        }
        Ok(())
    }

    fn normalize_content_lines(&self) -> Vec<String> {
        let mut normalized_content_lines = Vec::with_capacity(self.content_lines.len());
        for ln in &self.content_lines {
            let trimed_ln = ln.trim();
            if !trimed_ln.is_empty() {
                normalized_content_lines.push(trimed_ln.to_string());
            }
        }
        normalized_content_lines
    }

    fn normalize_content_lines_and_sha256_with_current_time(&self) -> (String, Vec<u8>) {
        let current_time = current_time();
        (
            current_time.clone(),
            self.normalize_content_lines_and_sha256(&current_time),
        )
    }

    fn normalize_content_lines_and_sha256(&self, current_time: &str) -> Vec<u8> {
        let normalized_content_lines = self.normalize_content_lines();
        let joined_normalized_content_lines = normalized_content_lines.join("\n");
        let mut hasher = Sha256::new();
        hasher.update(current_time.as_bytes());
        hasher.update(joined_normalized_content_lines.as_bytes());
        hasher.finalize().to_vec()
    }
}

#[test]
fn test_script_parse_01() {
    let script = Script::parse("test script").unwrap();
    assert_eq!(1, script.content_lines.len());
    assert!(script.signature.is_none());
}

#[test]
fn test_script_parse_02() {
    use base64::engine::general_purpose::STANDARD as standard_base64;
    use base64::Engine;
    let script =
        Script::parse("test script\n\n// @SCRIPT-SIGNATURE-V1: key-id.RS256.2025-01-05T20:57:14+08:00.aGVsbG93b3JsZA==\n")
            .unwrap();
    assert_eq!(2, script.content_lines.len());
    assert_eq!("test script", script.content_lines[0]);
    assert_eq!("", script.content_lines[1]);
    let current_time = "2025-01-05T20:57:14+08:00";
    let digest_sha256 = script.normalize_content_lines_and_sha256(&current_time);
    assert_eq!(
        "sybQ8O5TgRlkQ0i8pNIA6huHvAd5XbVZF+U60WMrdco=",
        standard_base64.encode(&digest_sha256)
    );
    assert!(script.signature.is_some());
    let s = script.signature.unwrap();
    assert_eq!("key-id", s.key_id);
    assert_eq!(ScriptSignatureAlgorithm::RS256, s.algorithm);
    assert_eq!("2025-01-05T20:57:14+08:00", s.time);
    assert_eq!(b"helloworld".to_vec(), s.signature);
}

#[test]
fn test_script_parse_03() {
    let script =
        Script::parse(r##"#!/usr/bin/env -S deno run --allow-env

console.log("Hello world.");

// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20250122T233410+08:00.MEQCIGogDudoVpCVfGiNPu8Wn6YPDtFX5OXC4bKtsN1nw414AiAq+5EVdvOuKAlXdVeeE1d91mKX9TaSTR25jliUx0km6A=="##)
            .unwrap();
    assert!(script.verify(&KeyMap::system()).unwrap());
}

#[test]
fn test_script_parse_04() {
    let script =
        Script::parse(r##"#!/usr/bin/env -S deno run --allow-env

console.log("Hello world.");

// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20250122T233410+08:00.MEQCIGogDudoVpCVfGiNPu8Wn6YPDtFX5OXC4bKtsN1nw414AiAq+5EVdvOuKAlXdVeeE1d91mKX9TaSTR25jliUx0km6A=="##)
            .unwrap();
    let script_str = script.as_string();
    assert_eq!(
        r##"#!/usr/bin/env -S deno run --allow-env

console.log("Hello world.");

// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20250122T233410+08:00.MEQCIGogDudoVpCVfGiNPu8W
// n6YPDtFX5OXC4bKtsN1nw414AiAq+5EVdvOuKAlXdVeeE1d91mKX9TaSTR25jliUx0km6A==
"##,
        script_str
    );
}

#[test]
fn test_script_parse_05() {
    let script = Script::parse(
        r##"#!/usr/bin/env -S deno run --allow-env

console.log("Hello world.");

// @SCRIPT-SIGNATURE-V1: yk-r1.ES256.20250122T233410+08:00.MEQCIGogDudoVpCVfGiNPu
// 8Wn6YPDtFX5OXC4bKtsN1nw414AiAq+5EVdvOuKAlXdVeeE1d91mKX9TaSTR25jliUx0km6A=="##,
    )
    .unwrap();
    assert!(script.verify(&KeyMap::system()).unwrap());
}