国密证书透明度日志的 Python 实现：基于 SM3 Merkle 树的审计与验证

"""CT 日志的 Merkle 树节点定义。"""
from __future__ import annotations
import hashlib
import struct
from dataclasses import dataclass, field
from typing import Optional

# CT 日志常量（RFC 6962 §3.1 定义）
LOG_ENTRY_TYPE_X509 = 0  # X.509 证书条目
LOG_ENTRY_TYPE_PRECERT = 1  # 预证书条目
MERKLE_LEAF_VERSION = 0x00  # 叶子版本号
MERKLE_NODE_VERSION = 0x01  # 内部节点版本号

# 国密 CT 日志用 SM3 替代 SHA-256
HASH_LENGTH = 32  # SM3 输出 32 字节，与 SHA-256 兼容
ZERO_HASH = b"\x00" * HASH_LENGTH


@dataclass
class MerkleNode:
    """Merkle 树节点。"""
    left: Optional[MerkleNode] = None
    right: Optional[MerkleNode] = None
    hash: bytes = field(default=ZERO_HASH)
    is_leaf: bool = False
    
    @classmethod
    def leaf_node(cls, data: bytes) -> MerkleNode:
        """创建叶子节点：H(0x00 || data)"""
        node = cls(hash=_sm3_hash(b"\x00" + data), is_leaf=True)
        return node
    
    @classmethod
    def internal_node(cls, left: MerkleNode, right: MerkleNode) -> MerkleNode:
        """创建内部节点：H(0x01 || left.hash || right.hash)"""
        combined = b"\x01" + left.hash + right.hash
        return cls(left=left, right=right, hash=_sm3_hash(combined))


def _sm3_hash(data: bytes) -> bytes:
    """
    使用 cryptography 的 SM3 实现计算哈希。
    替代 RFC 6962 中的 SHA-256。
    """
    from cryptography.hazmat.primitives import hashes
    digest = hashes.Hash(hashes.SM3())
    digest.update(data)
    return digest.finalize()

@dataclass
class CTLogEntry:
    """CT 日志条目结构（简化版）。"""
    entry_type: int  # 0=x509, 1=precert
    certificate_data: bytes  # DER 编码的证书（或预证书）
    timestamp: int  # Unix 毫秒时间戳
    
    def to_leaf_data(self) -> bytes:
        """序列化为叶子节点输入数据。"""
        # RFC 6962 §3.4: MerkleTreeLeaf 结构
        leaf = b""
        leaf += struct.pack(">B", MERKLE_LEAF_VERSION)  # version
        leaf += struct.pack(">B", self.entry_type)  # entry_type
        leaf += struct.pack(">Q", self.timestamp)  # timestamp
        # 证书数据（2字节长度前缀 + 数据）
        leaf += struct.pack(">H", len(self.certificate_data))
        leaf += self.certificate_data
        return leaf

class SM3MerkleTree:
    """
    基于 SM3 的 Merkle 树，支持：
    - 追加叶子节点
    - 计算树根哈希
    - 生成包含性证明（Inclusion Proof）
    - 生成一致性证明（Consistency Proof）
    - 验证包含性证明
    """
    
    def __init__(self):
        self.leaves: list[MerkleNode] = []
        self._root: Optional[MerkleNode] = None
        self._dirty: bool = False
    
    def add_leaf(self, data: bytes) -> int:
        """追加叶子节点，返回叶子索引。"""
        leaf = MerkleNode.leaf_node(data)
        self.leaves.append(leaf)
        self._dirty = True
        return len(self.leaves) - 1
    
    @property
    def root_hash(self) -> bytes:
        """计算 Merkle 树根哈希。"""
        if not self._dirty and self._root is not None:
            return self._root.hash
        if not self.leaves:
            return ZERO_HASH
        self._root = self._build_tree(list(self.leaves))
        self._dirty = False
        return self._root.hash
    
    def _build_tree(self, nodes: list[MerkleNode]) -> MerkleNode:
        """递归构建 Merkle 树。"""
        if len(nodes) == 1:
            return nodes[0]
        if len(nodes) % 2 == 1:
            # 奇数个节点，复制最后一个（RFC 6962 做法）
            nodes.append(nodes[-1])
        
        next_level = []
        for i in range(0, len(nodes), 2):
            parent = MerkleNode.internal_node(nodes[i], nodes[i + 1])
            next_level.append(parent)
        
        return self._build_tree(next_level)
    
    def get_inclusion_proof(self, leaf_index: int, tree_size: int) -> list[bytes]:
        """
        生成包含性证明。
        
        证明目标：验证某个叶子节点存在于树中的位置 leaf_index。
        返回：从该叶子到根路径上的兄弟节点哈希列表。
        
        Args:
            leaf_index: 叶子索引（0-based）
            tree_size: 证明时的树大小（如果 tree_size < 当前大小，使用快照）
        
        Returns:
            按从叶到根顺序的兄弟哈希列表
        """
        if leaf_index >= tree_size or leaf_index >= len(self.leaves):
            raise ValueError(f"Leaf index {leaf_index} out of range (size={tree_size})")
        
        # 只取前 tree_size 个叶子
        leaves_snapshot = self.leaves[:tree_size]
        return self._inclusion_proof_recursive(leaves_snapshot, leaf_index)
    
    def _inclusion_proof_recursive(
        self, nodes: list[MerkleNode], leaf_index: int
    ) -> list[bytes]:
        """递归生成包含性证明。"""
        if len(nodes) == 1:
            return []
        
        if len(nodes) % 2 == 1:
            nodes = list(nodes) + [nodes[-1]]
        
        half = len(nodes) // 2
        proof = []
        
        if leaf_index < half:
            # 叶子在左子树，右子树的根是兄弟
            right = self._build_tree(nodes[half:])
            proof.append(right.hash)
            proof.extend(self._inclusion_proof_recursive(nodes[:half], leaf_index))
        else:
            # 叶子在右子树，左子树的根是兄弟
            left = self._build_tree(nodes[:half])
            proof.append(left.hash)
            proof.extend(self._inclusion_proof_recursive(nodes[half:], leaf_index - half))
        
        return proof
    
    def verify_inclusion_proof(
        self, leaf_hash: bytes, leaf_index: int, tree_size: int, proof: list[bytes]
    ) -> bool:
        """
        验证包含性证明。
        
        Args:
            leaf_hash: 叶子节点的哈希值（用于验证）
            leaf_index: 叶子索引
            tree_size: 证明时的树大小
            proof: 包含性证明（兄弟哈希列表）
        
        Returns:
            验证是否通过
        """
        if leaf_index >= tree_size:
            return False
        
        current = leaf_hash
        index = leaf_index
        remaining_size = tree_size
        
        for sibling_hash in proof:
            if remaining_size == 1:
                # 只剩一个节点，当前就是根
                break
            
            is_right_leaf = (index % 2 == 1)
            
            if remaining_size % 2 == 1:
                # 当前层为奇数，最后一个节点是复制
                last_duplicate_index = remaining_size - 1
                if index == last_duplicate_index:
                    # 当前是复制节点，只计算内部到父节点
                    # 兄弟也是自身，已在下一轮处理
                    remaining_size = last_duplicate_index
                    # 不计算，直接跳过这个 index
                    index = index // 2
                    is_right_leaf = False
                    # 重新计算 current 的"兄弟"（即下一个证明）
                    if len(proof) > 1:
                        # 下一个 proof 节点才是真正需要的兄弟
                        # 实际上 RFC 6962 的算法更复杂，这里简化处理
                        pass
                    current = MerkleNode.internal_node(
                        MerkleNode(hash=current), MerkleNode(hash=current)
                    ).hash
                    continue
            
            if is_right_leaf:
                # 当前是右节点，兄弟在左
                combined = b"\x01" + sibling_hash + current
            else:
                # 当前是左节点，兄弟在右
                combined = b"\x01" + current + sibling_hash
            
            current = _sm3_hash(combined)
            index = index // 2
            remaining_size = (remaining_size + 1) // 2
        
        return current == self.root_hash
    
    def get_consistency_proof(self, old_size: int, new_size: int) -> list[bytes]:
        """
        生成一致性证明。
        
        证明目标：证明新树是旧树的扩展（没有修改旧条目）。
        返回：证明所需的节点哈希列表。
        
        Args:
            old_size: 先前的树大小
            new_size: 当前的树大小
        """
        if old_size > new_size or new_size > len(self.leaves):
            raise ValueError(f"Invalid sizes: old={old_size}, new={new_size}, total={len(self.leaves)}")
        if old_size == 0:
            return []
        if old_size == new_size:
            return []
        
        old_leaves = self.leaves[:old_size]
        new_leaves = self.leaves[:new_size]
        
        # RFC 6962 §2.1.2: consistency(old_root, new_root, proof)
        proof = []
        self._consistency_recursive(old_leaves, new_leaves, 0, proof)
        return proof
    
    def _consistency_recursive(
        self, old_nodes: list[MerkleNode], new_nodes: list[MerkleNode],
        depth: int, proof: list[bytes]
    ) -> None:
        """递归生成一致性证明。"""
        old_size = len(old_nodes)
        new_size = len(new_nodes)
        
        if old_size == new_size:
            # 大小相同，添加新树的根到证明
            new_root = self._build_tree(list(new_nodes))
            proof.append(new_root.hash)
            return
        
        if old_size == 1:
            # 旧树只有一个节点，新树有多个
            new_root = self._build_tree(list(new_nodes))
            proof.append(new_root.hash)
            return
        
        half_old = old_size // 2
        if old_size % 2 == 1:
            half_old = old_size  # 奇数时不分割
        
        if half_old >= old_size:
            # 无法再分割
            new_root = self._build_tree(list(new_nodes))
            proof.append(new_root.hash)
            return
        
        # 分割
        left_old = old_nodes[:half_old]
        right_old = old_nodes[half_old:]
        
        self._consistency_recursive(left_old, new_nodes[:half_old], depth + 1, proof)
        
        # 右子树贡献：右子树的根哈希
        right_new_root = self._build_tree(list(new_nodes[half_old:]))
        proof.append(right_new_root.hash)

def demonstrate_ct_log_inclusion():
    """演示 CT 日志的包含性证明流程。"""
    import time
    
    # 创建一个 CT 日志实例
    ct_log = SM3MerkleTree()
    
    # 模拟提交证书
    # 实际应用中 certificate_data 是 DER 编码的 X.509 证书
    certs = [
        b"\x30\x82\x01\x00" + bytes([i] * 200) for i in range(7)  # 模拟 7 个证书
    ]
    
    entries = []
    for i, cert_der in enumerate(certs):
        entry = CTLogEntry(
            entry_type=LOG_ENTRY_TYPE_X509,
            certificate_data=cert_der,
            timestamp=int(time.time() * 1000) + i * 1000,
        )
        leaf_data = entry.to_leaf_data()
        idx = ct_log.add_leaf(leaf_data)
        entries.append((idx, entry))
        print(f"[提交] 证书 #{idx}: 时间戳={entry.timestamp}")
    
    # 计算树根
    root = ct_log.root_hash
    print(f"\n树根哈希 (SM3): {root.hex()}")
    print(f"树大小: {len(ct_log.leaves)}")
    
    # 为证书 #3 生成包含性证明
    target_index = 3
    proof = ct_log.get_inclusion_proof(target_index, len(ct_log.leaves))
    
    print(f"\n--- 包含性证明 (证书 #{target_index}) ---")
    print(f"证明路径长度: {len(proof)}")
    for i, ph in enumerate(proof):
        print(f"  第 {i} 层兄弟: {ph.hex()[:32]}...")
    
    # 验证包含性证明
    leaf_node = ct_log.leaves[target_index]
    is_valid = ct_log.verify_inclusion_proof(
        leaf_node.hash, target_index, len(ct_log.leaves), proof
    )
    print(f"\n包含性证明验证: {'✅ 通过' if is_valid else '❌ 失败'}")
    
    # 篡改叶子数据后再次验证
    tampered_hash = _sm3_hash(b"tampered data")
    is_valid_tampered = ct_log.verify_inclusion_proof(
        tampered_hash, target_index, len(ct_log.leaves), proof
    )
    print(f"篡改后验证: {'✅ 通过' if is_valid_tampered else '❌ 失败'} (预期失败)")
    
    return ct_log


if __name__ == "__main__":
    demonstrate_ct_log_inclusion()

import json
from dataclasses import asdict

@dataclass
class SignedTreeHead:
    """签名树头（STH），RFC 6962 定义。"""
    tree_size: int  # 树大小
    timestamp: int  # 毫秒时间戳
    sha256_root_hash: bytes  # 树根哈希（国密版使用 SM3）
    signature: bytes  # 日志签名的 STH
    
    def to_signing_data(self) -> bytes:
        """生成待签名数据。"""
        # RFC 6962 §3.4: 版本(1) + 签名类型(1) + 时间(8) + tree_size(8) + root_hash(32)
        data = b""
        data += struct.pack(">B", 0)  # version = 0 (RFC 6962 v1)
        data += struct.pack(">B", 1)  # signature_type = 1 (tree_hash signature)
        data += struct.pack(">Q", self.timestamp)
        data += struct.pack(">Q", self.tree_size)
        data += self.sha256_root_hash  # 32 字节 SM3 根哈希
        return data


def sign_sth_with_sm2(
    ct_log: SM3MerkleTree, 
    sm2_private_key: str,
    sm2_public_key: str
) -> SignedTreeHead:
    """
    使用 SM2 私钥签名 STH。
    
    Args:
        ct_log: CT 日志实例（已添加条目）
        sm2_private_key: SM2 私钥（64 字符十六进制）
        sm2_public_key: SM2 公钥（130 字符十六进制，含 04 前缀）
    
    Returns:
        签名后的 STH
    """
    import time
    from gmssl import sm2
    
    timestamp = int(time.time() * 1000)
    tree_size = len(ct_log.leaves)
    root_hash = ct_log.root_hash
    
    # 构建 STH
    sth = SignedTreeHead(
        tree_size=tree_size,
        timestamp=timestamp,
        sha256_root_hash=root_hash,
        signature=b""
    )
    
    # 获取待签名数据
    signing_data = sth.to_signing_data()
    
    # SM2-SM3 签名：使用 sign_with_sm3（标准 API）
    sm2_crypt = sm2.CryptSM2(
        public_key=sm2_public_key, 
        private_key=sm2_private_key
    )
    
    # sign_with_sm3 自动完成：用户 ID → ZA 计算 → SM3 哈希 → SM2 签名
    # 等价于：SM2_SignWithZA(userID, data) = SM2_Sign(ZA || data)
    # 默认用户 ID 为 ASCII '1234567812345678'（GMSSL 默认值）
    signature = sm2_crypt.sign_with_sm3(signing_data)
    
    # 更新签名
    sth.signature = bytes.fromhex(signature) if isinstance(signature, str) else signature
    
    return sth


def verify_sth_with_sm2(
    sth: SignedTreeHead,
    sm2_public_key: str
) -> bool:
    """验证 SM2 签名的 STH。"""
    from gmssl import sm2
    
    # 重新计算签名前的数据
    data_to_verify = sth.to_signing_data()
    
    sm2_crypt = sm2.CryptSM2(public_key=sm2_public_key, private_key="")
    return sm2_crypt.verify_with_sm3(sth.signature, data_to_verify)

def demonstrate_consistency_proof():
    """演示 CT 日志的一致性证明。"""
    ct_log = SM3MerkleTree()
    
    # 第一批条目
    for i in range(5):
        entry = CTLogEntry(
            entry_type=LOG_ENTRY_TYPE_X509,
            certificate_data=bytes([i] * 200),
            timestamp=1_720_000_000_000 + i * 1000,
        )
        ct_log.add_leaf(entry.to_leaf_data())
    
    root_5 = ct_log.root_hash
    print(f"5 个条目时的树根: {root_5.hex()[:32]}...")
    
    # 追加更多条目
    for i in range(5, 8):
        entry = CTLogEntry(
            entry_type=LOG_ENTRY_TYPE_X509,
            certificate_data=bytes([i] * 200),
            timestamp=1_720_000_000_000 + i * 1000,
        )
        ct_log.add_leaf(entry.to_leaf_data())
    
    root_8 = ct_log.root_hash
    print(f"8 个条目时的树根: {root_8.hex()[:32]}...")
    
    # 生成从 size=5 到 size=8 的一致性证明
    proof = ct_log.get_consistency_proof(5, 8)
    print(f"\n一致性证明包含 {len(proof)} 个节点")
    for i, node in enumerate(proof):
        print(f"  节点 {i}: {node.hex()[:32]}...")
    
    # 验证一致性证明
    # 简化验证：用 proof 能计算出新的根
    # 完整实现需要对比 proof 生成的根与 sth.root_hash
    print("\n一致性证明生成完成 ✅")


if __name__ == "__main__":
    demonstrate_ct_log_inclusion()
    print("\n" + "=" * 60 + "\n")
    demonstrate_consistency_proof()