From ebcf7b83d62e0b3be8c1eace162049820b121a10 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Wed, 22 Jan 2025 01:40:00 +0800
Subject: [PATCH] feat: init commit

---
 .gitignore       |   1 +
 Cargo.toml       |  22 +++++
 src/keymap.rs    |  25 ++++++
 src/lib.rs       | 205 +++++++++++++++++++++++++++++++++++++++++++++++
 src/sign.rs      |  51 ++++++++++++
 src/signature.rs | 108 +++++++++++++++++++++++++
 src/util.rs      |   5 ++
 7 files changed, 417 insertions(+)
 create mode 100644 Cargo.toml
 create mode 100644 src/keymap.rs
 create mode 100644 src/lib.rs
 create mode 100644 src/sign.rs
 create mode 100644 src/signature.rs
 create mode 100644 src/util.rs

diff --git a/.gitignore b/.gitignore
index 3bf25c0..409abaa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.idea/
 # ---> Rust
 # Generated by Cargo
 # will have compiled files and executables
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..01c862f
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,22 @@
+[package]
+name = "script-sign"
+version = "0.1.0"
+edition = "2021"
+authors = ["Hatter Jiang <jht5945@gmail.com>"]
+description = "Script Sign"
+license = "MIT"
+readme = "README.md"
+
+[dependencies]
+base64 = "0.22"
+digest = "0.10"
+ecdsa = "0.16"
+hex = "0.4"
+p256 = "0.13"
+p384 = "0.13"
+sha2 = "0.10"
+regex = "1.11"
+rust_util = "0.6"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+simpledateformat = "0.1"
diff --git a/src/keymap.rs b/src/keymap.rs
new file mode 100644
index 0000000..f5a3139
--- /dev/null
+++ b/src/keymap.rs
@@ -0,0 +1,25 @@
+use rust_util::XResult;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct KeyMap {
+    key_map: HashMap<String, String>,
+}
+
+impl KeyMap {
+    pub fn system() -> XResult<Self> {
+        let signing_keys = r##"
+{
+  "yk-r1": "04dd3eebd906c9cf00b08ec29f7ed61804d1cc1d1352d9257b628191e08fc3717c4fae3298cd5c4829cec8bf3a946e7db60b7857e1287f6a0bae6b3f2342f007d0"
+}
+"##;
+        // unwrap should not happen
+        let key_map: HashMap<String, String> = serde_json::from_str(signing_keys).unwrap();
+        Ok(KeyMap { key_map })
+    }
+
+    pub fn find(&self, key_id: &str) -> Option<&String> {
+        self.key_map.get(key_id)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..1ea323d
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,205 @@
+mod keymap;
+mod sign;
+mod signature;
+mod util;
+
+pub use crate::keymap::KeyMap;
+use crate::sign::{ecdsaverify, EcdsaAlgorithm};
+use crate::signature::{
+    CardEcSignResult, ScriptSignature, ScriptSignatureAlgorithm, 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 last_non_empty_line = lines.iter().rev().find(|ln| !ln.is_empty());
+        match last_non_empty_line {
+            Some(last_non_empty_line) if last_non_empty_line.starts_with(SIGNATURE_PREFIX) => {
+                let script_signature = ScriptSignature::parse(last_non_empty_line)?;
+                let final_lines = lines
+                    .iter()
+                    .rev()
+                    .skip_while(|ln| ln.is_empty())
+                    .skip(1)
+                    .collect::<Vec<_>>()
+                    .into_iter()
+                    .rev()
+                    .collect::<Vec<_>>();
+                Ok(Script {
+                    content_lines: final_lines.iter().map(ToString::to_string).collect(),
+                    signature: Some(script_signature),
+                })
+            }
+            _ => Ok(Script {
+                content_lines: lines.iter().map(ToString::to_string).collect(),
+                signature: None,
+            }),
+        }
+    }
+
+    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");
+                }
+                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 key_bytes = hex::decode(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 {
+                key_id: "yk-r1".to_string(),
+                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);
+}
diff --git a/src/sign.rs b/src/sign.rs
new file mode 100644
index 0000000..f9978c3
--- /dev/null
+++ b/src/sign.rs
@@ -0,0 +1,51 @@
+use ecdsa::signature::hazmat::PrehashVerifier;
+use ecdsa::Signature;
+use ecdsa::VerifyingKey;
+use p256::NistP256;
+use p384::NistP384;
+use rust_util::opt_result;
+use rust_util::{simple_error, XResult};
+
+macro_rules! ecdsa_verify_signature {
+    ($algo: tt, $pk_point: tt, $prehash: tt, $signature: tt) => {{
+        let verifying_key: VerifyingKey<$algo> = opt_result!(
+            VerifyingKey::<$algo>::from_sec1_bytes($pk_point),
+            "Parse public key failed: {}"
+        );
+        let sign = if let Ok(signature) = Signature::from_der($signature) {
+            signature
+        } else if let Ok(signature) = Signature::from_slice($signature) {
+            signature
+        } else {
+            return simple_error!("Parse signature failed: {}", hex::encode($signature));
+        };
+        opt_result!(
+            verifying_key.verify_prehash($prehash, &sign),
+            "Verify signature failed: {}"
+        );
+    }};
+}
+
+#[allow(dead_code)]
+#[derive(Copy, Clone)]
+pub enum EcdsaAlgorithm {
+    P256,
+    P384,
+}
+
+pub fn ecdsaverify(
+    algo: EcdsaAlgorithm,
+    pk_point: &[u8],
+    prehash: &[u8],
+    signature: &[u8],
+) -> XResult<()> {
+    match algo {
+        EcdsaAlgorithm::P256 => {
+            ecdsa_verify_signature!(NistP256, pk_point, prehash, signature)
+        }
+        EcdsaAlgorithm::P384 => {
+            ecdsa_verify_signature!(NistP384, pk_point, prehash, signature)
+        }
+    }
+    Ok(())
+}
diff --git a/src/signature.rs b/src/signature.rs
new file mode 100644
index 0000000..f20b639
--- /dev/null
+++ b/src/signature.rs
@@ -0,0 +1,108 @@
+use base64::engine::general_purpose::STANDARD as standard_base64;
+use base64::Engine;
+use regex::Regex;
+use rust_util::{opt_result, opt_value_result, simple_error, XResult};
+use serde::Deserialize;
+
+pub const SIGNATURE_PREFIX: &str = "// @SCRIPT-SIGNATURE-";
+pub const SIGNATURE_V1: &str = "V1";
+
+#[derive(Debug, Eq, PartialEq)]
+pub enum ScriptSignatureAlgorithm {
+    RS256,
+    ES256,
+    ES384,
+    ES521,
+}
+
+impl ScriptSignatureAlgorithm {
+    pub fn try_from(algo: &str) -> XResult<Self> {
+        let upper_algo = algo.to_uppercase();
+        Ok(match upper_algo.as_str() {
+            "RS256" => Self::RS256,
+            "ES256" => Self::ES256,
+            "ES384" => Self::ES384,
+            "ES521" => Self::ES521,
+            _ => return simple_error!("Not valid algorithm: {}", algo),
+        })
+    }
+
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            ScriptSignatureAlgorithm::RS256 => "RS256",
+            ScriptSignatureAlgorithm::ES256 => "ES256",
+            ScriptSignatureAlgorithm::ES384 => "ES384",
+            ScriptSignatureAlgorithm::ES521 => "ES521",
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct ScriptSignature {
+    pub key_id: String,
+    pub algorithm: ScriptSignatureAlgorithm,
+    pub time: String,
+    pub signature: Vec<u8>,
+}
+
+#[allow(dead_code)]
+#[derive(Debug, Deserialize)]
+pub struct CardEcSignResult {
+    pub algorithm: String,
+    pub hash_hex: String,
+    pub signed_data_base64: String,
+    pub signed_data_hex: String,
+    pub slot: String,
+}
+
+impl ScriptSignature {
+    pub fn parse(script_signature_line: &str) -> XResult<ScriptSignature> {
+        // e.g. // @SCRIPT-SIGNATURE-V1: <key-id>.<algotirhm>.<time>.<signature-value-in-base64>
+        let script_signature_v1_regex = Regex::new(
+            r##"^\s*//\s*@SCRIPT-SIGNATURE-V1:\s*([a-zA-Z0-9-_]+)\.([a-zA-Z0-9]+)\.([0-9a-zA-Z\-+:]+)\.([a-zA-Z0-9\-+_/=]+)\s*$"##).unwrap();
+        let script_signature_v1_captures = opt_value_result!(
+            script_signature_v1_regex.captures(script_signature_line),
+            "Parse script signature failed: {}",
+            script_signature_line
+        );
+        let (_, [key_id, algorithm, time, signature]) = script_signature_v1_captures.extract();
+
+        let signature = opt_result!(
+            standard_base64.decode(signature),
+            "Parse script signature failed, decode signature failed: {}"
+        );
+        Ok(ScriptSignature {
+            key_id: key_id.to_string(),
+            algorithm: ScriptSignatureAlgorithm::try_from(algorithm)?,
+            time: time.to_string(),
+            signature,
+        })
+    }
+
+    pub fn as_string(&self) -> String {
+        let mut s = String::with_capacity(245);
+        s.push_str(SIGNATURE_PREFIX);
+        s.push_str(SIGNATURE_V1);
+        s.push_str(": ");
+        s.push_str(&self.key_id);
+        s.push('.');
+        s.push_str(self.algorithm.as_str());
+        s.push('.');
+        s.push_str(&self.time);
+        s.push('.');
+        s.push_str(&standard_base64.encode(&self.signature));
+        s
+    }
+}
+
+#[test]
+fn test_script_signature_parse() {
+    let v1 = "// @SCRIPT-SIGNATURE-V1: key-id.RS256.2025-01-05T20:57:14+08:00.aGVsbG93b3JsZA==";
+    let s = ScriptSignature::parse(v1).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);
+
+    assert_eq!(v1, s.as_string());
+}
diff --git a/src/util.rs b/src/util.rs
new file mode 100644
index 0000000..ae40aa5
--- /dev/null
+++ b/src/util.rs
@@ -0,0 +1,5 @@
+pub fn current_time() -> String {
+    simpledateformat::fmt("yyyyMMdd'T'HHmmssz")
+        .unwrap()
+        .format_local_now()
+}