国密 API 网关级签名认证：防重放攻击与密钥轮转生产方案

"""
国密 API 网关认证中间件
依赖：fastapi, redis, cryptography, gmssl
"""
import time
import uuid
import hashlib
import hmac as stdlib_hmac
from datetime import datetime, timedelta
from typing import Optional
from dataclasses import dataclass
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.hmac import HMAC as CryHMAC
from gmssl import sm3, func


@dataclass
class SignatureConfig:
    """签名配置"""
    secret_key: bytes           # HMAC 密钥（32字节）
    timestamp_window: int = 300   # 时间窗口（秒），默认 5 分钟
    nonce_ttl: int = 600        # Nonce 缓存 TTL（秒），默认 10 分钟


class GmGatewayAuth:
    """国密 API 网关认证处理器"""
    
    def __init__(self, config: SignatureConfig, redis_client):
        self.config = config
        self.redis = redis_client
    
    def compute_signature(
        self,
        method: str,
        path: str,
        query: dict,
        body: bytes,
        timestamp: int,
        nonce: str
    ) -> str:
        """
        计算 HMAC-SM3 签名
        
        签名串格式（严格按此顺序，避免签名不一致）：
        METHOD\\nPATH\\nSORTED_QUERY\\nTIMESTAMP\\nNONCE\\nBODY_HASH
        
        参考：GM/T 0054-2018 第 5.3 条"应用层签名要求"
        """
        # 1. 对查询参数排序后拼接
        sorted_query = '&'.join(
            f'{k}={v}' for k, v in sorted(query.items())
        ) if query else ''
        
        # 2. 计算请求体哈希（空体时为固定值）
        body_hash = (
            sm3.sm3_hash(func.bytes_to_list(body))
            if body else 'e3b0c44298fc1c149afbf4c8996fb924'
        )
        
        # 3. 构建待签名字符串
        sign_str = '\n'.join([
            method.upper(),
            path,
            sorted_query,
            str(timestamp),
            nonce,
            body_hash
        ])
        
        # 4. 计算 HMAC-SM3
        h = CryHMAC(self.config.secret_key, hashes.SM3())
        h.update(sign_str.encode('utf-8'))
        return h.finalize().hex()
    
    def verify_signature(
        self,
        method: str,
        path: str,
        query: dict,
        body: bytes,
        timestamp: int,
        nonce: str,
        signature: str
    ) -> tuple[bool, str]:
        """
        验证签名 + 时间戳 + Nonce 三重检查
        返回：(是否有效, 错误原因)
        """
        # === 第一层：时间戳有效性 ===
        now = int(time.time())
        if abs(now - timestamp) > self.config.timestamp_window:
            return False, f"请求已过期：时间戳差 {abs(now - timestamp)}s > {self.config.timestamp_window}s"
        
        # === 第二层：Nonce 唯一性（防重放）===
        nonce_key = f"gm:nonce:{self._safe_app_id(method, path)}:{nonce}"
        if self.redis.exists(nonce_key):
            return False, "请求已处理（Nonce 重复）"
        
        # 设置 Nonce 缓存（原子操作）
        self.redis.setex(nonce_key, self.config.nonce_ttl, '1')
        
        # === 第三层：签名验证 ===
        expected = self.compute_signature(
            method, path, query, body, timestamp, nonce
        )
        
        # 恒定时间比较，防时序攻击
        if not stdlib_hmac.compare_digest(expected, signature):
            return False, "签名验证失败"
        
        return True, "OK"
    
    def _safe_app_id(self, method: str, path: str) -> str:
        """生成安全的缓存键前缀"""
        return hashlib.md5(f"{method}:{path}".encode()).hexdigest()[:8]


# ========== 使用示例 ==========
if __name__ == '__main__':
    import redis
    
    # 初始化（生产环境从 HSM/密钥管理系统读取）
    config = SignatureConfig(
        secret_key=os.urandom(32),  # 32 字节符合 GM/T 0054 要求
        timestamp_window=300,
        nonce_ttl=600
    )
    redis_client = redis.Redis(host='localhost', port=6379, db=0)
    
    auth = GmGatewayAuth(config, redis_client)
    
    # 模拟客户端请求
    method = 'POST'
    path = '/api/v1/transfer'
    query = {'page': '1', 'size': '10'}
    body = b'{"to":"ACC-999","amount":10000}'
    timestamp = int(time.time())
    nonce = uuid.uuid4().hex
    
    # 计算签名
    signature = auth.compute_signature(
        method, path, query, body, timestamp, nonce
    )
    print(f"签名: {signature[:32]}...")
    
    # 验证签名（同一请求）
    ok, reason = auth.verify_signature(
        method, path, query, body, timestamp, nonce, signature
    )
    print(f"第一次验证: {ok}, {reason}")
    
    # 重放攻击检测
    ok2, reason2 = auth.verify_signature(
        method, path, query, body, timestamp, nonce, signature
    )
    print(f"重放检测: {ok2}, {reason2}")
    
    # 参数篡改检测
    bad_body = b'{"to":"ACC-999","amount":100000}'
    ok3, reason3 = auth.verify_signature(
        method, path, query, bad_body, timestamp, nonce, signature
    )
    print(f"篡改检测: {ok3}, {reason3}")
    
    # 时间戳过期检测
    old_timestamp = timestamp - 600  # 10 分钟前
    ok4, reason4 = auth.verify_signature(
        method, path, query, body, old_timestamp, nonce, signature
    )
    print(f"过期检测: {ok4}, {reason4}")

"""
密钥版本管理：支持多版本密钥同时验证
"""
import json
from datetime import datetime
from typing import Optional


class KeyRotationManager:
    """密钥轮转管理器"""
    
    def __init__(self):
        # 密钥版本库：version -> SecretKeyConfig
        self._keys: dict[str, dict] = {}
        self._active_version: Optional[str] = None
    
    def add_key_version(
        self,
        version: str,
        secret_key: bytes,
        created_at: datetime,
        expires_at: datetime,
       备注: str = ""
    ) -> None:
        """添加密钥版本"""
        self._keys[version] = {
            'secret_key': secret_key,
            'created_at': created_at.isoformat(),
            'expires_at': expires_at.isoformat(),
            '备注': 备注,
            'active': version == self._active_version
        }
    
    def set_active(self, version: str) -> None:
        """设置为当前活跃版本"""
        if version not in self._keys:
            raise ValueError(f"密钥版本 {version} 不存在")
        for v, cfg in self._keys.items():
            cfg['active'] = (v == version)
        self._active_version = version
    
    def get_active_key(self) -> tuple[str, bytes]:
        """获取当前活跃密钥"""
        if not self._active_version:
            raise RuntimeError("未设置活跃密钥版本")
        cfg = self._keys[self._active_version]
        return self._active_version, cfg['secret_key']
    
    def verify_with_all_keys(
        self,
        signature: str,
        sign_str: str
    ) -> tuple[bool, str]:
        """
        用所有有效密钥验证签名
        用于密钥切换期间的兼容验证
        返回：(是否验证通过, 命中的密钥版本)
        """
        now = datetime.now()
        for version, cfg in self._keys.items():
            # 检查密钥有效期
            expires = datetime.fromisoformat(cfg['expires_at'])
            if now > expires:
                continue
            
            # 尝试验证
            auth = GmGatewayAuth(
                SignatureConfig(secret_key=cfg['secret_key']),
                None  # 跳过 Nonce 检查，仅验证签名
            )
            expected = auth.compute_signature_raw(sign_str, cfg['secret_key'])
            if stdlib_hmac.compare_digest(expected, signature):
                return True, version
        
        return False, "未命中任何有效密钥"
    
    def rotate_key(self, new_version: str, new_key: bytes) -> str:
        """
        执行密钥轮转
        返回：新密钥版本号
        """
        now = datetime.now()
        # 新密钥有效期 1 年
        expires = now + timedelta(days=365)
        
        # 激活新密钥
        self.add_key_version(new_version, new_key, now, expires)
        self.set_active(new_version)
        
        # 标记旧密钥为待废弃（保留 30 天用于兼容验证）
        for v, cfg in self._keys.items():
            if v != new_version:
                cfg['deprecated_at'] = now.isoformat()
                cfg['deprecation_ttl'] = 30 * 24 * 3600  # 30 天
        
        return new_version

import timeit
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.hmac import HMAC as CryHMAC
from gmssl import sm3, func

key = b'secret-key-32-bytes-length!!!!'
message = b'test message for performance benchmark'

# 方案 1: cryptography 库
def test_cryptography():
    h = CryHMAC(key, hashes.SM3())
    h.update(message)
    return h.finalize()

# 方案 2: gmssl 手动实现
def test_gmssl_manual():
    block_size = 64
    k = key if len(key) <= block_size else sm3.sm3_hash(func.bytes_to_list(key))
    k = k.ljust(block_size, b'\x00')
    i_key_pad = bytes(k[i] ^ 0x36 for i in range(block_size))
    o_key_pad = bytes(k[i] ^ 0x5c for i in range(block_size))
    inner = sm3.sm3_hash(func.bytes_to_list(i_key_pad + message))
    return sm3.sm3_hash(func.bytes_to_list(o_key_pad + bytes.fromhex(inner)))

# 运行测试
n = 10000
t1 = timeit.timeit(test_cryptography, number=n)
t2 = timeit.timeit(test_gmssl_manual, number=n)

print(f"cryptography: {n/t1:.0f} ops/sec (单次 {t1/n*1e6:.1f}μs)")
print(f"gmssl 手工:   {n/t2:.0f} ops/sec (单次 {t2/n*1e6:.1f}μs)")
print(f"差距: {(t2/t1-1)*100:.0f}%")

"""
端到端认证流程性能测试
包含：时间戳检查 + Nonce 去重 + HMAC-SM3 签名验证
"""
import asyncio
from concurrent.futures import ThreadPoolExecutor

async def benchmark_full_auth():
    """模拟高并发认证请求"""
    auth = GmGatewayAuth(
        config=SignatureConfig(secret_key=os.urandom(32)),
        redis_client=redis.Redis()
    )
    
    n = 1000
    tasks = []
    
    for i in range(n):
        method = 'POST'
        path = '/api/v1/transfer'
        query = {'page': str(i)}
        body = f'{{"amount": {i}}}'.encode()
        timestamp = int(time.time())
        nonce = uuid.uuid4().hex
        signature = auth.compute_signature(
            method, path, query, body, timestamp, nonce
        )
        
        async def verify():
            return await asyncio.to_thread(
                auth.verify_signature,
                method, path, query, body, timestamp, nonce, signature
            )
        
        tasks.append(verify())
    
    start = time.perf_counter()
    results = await asyncio.gather(*tasks)
    elapsed = time.perf_counter() - start
    
    success = sum(1 for r, _ in results if r)
    print(f"并发认证: {n} 次请求, {elapsed:.2f}s")
    print(f"吞吐量: {n/elapsed:.0f} req/sec")
    print(f"成功率: {success/n*100:.1f}%")
    print(f"平均延迟: {elapsed/n*1000:.1f} ms/request")

# 运行测试
asyncio.run(benchmark_full_auth())