量子密碼學 - Leighton-Micali Signature 簽章

import hashlib
import math
import os

# --- 常數定義 ---
# 使用 SHA-256 作為我們的雜湊函數
HASH_FUNC = hashlib.sha256
# 每個雜湊輸出的位元組長度
N = HASH_FUNC().digest_size
# Winternitz 參數 (w)，這裡使用 16 (即 4-bit)
# w 越大，簽章越小，但計算量越大
W = 16

height = 8  # Merkle 樹的高度，這裡設定為 8 (即 2^8 = 256 個葉節點)

# --- WOTS+ 參數計算 ---
# 計算簽署一個 N 字節雜湊需要多少個鍊 (chain)
L1 = math.ceil((N * 8) / math.log2(W))
# 計算 checksum 需要多少個鍊
L2 = math.floor(math.log2(L1 * (W - 1)) / math.log2(W)) + 1
# WOTS+ 金鑰對中的鍊總數
L = L1 + L2


class WOTS:
    """
    一個簡化的 WOTS+ 一次性簽章實現。
    """

    @staticmethod
    def _chain(x, start, steps, pub_seed, addr):
        """對輸入 x 進行雜湊鍊計算"""
        if start + steps >= W:
            return None  # 步數超出範圍

        output = x
        for i in range(start, start + steps):
            # 在實際的 NIST 標準中，這裡會更複雜，會使用金鑰化的雜湊 (PRF)
            # 為了簡化，我們只將位址、種子和當前值串接後雜湊
            input_data = pub_seed + addr.to_bytes(4, 'big') + output
            output = HASH_FUNC(input_data).digest()
        return output

    @staticmethod
    def gen_keys(seed, pub_seed, addr):
        """
        生成一個 WOTS+ 金鑰對。
        在實際應用中，私鑰是從種子(seed)派生出來的，而不是儲存的。
        """
        # 私鑰是 L 個 N 字節的隨機字串
        private_key = [os.urandom(N) for _ in range(L)]

        # 公鑰是將私鑰的每個元素都執行 W-1 次雜湊鍊計算得到的
        public_key = [WOTS._chain(private_key[i], 0, W - 1, pub_seed, addr + i) for i in range(L)]

        # 將公鑰的所有部分雜湊成一個值，方便在 Merkle Tree 中使用
        pk_hash_input = b"".join(public_key)
        public_key_hash = HASH_FUNC(pk_hash_input).digest()

        return private_key, public_key, public_key_hash

    @staticmethod
    def sign(message, private_key, pub_seed, addr):
        """使用私鑰簽署一個訊息"""
        # 1. 將訊息雜湊
        msg_hash = HASH_FUNC(message).digest()

        # 2. 將雜湊值和 checksum 轉換為基於 W 的數字陣列
        base_w_msg = []
        checksum = 0

        # 處理訊息雜湊
        for i in range(L1):
            # 這裡用一個簡化的方式來切分雜湊值
            val = msg_hash[i // 2] if i % 2 == 0 else msg_hash[i // 2] >> 4
            val &= 0x0F  # 取 4 bits (因為 W=16)
            base_w_msg.append(val)
            checksum += (W - 1) - val

        # 處理 checksum
        checksum <<= (8 - ((L2 * math.floor(math.log2(W))) % 8)) % 8
        for i in range(L2):
            base_w_msg.append((checksum >> (i * int(math.log2(W)))) & (W - 1))

        # 3. 根據數字陣列，對私鑰的每個部分進行對應次數的鍊計算
        signature = [WOTS._chain(private_key[i], 0, base_w_msg[i], pub_seed, addr + i) for i in range(L)]
        return signature

    @staticmethod
    def pk_from_signature(signature, message, pub_seed, addr):
        """從簽章和訊息中，反推出公鑰"""
        # 這是在驗證時使用的關鍵步驟
        msg_hash = HASH_FUNC(message).digest()

        base_w_msg = []
        checksum = 0

        for i in range(L1):
            val = msg_hash[i // 2] if i % 2 == 0 else msg_hash[i // 2] >> 4
            val &= 0x0F
            base_w_msg.append(val)
            checksum += (W - 1) - val

        checksum <<= (8 - ((L2 * math.floor(math.log2(W))) % 8)) % 8
        for i in range(L2):
            base_w_msg.append((checksum >> (i * int(math.log2(W)))) & (W - 1))

        # 從簽章的每個部分，繼續計算鍊剩下的部分，以得到公鑰的每個部分
        recovered_pk = [WOTS._chain(signature[i], base_w_msg[i], (W - 1) - base_w_msg[i], pub_seed, addr + i) for i in
                        range(L)]

        pk_hash_input = b"".join(recovered_pk)
        return HASH_FUNC(pk_hash_input).digest()


class MerkleTree:
    """一個簡化的 Merkle Tree 實現"""

    def __init__(self, leaves):
        self.leaves = leaves
        self.tree = self._build_tree()

    def _build_tree(self):
        tree = [self.leaves]
        current_level = self.leaves
        while len(current_level) > 1:
            next_level = []
            for i in range(0, len(current_level), 2):
                left = current_level[i]
                # 如果是奇數個節點，最後一個節點就和自己雜湊
                right = current_level[i + 1] if i + 1 < len(current_level) else left
                parent = HASH_FUNC(left + right).digest()
                next_level.append(parent)
            tree.append(next_level)
            current_level = next_level
        return tree

    def get_root(self):
        return self.tree[-1][0]

    def get_authentication_path(self, index):
        path = []
        for level in range(len(self.tree) - 1):
            is_right_node = index % 2
            sibling_index = index - 1 if is_right_node else index + 1
            if sibling_index < len(self.tree[level]):
                path.append(self.tree[level][sibling_index])
            else:  # 如果沒有兄弟，就用自己
                path.append(self.tree[level][index])
            index //= 2
        return path

    @staticmethod
    def verify_path(root, leaf, path, index):
        """驗證一個葉節點和路徑是否能構成樹根"""
        node = leaf
        for i, sibling in enumerate(path):
            is_right_node = (index >> i) % 2
            if is_right_node:
                node = HASH_FUNC(sibling + node).digest()
            else:
                node = HASH_FUNC(node + sibling).digest()
        return node == root


class LMS:
    """
    LMS 簽章系統的主類別。
    """

    def __init__(self, height):
        self.height = height
        self.num_leaves = 1 << height

        # --- 狀態管理 ---
        # 這是 LMS 最關鍵的部分！
        # 在真實世界中，這個計數器必須被持久化且原子性地更新。
        self.q = 0  # 下一個可用的 OTS 金鑰索引

        print(f"初始化 LMS 系統，樹高 {height}，可用簽章數: {self.num_leaves}")

        # 整個 LMS 系統的唯一識別符
        self.I = os.urandom(N)
        self.pub_seed = os.urandom(N)

        # 生成所有底層的 WOTS+ 金鑰對
        print("正在生成所有 WOTS+ 金鑰對並建構 Merkle 樹...")
        self.wots_pks = []
        self.wots_sks = []
        wots_pk_hashes = []

        for i in range(self.num_leaves):
            # 每個 OTS 金鑰有唯一的位址
            addr = i
            sk, pk, pk_hash = WOTS.gen_keys(self.I, self.pub_seed, addr)
            self.wots_sks.append(sk)
            self.wots_pks.append(pk)
            wots_pk_hashes.append(pk_hash)

        # 建構 Merkle Tree
        self.merkle_tree = MerkleTree(wots_pk_hashes)
        self.public_key = self.merkle_tree.get_root()
        print("LMS 公鑰 (Merkle Root):", self.public_key.hex())
        print("-" * 20)

    def sign(self, message):
        """簽署一則訊息"""
        if self.q >= self.num_leaves:
            raise Exception("錯誤：沒有可用的簽章了！")

        # --- 關鍵的狀態更新 ---
        current_q = self.q
        self.q += 1
        # 在真實應用中，這裡必須有持久化儲存，確保 self.q 的值永不回退
        print(f"使用第 {current_q} 個 OTS 金鑰進行簽章 (下一個可用: {self.q})")

        # 1. 取得對應的 WOTS+ 私鑰並簽署
        ots_sk = self.wots_sks[current_q]
        ots_signature = WOTS.sign(message, ots_sk, self.pub_seed, current_q)

        # 2. 取得驗證路徑
        auth_path = self.merkle_tree.get_authentication_path(current_q)

        # 3. 組合 LMS 簽章
        lms_signature = {
            "q": current_q,
            "ots_signature": ots_signature,
            "auth_path": auth_path
        }
        return lms_signature

    def verify(self, message, signature):
        """驗證一個 LMS 簽章"""
        q = signature["q"]
        ots_signature = signature["ots_signature"]
        auth_path = signature["auth_path"]

        # 步驟 1: 從 OTS 簽章中反推出 WOTS+ 公鑰的雜湊值
        recovered_wots_pk_hash = WOTS.pk_from_signature(ots_signature, message, self.pub_seed, q)

        # 步驟 2: 使用驗證路徑，驗證這個 WOTS+ 公鑰的雜湊值是否真的在 Merkle Tree 中
        is_valid = MerkleTree.verify_path(self.public_key, recovered_wots_pk_hash, auth_path, q)

        return is_valid


# --- 主執行區塊 ---
if __name__ == "__main__":
    # 建立一個高度為 4 的 LMS 系統 (總共 2^4 = 16 個可用簽章)
    lms_system = LMS(height=height)

    # 要簽署的訊息
    my_message = b"This is a very important document."

    # --- 第一次簽章 ---
    print("\n[測試 1: 正常簽章與驗證]")
    signature1 = lms_system.sign(my_message)
    is_verified = lms_system.verify(my_message, signature1)
    print(f"訊息: {my_message.decode()}")
    print(f"驗證結果: {'成功' if is_verified else '失敗'}")

    # --- 第二次簽章 (使用下一個金鑰) ---
    print("\n[測試 2: 簽署另一則訊息]")
    my_second_message = b"This is another document."
    signature2 = lms_system.sign(my_second_message)
    is_verified_2 = lms_system.verify(my_second_message, signature2)
    print(f"訊息: {my_second_message.decode()}")
    print(f"驗證結果: {'成功' if is_verified_2 else '失敗'}")

    # --- 失敗測試 ---
    print("\n[測試 3: 竄改訊息後驗證]")
    tampered_message = b"This is a NOT important document."
    is_verified_tampered = lms_system.verify(tampered_message, signature1)
    print(f"訊息: {tampered_message.decode()}")
    print(f"驗證結果 (使用原始簽章): {'成功' if is_verified_tampered else '失敗'}")

    print("\n狀態檢查:")
    print(f"當前 LMS 系統已使用 {lms_system.q} 個簽章。")