Technical Audit Document

def calculate_file_hash(file_path: Path, algorithm: str = 'sha256') -> str:
    """
    计算文件的哈希值（极简版，仅用于非确权场景）
    
    Args:
        file_path: 文件路径
        algorithm: 哈希算法，默认 'sha256'
    
    Returns:
        文件的十六进制哈希值
    
    Raises:
        RightsError: 文件不存在或读取失败
    """
    if not file_path or not isinstance(file_path, Path):
        raise RightsError(f"无效的文件路径: {file_path}")
    
    if not file_path.exists():
        raise RightsError(f"文件不存在: {file_path}")
    
    if not file_path.is_file():
        raise RightsError(f"路径不是文件: {file_path}")
    
    try:
        file_size = file_path.stat().st_size
        if file_size > 10 * 1024 * 1024 * 1024:  # 10GB 限制
            raise RightsError(f"文件过大（超过10GB）: {file_path}")
        
        hash_obj = hashlib.new(algorithm)
        with open(file_path, 'rb') as f:
            chunk_size = 8192
            while True:
                chunk = f.read(chunk_size)
                if not chunk:
                    break
                hash_obj.update(chunk)
        
        return hash_obj.hexdigest()
    except PermissionError:
        raise RightsError(f"没有权限读取文件: {file_path}")
    except OSError as e:
        raise RightsError(f"读取文件失败: {file_path}, 错误: {str(e)}")
    except Exception as e:
        raise RightsError(f"计算文件哈希失败: {file_path}, 错误: {str(e)}")

class MerkleTree:
    """
    Merkle Tree 实现，用于将多个文件的哈希值聚合成单个根哈希

    优势：
    - 支持 100+ 文件批量处理
    - 单个根哈希可代表整个批次
    - 节省 TSA 调用成本（1 次调用 vs N 次调用）

    Known limitation (second-preimage resistance):
        This implementation does NOT use domain separation prefixes for leaf vs
        internal nodes (i.e. b'\\x00' for leaves, b'\\x01' for internal nodes as
        recommended by RFC 6962 §2.1).  Adding prefixes would change the root hash
        computation and break backward compatibility with all existing .nep files
        and the verification website (verify.arisfusion.com).  A future tree_version
        bump can introduce domain separation; the current version is safe for our
        threat model (user-submitted files, not adversarial tree construction).
    """
    
    def __init__(self, hash_algorithm: str = 'sha256'):
        """
        初始化 Merkle Tree
        
        Args:
            hash_algorithm: 哈希算法，默认 'sha256'
        """
        self.hash_algorithm = hash_algorithm
        self.leaves: List[str] = []
        self.tree: List[List[str]] = []
        self.root_hash: Optional[str] = None
    
    def add_leaf(self, data: bytes) -> str:
        """
        添加叶子节点（文件哈希）
        
        Args:
            data: 文件数据或哈希值（bytes）
        
        Returns:
            叶子节点的哈希值
        """
        hash_obj = hashlib.new(self.hash_algorithm)
        hash_obj.update(data)
        leaf_hash = hash_obj.hexdigest()
        self.leaves.append(leaf_hash)
        return leaf_hash
    
    def add_file_hash(self, file_hash: str) -> None:
        """
        直接添加文件哈希值（已计算好的）
        
        Args:
            file_hash: 文件的十六进制哈希值
        """
        self.leaves.append(file_hash)
    
    def build(self) -> str:
        """
        构建 Merkle Tree 并返回根哈希
        
        Returns:
            根哈希值（十六进制字符串）
        """
        if not self.leaves:
            raise ValueError("Merkle Tree 没有叶子节点")
        
        # 如果只有一个叶子节点，直接返回
        if len(self.leaves) == 1:
            self.root_hash = self.leaves[0]
            return self.root_hash
        
        # 构建树：从叶子节点开始，逐层向上
        current_level = self.leaves.copy()
        self.tree = [current_level]
        
        while len(current_level) > 1:
            next_level = []
            
            # 成对处理节点
            for i in range(0, len(current_level), 2):
                if i + 1 < len(current_level):
                    # 两个节点：合并哈希
                    combined = current_level[i] + current_level[i + 1]
                else:
                    # 奇数个节点：最后一个节点复制后与自己合并
                    combined = current_level[i] + current_level[i]
                
                # 计算父节点哈希
                hash_obj = hashlib.new(self.hash_algorithm)
                hash_obj.update(combined.encode('utf-8'))
                parent_hash = hash_obj.hexdigest()
                next_level.append(parent_hash)
            
            self.tree.append(next_level)
            current_level = next_level
        
        # 根哈希是最后一层的唯一节点
        self.root_hash = current_level[0]
        return self.root_hash
    
    def get_proof(self, leaf_index: int) -> List[Dict]:
        """
        获取指定叶子节点的 Merkle Proof（用于验证）

        Args:
            leaf_index: 叶子节点索引

        Returns:
            Merkle Proof 路径，每个元素为 {'hash': str, 'position': 'left'|'right'}
            position 表示兄弟节点在合并时的位置
        """
        if not self.tree:
            self.build()

        if leaf_index >= len(self.leaves):
            raise IndexError(f"叶子节点索引超出范围: {leaf_index}")

        proof = []
        current_index = leaf_index
        current_level = 0

        while current_level < len(self.tree) - 1:
            level = self.tree[current_level]

            # 找到兄弟节点并记录位置
            if current_index % 2 == 0:
                # 当前是左节点，兄弟在右侧
                sibling_index = current_index + 1
                if sibling_index < len(level):
                    proof.append({'hash': level[sibling_index], 'position': 'right'})
                else:
                    # 奇数情况，兄弟是自己（已复制）
                    proof.append({'hash': level[current_index], 'position': 'right'})
            else:
                # 当前是右节点，兄弟在左侧
                sibling_index = current_index - 1
                proof.append({'hash': level[sibling_index], 'position': 'left'})

            # 移动到上一层
            current_index = current_index // 2
            current_level += 1

        return proof

    def verify_proof(self, leaf_hash: str, proof: List[Dict], root_hash: str) -> bool:
        """
        验证 Merkle Proof

        Args:
            leaf_hash: 叶子节点哈希
            proof: Merkle Proof 路径（由 get_proof 返回）
            root_hash: 根哈希

        Returns:
            验证是否通过
        """
        current_hash = leaf_hash

        for step in proof:
            sibling_hash = step['hash']
            position = step['position']

            # 按照 build() 相同的位置顺序合并：左 + 右
            if position == 'right':
                combined = current_hash + sibling_hash
            else:
                combined = sibling_hash + current_hash

            hash_obj = hashlib.new(self.hash_algorithm)
            hash_obj.update(combined.encode('utf-8'))
            current_hash = hash_obj.hexdigest()

        return current_hash == root_hash

class TSAClient:
    """
    RFC 3161 时间戳服务客户端

    支持的服务：
    - FreeTSA (https://freetsa.org/tsr) - 免费，无需注册
    - DigiCert (http://timestamp.digicert.com)
    - IdenTrust (http://timestamp.identrust.com)
    - 其他 RFC 3161 兼容服务
    """

    # 预定义的 TSA 服务提供商
    PROVIDERS = {
        'freetsa': {
            'name': 'FreeTSA',
            'url': 'https://freetsa.org/tsr',
            'hashname': 'sha256',
            'description': '免费时间戳服务，国际标准 RFC 3161',
            'requires_auth': False,
            'legal_strength': 3,  # 1-5 评分
            'price': 0
        },
        'digicert': {
            'name': 'DigiCert',
            'url': 'http://timestamp.digicert.com',
            'hashname': 'sha256',
            'description': 'DigiCert 免费时间戳服务',
            'requires_auth': False,
            'legal_strength': 4,
            'price': 0
        },
        'identrust': {
            'name': 'IdenTrust',
            'url': 'http://timestamp.identrust.com',
            'hashname': 'sha256',
            'description': 'IdenTrust 免费时间戳服务',
            'requires_auth': False,
            'legal_strength': 4,
            'price': 0
        }
    }

    def __init__(
        self,
        provider: str = 'freetsa',
        custom_url: Optional[str] = None,
        hashname: str = 'sha256',
        timeout: int = 30
    ):
        """
        初始化 TSA 客户端

        Args:
            provider: 预定义的服务提供商名称 ('freetsa', 'digicert', 'identrust')
            custom_url: 自定义 TSA URL（如果指定，则忽略 provider）
            hashname: 哈希算法 ('sha256', 'sha512' 等)
            timeout: 请求超时时间（秒）
        """
        if not RFC3161_AVAILABLE:
            raise ImportError(
                "rfc3161ng 库未安装。请运行: pip install rfc3161ng"
            )

        if custom_url:
            self.url = custom_url
            self.provider_name = "Custom TSA"
            self.provider_key = None
        elif provider in self.PROVIDERS:
            config = self.PROVIDERS[provider]
            self.url = config['url']
            self.provider_name = config['name']
            self.provider_key = provider
            hashname = config['hashname']
        else:
            raise ValueError(
                f"未知的 TSA 提供商: {provider}。"
                f"支持的提供商: {', '.join(self.PROVIDERS.keys())}"
            )

        self.hashname = hashname
        self.timeout = timeout

        # 初始化 rfc3161ng 时间戳器
        # 注意：某些环境可能遇到 SSL 握手问题，这是正常的
        # 我们的设计会自动降级到本地哈希
        try:
            self.stamper = rfc3161ng.RemoteTimestamper(
                url=self.url,
                hashname=self.hashname,
                timeout=self.timeout
            )
        except Exception as e:
            # 如果初始化失败，记录错误但不抛出异常
            # 后续调用时会返回失败状态
            self.stamper = None
            self._init_error = str(e)

    FAILOVER_ORDER: List[str] = ['digicert', 'freetsa', 'identrust']

    def _call_with_retry(self, hash_bytes: bytes, max_retries: int = 3) -> bytes:
        """
        带指数退避和提供商故障转移的 TSA 调用

        Args:
            hash_bytes: 要签名的哈希值（二进制）
            max_retries: 每个提供商的最大重试次数

        Returns:
            TSR 令牌（二进制）

        Raises:
            Exception: 所有提供商均失败
        """
        # 构建尝试顺序：当前提供商优先，然后是其他提供商
        providers_to_try = []
        if self.provider_key:
            providers_to_try.append(self.provider_key)
            for p in self.FAILOVER_ORDER:
                if p != self.provider_key:
                    providers_to_try.append(p)
        else:
            # 自定义 URL，无故障转移
            providers_to_try = [None]

        last_error = None
        for provider_key in providers_to_try:
            if provider_key is not None:
                config = self.PROVIDERS[provider_key]
                url = config['url']
                hashname = config['hashname']
                provider_name = config['name']
            else:
                url = self.url
                hashname = self.hashname
                provider_name = self.provider_name

            for attempt in range(max_retries):
                try:
                    stamper = rfc3161ng.RemoteTimestamper(
                        url=url,
                        hashname=hashname,
                        timeout=self.timeout
                    )
                    tsr_token = stamper(digest=hash_bytes)
                    # 成功后更新当前提供商信息
                    if provider_key is not None:
                        self.provider_name = provider_name
                        self.provider_key = provider_key
                        self.url = url
                        self.stamper = stamper
                    return tsr_token
                except Exception as e:
                    last_error = e
                    if attempt < max_retries - 1:
                        time.sleep(2 ** attempt)

        raise Exception(
            f"所有 TSA 提供商均失败 (尝试: {', '.join(p or 'custom' for p in providers_to_try)}): {last_error}"
        )

    def timestamp_data(self, data: bytes) -> bytes:
        """
        对原始数据生成时间戳

        Args:
            data: 要加时间戳的数据（二进制）

        Returns:
            TSR (Time-Stamp Response) 二进制令牌

        Raises:
            Exception: 时间戳请求失败
        """
        try:
            # 计算数据的哈希值
            hash_obj = hashlib.new(self.hashname)
            hash_obj.update(data)
            data_hash = hash_obj.digest()

            # 调用 TSA 服务获取时间戳（带重试和故障转移）
            tsr_token = self._call_with_retry(data_hash)

            return tsr_token
        except Exception as e:
            raise Exception(f"时间戳请求失败: {str(e)}")

    def timestamp_file(self, file_path: Path, output_path: Optional[Path] = None) -> Dict:
        """
        为文件生成时间戳（流式处理，不将整个文件载入内存）

        Args:
            file_path: 要加时间戳的文件路径
            output_path: TSR 文件输出路径（可选，默认为 file_path.tsr）

        Returns:
            包含时间戳信息的字典
        """
        file_path = Path(file_path)

        if not file_path.exists():
            return {
                'success': False,
                'message': f"文件不存在: {file_path}"
            }

        try:
            # 流式计算文件哈希（8KB 分块，避免大文件内存溢出）
            hash_obj = hashlib.new(self.hashname)
            with open(file_path, 'rb') as f:
                while True:
                    chunk = f.read(8192)
                    if not chunk:
                        break
                    hash_obj.update(chunk)

            file_hash = hash_obj.hexdigest()

            # 确定输出路径
            if output_path is None:
                output_path = file_path.parent / f"{file_path.stem}.tsr"
            else:
                output_path = Path(output_path)

            # 委托给 timestamp_hash（处理 TSA 调用、重试和文件写入）
            result = self.timestamp_hash(file_hash, output_path)

            # 确保返回文件哈希
            if result.get('success'):
                result['hash'] = file_hash

            return result

        except Exception as e:
            return {
                'success': False,
                'message': f"时间戳生成失败: {str(e)}"
            }

    def timestamp_hash(self, hash_value: str, output_path: Path) -> Dict:
        """
        为已知的哈希值生成时间戳（用于 Merkle Root 等场景）

        Args:
            hash_value: 十六进制哈希值字符串
            output_path: TSR 文件输出路径

        Returns:
            包含时间戳信息的字典
        """
        try:
            # 将十六进制哈希值转换为二进制
            hash_bytes = bytes.fromhex(hash_value)

            # 获取时间戳令牌（带重试和故障转移）
            tsr_token = self._call_with_retry(hash_bytes)

            # 保存 TSR 文件
            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with open(output_path, 'wb') as f:
                f.write(tsr_token)

            # 从 TSR 令牌中提取 TSA 认证时间（而非本机时钟）
            tsa_timestamp = datetime.now().isoformat()  # fallback
            tsa_issuer = self.provider_name
            try:
                from asn1crypto import tsp, cms
                # Try TimeStampResp first, then ContentInfo
                signed_data = None
                try:
                    ts_resp = tsp.TimeStampResp.load(tsr_token)
                    signed_data = ts_resp['time_stamp_token']['content']
                except (ValueError, KeyError, TypeError):
                    try:
                        ci = cms.ContentInfo.load(tsr_token)
                        if ci['content_type'].native == 'signed_data':
                            signed_data = ci['content']
                    except (ValueError, KeyError, TypeError):
                        pass

                if signed_data:
                    tst_info = signed_data['encap_content_info']['content'].parsed
                    gen_time = tst_info['gen_time'].native
                    if gen_time:
                        tsa_timestamp = gen_time.isoformat()

                    # Extract issuer CN from signer info
                    try:
                        signer_infos = signed_data['signer_infos']
                        if signer_infos:
                            sid = signer_infos[0]['sid']
                            if sid.name == 'issuer_and_serial_number':
                                for rdn in sid.chosen['issuer'].chosen:
                                    for attr in rdn:
                                        if attr['type'].dotted == '2.5.4.3':
                                            tsa_issuer = attr['value'].native
                                            break
                    except (KeyError, IndexError, ValueError):
                        pass
            except ImportError:
                pass  # asn1crypto not available, use fallback values

            return {
                'success': True,
                'timestamp': tsa_timestamp,
                'hash': hash_value,
                'issuer': tsa_issuer,
                'tsr_path': str(output_path),
                'algorithm': self.hashname.upper(),
                'message': f"时间戳生成成功，TSR 文件: {output_path.name}"
            }

        except Exception as e:
            return {
                'success': False,
                'message': f"时间戳生成失败: {str(e)}"
            }

    def verify_tsr(
        self,
        tsr_path: Path,
        data: Optional[bytes] = None,
        digest: Optional[bytes] = None,
    ) -> Dict:
        """
        验证 TSR 时间戳令牌

        Args:
            tsr_path: TSR 文件路径
            data: 原始数据（可选，库内部计算哈希后比对）
            digest: 预计算的哈希值（可选，直接与 TSR 中记录的哈希比对）
                    data 和 digest 二选一；都不传则仅验证 TSR 结构

        Returns:
            验证结果字典，包含 'valid', 'message', 'issuer' 等字段。
            当 data/digest 未提供时，结果包含 'partial_verification': True 标志。
        """
        tsr_path = Path(tsr_path)

        if not tsr_path.exists():
            return {
                'valid': False,
                'message': f"TSR 文件不存在: {tsr_path}"
            }

        try:
            with open(tsr_path, 'rb') as f:
                tsr_token = f.read()

            if digest is not None:
                # 直接用预计算的哈希比对（用于 Merkle Root 等场景）
                verified = rfc3161ng.check_timestamp(tsr_token, digest=digest)
                if verified:
                    return {
                        'valid': True,
                        'message': '时间戳验证通过（数据完整性已确认）',
                        'issuer': self._extract_issuer_from_tsr(tsr_token),
                    }
                else:
                    return {
                        'valid': False,
                        'message': '时间戳验证失败：哈希值不匹配'
                    }
            elif data is not None:
                # 传入原始数据，由库内部计算哈希后比对
                verified = rfc3161ng.check_timestamp(tsr_token, data=data)
                if verified:
                    return {
                        'valid': True,
                        'message': '时间戳验证通过（数据完整性已确认）',
                        'issuer': self._extract_issuer_from_tsr(tsr_token),
                    }
                else:
                    return {
                        'valid': False,
                        'message': '时间戳验证失败：哈希值不匹配'
                    }
            else:
                # 结构验证：无原始数据时，验证 TSR 文件结构有效性
                return self._verify_tsr_structure(tsr_token)

        except Exception as e:
            return {
                'valid': False,
                'message': f'时间戳验证失败: {str(e)}'
            }

    def _extract_issuer_from_tsr(self, tsr_token: bytes) -> str:
        """Extract the actual issuer CN from a TSR token using asn1crypto.

        Falls back to self.provider_name if parsing fails or asn1crypto
        is unavailable.  This ensures verify_tsr() reports the real issuer
        even after failover to a different TSA provider.
        """
        try:
            from asn1crypto import tsp, cms

            signed_data = None
            try:
                ts_resp = tsp.TimeStampResp.load(tsr_token)
                signed_data = ts_resp['time_stamp_token']['content']
            except (ValueError, KeyError, TypeError):
                try:
                    ci = cms.ContentInfo.load(tsr_token)
                    if ci['content_type'].native == 'signed_data':
                        signed_data = ci['content']
                except (ValueError, KeyError, TypeError):
                    pass

            if signed_data:
                signer_infos = signed_data['signer_infos']
                if signer_infos:
                    sid = signer_infos[0]['sid']
                    if sid.name == 'issuer_and_serial_number':
                        for rdn in sid.chosen['issuer'].chosen:
                            for attr in rdn:
                                if attr['type'].dotted == '2.5.4.3':  # CN
                                    return attr['value'].native
        except (ImportError, Exception):
            pass

        return self.provider_name

    def _verify_tsr_structure(self, tsr_token: bytes) -> Dict:
        """验证 TSR 令牌的 ASN.1 结构（不验证数据完整性）"""
        if len(tsr_token) < 20:
            return {'valid': False, 'message': 'TSR 文件过小，可能已损坏'}

        if tsr_token[0] != 0x30:
            return {'valid': False, 'message': 'TSR 文件不是有效的 ASN.1 DER 格式'}

        # 尝试 asn1crypto 深度结构验证
        try:
            from asn1crypto import tsp, cms

            # 格式 1: TimeStampResp
            try:
                ts_resp = tsp.TimeStampResp.load(tsr_token)
                status = ts_resp['status']['status'].native
                if status not in ('granted', 'granted_with_mods'):
                    return {
                        'valid': False,
                        'message': f'TSR 状态异常: {status}'
                    }
                return {
                    'valid': True,
                    'message': '时间戳结构验证通过（未验证数据完整性）',
                    'issuer': self.provider_name,
                    'partial_verification': True
                }
            except (ValueError, KeyError, TypeError):
                pass

            # 格式 2: ContentInfo/SignedData (DigiCert 等)
            try:
                content_info = cms.ContentInfo.load(tsr_token)
                if content_info['content_type'].native == 'signed_data':
                    return {
                        'valid': True,
                        'message': '时间戳结构验证通过（未验证数据完整性）',
                        'issuer': self.provider_name,
                        'partial_verification': True
                    }
            except (ValueError, KeyError, TypeError):
                pass

            return {'valid': False, 'message': '无法解析 TSR 结构'}

        except ImportError:
            # asn1crypto 不可用，基础结构检查已通过（size + 0x30 tag）
            return {
                'valid': True,
                'message': '时间戳基础结构检查通过（安装 asn1crypto 可获得深度验证）',
                'issuer': self.provider_name,
                'partial_verification': True
            }

    @classmethod
    def get_provider_info(cls, provider: str) -> Optional[Dict]:
        """获取预定义服务提供商的信息"""
        return cls.PROVIDERS.get(provider)

    @classmethod
    def list_providers(cls) -> Dict[str, Dict]:
        """列出所有预定义的服务提供商"""
        return cls.PROVIDERS.copy()

def batch_protect_works(
    file_paths: List[Path],
    author_name: str,
    inspiration: Optional[str] = None,
    output_dir: Optional[Path] = None,
    password: Optional[str] = None,
    progress_callback=None,
    tsa_provider: str = 'digicert',
    tsa_timeout: int = 30,
    cert_mode: str = 'simple',
) -> Dict:
    """
    批量保护作品（数字存证核心流程）
    
    流程：
    1. 计算所有文件哈希
    2. 构建 Merkle Tree
    3. 生成 manifest.json
    4. 生成缩略图拼贴
    5. 调用 TSA 获取时间戳（使用根哈希）
    6. 生成 PDF 报告
    7. 打包为 .nep 文件
    
    Args:
        file_paths: 文件路径列表
        author_name: 作者名称
        inspiration: 创作灵感（可选）
        output_dir: 输出目录
        password: .nep 文件密码（可选）
        progress_callback: 进度回调 (current, total, message)
    
    Returns:
        包含处理结果的字典
    """
    from .utils import build_merkle_tree_from_files
    from .rights_packer import RightsPacker
    from .pdf_generator import PDFGenerator
    
    if not file_paths:
        raise RightsError("文件列表为空")
    
    if output_dir is None:
        output_dir = Path.cwd() / "digital_evidence"
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        # 步骤 1: 计算文件哈希并构建 Merkle Tree
        if progress_callback:
            progress_callback(0, len(file_paths), "正在计算文件哈希...")
        
        def hash_progress(current, total):
            if progress_callback:
                progress_callback(current, total, f"计算哈希: {current}/{total}")
        
        tree = build_merkle_tree_from_files(file_paths, progress_callback=hash_progress)
        root_hash = tree.root_hash

        image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}

        # 步骤 2: 收集文件哈希和作品信息 + 计算感知哈希
        # Reuse hashes already computed by Merkle Tree (avoid double I/O)
        file_hashes = dict(tree.file_hashes)
        works = []
        fingerprints = []  # Perceptual hashes for image files

        for file_path in file_paths:
            file_hash = file_hashes.get(str(file_path), calculate_file_hash(file_path))

            # 获取文件创建时间
            try:
                creation_time = datetime.fromtimestamp(file_path.stat().st_mtime)
            except (OSError, ValueError, OverflowError):
                creation_time = datetime.now()

            works.append({
                'title': file_path.stem,
                'creation_date': creation_time.isoformat(),
                'file_path': str(file_path),
                'file_hash': file_hash
            })

            # Compute perceptual hash for image files (non-blocking, best-effort)
            if file_path.suffix.lower() in image_extensions:
                try:
                    from .fingerprint import compute_fingerprint
                    fp = compute_fingerprint(file_path, file_sha256=file_hash)
                    fingerprints.append(fp)
                except Exception as e:
                    import logging as _logging
                    _logging.getLogger(__name__).debug(
                        "Skipping perceptual hash for %s: %s", file_path.name, e
                    )
        
        # 步骤 3: 生成 manifest.json
        if progress_callback:
            progress_callback(len(file_paths), len(file_paths), "正在生成清单...")
        
        packer = RightsPacker(output_dir / "evidence.nep", password=password)
        manifest_data = packer.create_manifest(
            author_name=author_name,
            inspiration=inspiration,
            works=works,
            file_hashes=file_hashes
        )
        
        # 步骤 4: 生成缩略图（仅图片文件）
        image_paths = [p for p in file_paths if p.suffix.lower() in image_extensions]
        
        thumbnail_path = None
        if image_paths:
            if progress_callback:
                progress_callback(len(file_paths), len(file_paths), "正在生成缩略图...")
            
            thumbnail_path = output_dir / "thumbnail.jpg"
            packer.generate_thumbnail(image_paths, thumbnail_path)
        
        # 步骤 5: 调用 DigiCert 获取真实 RFC 3161 时间戳
        if progress_callback:
            progress_callback(len(file_paths), len(file_paths), "正在获取 DigiCert 时间戳...")

        tsa_binary_path = output_dir / "proof.tsa"    # RFC 3161 二进制
        local_json_path = output_dir / "proof.json"   # 本地降级 JSON
        timestamp_file = tsa_binary_path              # 最终实际写入的文件

        try:
            from .tsa_client import TSAClient

            # 初始化 TSA 客户端（可配置提供商）
            tsa_client = TSAClient(provider=tsa_provider, timeout=tsa_timeout)

            # 用 Merkle Root Hash 获取时间戳
            tsa_result = tsa_client.timestamp_hash(root_hash, tsa_binary_path)

            if tsa_result['success']:
                timestamp_file = tsa_binary_path
                timestamp_info = {
                    'timestamp': tsa_result['timestamp'],
                    'hash': root_hash,
                    'issuer': tsa_result['issuer'],
                    'algorithm': tsa_result['algorithm'],
                    'valid': True,
                    'tsr_path': str(tsa_binary_path)
                }
            else:
                # TSA 失败，使用本地时间戳（写 .json，不混淆 .tsa）
                timestamp_file = local_json_path
                timestamp_info = {
                    'timestamp': datetime.now().isoformat(),
                    'hash': root_hash,
                    'issuer': 'Nephele Workshop (本地)',

    def _write_zip_contents(zipf, manifest_data, thumbnail_path, timestamp_file,
                            pdf_report, additional_files, source_files):
        """将内容写入 ZIP 文件（供 pack() 内部使用）"""
        # 1. 添加 manifest.json
        manifest_json = json.dumps(manifest_data, indent=2, ensure_ascii=False)
        zipf.writestr('manifest.json', manifest_json.encode('utf-8'))

        # 2. 添加缩略图
        if thumbnail_path and thumbnail_path.exists():
            zipf.write(thumbnail_path, 'thumbnail.jpg')

        # 3. 添加时间戳文件（保持原始扩展名：.tsa=RFC 3161 二进制，.json=本地降级）
        if timestamp_file and timestamp_file.exists():
            archive_name = 'proof.tsa' if timestamp_file.suffix == '.tsa' else 'proof.json'
            zipf.write(timestamp_file, archive_name)

        # 4. 添加 PDF 报告
        if pdf_report and pdf_report.exists():
            zipf.write(pdf_report, 'VerificationReport.pdf')

        # 5. 原始作品文件 → works/ 目录
        # Use indexed filenames (e.g. works/000_photo.jpg) to avoid collisions
        # when multiple source files share the same basename.  Must match the
        # indexed keys in manifest_data['works_map'].
        if source_files:
            for idx, file_path in enumerate(source_files):
                if isinstance(file_path, str):
                    file_path = Path(file_path)
                if file_path.exists() and file_path.is_file():
                    zipf.write(file_path, f'works/{idx:03d}_{file_path.name}')

        # 6. 添加额外文件
        if additional_files:
            for file_path in additional_files:
                if file_path.exists():
                    zipf.write(file_path, f'additional/{file_path.name}')

def verify_evidence_package(
    tsa_path: Path,
    file_paths: List[Path],
) -> Tuple[bool, Dict]:
    """
    深度验证存证包：重算文件哈希 → 重建 Merkle Tree → 与 TSR 中签名的哈希比对。

    验证链路:
      文件列表 → SHA-256 → Merkle Tree → Root Hash → 与 TSR 内 messageImprint 比对

    Args:
        tsa_path: .tsa/.tsr 时间戳文件路径
        file_paths: 被存证的原始文件路径列表（顺序必须与存证时一致）

    Returns:
        (是否通过, 验证结果字典)
    """
    from .utils import build_merkle_tree_from_files

    if not tsa_path.exists():
        return False, {'valid': False, 'message': f"时间戳文件不存在: {tsa_path}"}

    missing = [str(p) for p in file_paths if not p.exists()]
    if missing:
        return False, {'valid': False, 'message': f"原始文件缺失: {', '.join(missing)}"}

    try:
        # Step 1: 重算文件哈希，重建 Merkle Tree
        tree = build_merkle_tree_from_files(file_paths)
        computed_root = tree.root_hash

        # Step 2: 检查是否为本地时间戳（无第三方签名，不具备证明力）
        suffix = tsa_path.suffix.lower()
        if suffix == '.json':
            tsr_result = parse_tsr(tsa_path)
            tsr_hash = tsr_result.get('hash', tsr_result.get('work_identity', ''))
            if computed_root.lower() == tsr_hash.lower():
                return False, {
                    'valid': False,
                    'message': "本地时间戳无第三方签名，不具备密码学证明力。文件哈希一致但无法证明时间。",
                    'root_hash': computed_root,
                    'local_only': True,
                    'file_count': len(file_paths),
                }
            else:
                return False, {
                    'valid': False,
                    'message': "验证失败：Merkle Root 与本地时间戳记录不匹配",
                    'computed_root': computed_root,
                    'file_count': len(file_paths),
                }

        # Step 3: RFC 3161 TSR — 用 rfc3161ng 做真正的密码学签名验证
        computed_digest = bytes.fromhex(computed_root)
        try:
            from .tsa_client import TSAClient
            tsa_client = TSAClient()
            verify_result = tsa_client.verify_tsr(tsa_path, digest=computed_digest)

            if verify_result.get('valid'):
                # 同时解析 TSR 获取时间戳详情（签发时间、签发方）
                tsr_result = parse_tsr(tsa_path)
                return True, {
                    'valid': True,
                    'message': f"深度验证通过：{len(file_paths)} 个文件的 Merkle Root 与 TSA 签名匹配（密码学验证）",
                    'timestamp': tsr_result.get('timestamp'),
                    'issuer': tsr_result.get('issuer'),
                    'root_hash': computed_root,
                    'file_count': len(file_paths),
                }
            else:
                return False, {
                    'valid': False,
                    'message': f"TSA 签名验证失败：{verify_result.get('message', 'unknown')}",
                    'root_hash': computed_root,
                    'file_count': len(file_paths),
                }

        except ImportError:
            # rfc3161ng 未安装，降级为结构性比对（明确标注）
            tsr_result = parse_tsr(tsa_path)
            tsr_hash_raw = tsr_result.get('hash', '')
            tsr_hash = tsr_hash_raw.split(':', 1)[1] if ':' in tsr_hash_raw else tsr_hash_raw

            if computed_root.lower() == tsr_hash.lower():
                return True, {
                    'valid': True,
                    'message': f"结构验证通过（安装 rfc3161ng 可启用密码学签名验证）",
                    'timestamp': tsr_result.get('timestamp'),
                    'issuer': tsr_result.get('issuer'),
                    'root_hash': computed_root,
                    'file_count': len(file_paths),
                    'partial_verification': True,
                }
            else:
                return False, {
                    'valid': False,
                    'message': "验证失败：Merkle Root 与 TSR 记录的哈希不匹配",
                    'computed_root': computed_root,
                    'tsr_hash': tsr_hash,
                    'file_count': len(file_paths),
                }

    except Exception as e:
        return False, {'valid': False, 'message': f"深度验证失败: {e}"}


def batch_protect_works(

    def screenshot(self, path: Optional[str] = None) -> dict:
        """截图保存"""
        if not self._page:
            return {"success": False, "message": "浏览器未打开"}
        try:
            return self._run(self._async_screenshot(path))
        except Exception as e:
            return {"success": False, "message": f"截图失败: {e}"}

    async def _async_screenshot(self, path: Optional[str]) -> dict:
        import tempfile
        if not path:
            path = os.path.join(tempfile.gettempdir(), "nephele_screenshot.png")
        await self._page.screenshot(path=path, full_page=False)
        return {"success": True, "message": f"截图已保存: {path}", "output_path": path}

    def capture(
        self,
        url: str,
        progress_callback=None,
    ) -> Dict:
        """
        Execute the full URL evidence capture pipeline.

        Args:
            url: Target URL to capture
            progress_callback: Optional (step, total, message) callback

        Returns:
            {
                "success": bool,
                "evidence_id": str,
                "output_dir": str,
                "manifest": dict,
                "timestamp_info": dict,
                "message": str,
            }
        """
        total_phases = 5
        self._record(f"Starting evidence capture for: {url}")

        try:
            # Phase 1: Environment + DNS + TLS certificate
            if progress_callback:
                progress_callback(1, total_phases, "Collecting environment info...")
            environment = self.collect_environment()
            dns_info = self.resolve_dns(url)
            tls_info = self.capture_tls_certificate(url)

            # Phase 2: Browser capture (navigate + screenshot + HTML + images)
            if progress_callback:
                progress_callback(2, total_phases, "Capturing page...")
            artifacts = self.capture_page(url, progress_callback=None)

            # Phase 3: Hash all artifacts (including TLS cert files)
            if progress_callback:
                progress_callback(3, total_phases, "Computing hashes...")
            file_hashes = self.hash_artifacts(artifacts)

            # Also hash TLS cert + response headers files
            for extra_key in ("der_path", "pem_path"):
                p = Path(tls_info.get(extra_key, ""))
                if p.exists():
                    file_hashes[p.name] = {
                        "sha256": self._sha256_file(p),
                        "size": p.stat().st_size,
                        "type": "tls_certificate",
                    }
            resp_headers_path = artifacts.get("response_headers_path")
            if resp_headers_path and Path(resp_headers_path).exists():
                p = Path(resp_headers_path)
                file_hashes[p.name] = {
                    "sha256": self._sha256_file(p),
                    "size": p.stat().st_size,
                    "type": "response_headers",
                }

            # Phase 4: Generate manifest (save log first so it's included)
            log_path = Path(self.save_log())
            self._log_committed = True  # Log content is now frozen for hashing
            file_hashes[log_path.name] = {
                "sha256": self._sha256_file(log_path),
                "size": log_path.stat().st_size,
                "type": "operation_log",
            }

            if progress_callback:
                progress_callback(4, total_phases, "Generating manifest...")
            manifest = self.generate_manifest(
                target_url=url,
                environment=environment,
                dns_info=dns_info,
                artifacts=artifacts,
                file_hashes=file_hashes,
                tls_info=tls_info,
            )

            # Phase 5: Timestamp
            if progress_callback:
                progress_callback(5, total_phases, "Requesting timestamp...")
            ts_info = self.timestamp_manifest(manifest)

            # DO NOT re-save operation log — the version already hashed in manifest
            # is the authoritative one. Any further _record() calls only live in memory.

            image_count = len(artifacts.get("images", []))

            # Phase 6: Package as .nep (tamper-proof archive)
            # Uses files on disk (which match manifest hashes)
            nep_path = self._package_nep()

            return {
                "success": True,
                "evidence_id": self._evidence_id,
                "output_dir": str(self._output_dir),
                "nep_path": str(nep_path),
                "manifest": manifest,
                "timestamp_info": ts_info,
                "page_title": artifacts.get("page_title", ""),
                "image_count": image_count,
                "file_count": len(file_hashes),
                "message": (
                    f"URL evidence captured: {len(file_hashes)} files, "
                    f"{image_count} images, timestamp by {ts_info.get('issuer', 'N/A')}"
                ),
            }

        except URLEvidenceError as e:
            self._record(f"FATAL: {e}")
            if not self._log_committed:
                self.save_log()
            return {
                "success": False,
                "evidence_id": self._evidence_id,
                "output_dir": str(self._output_dir),
                "message": str(e),
            }
        except Exception as e:
            self._record(f"UNEXPECTED ERROR: {e}")
            if not self._log_committed:
                self.save_log()
            logger.exception("URL evidence capture failed")
            return {
                "success": False,
                "evidence_id": self._evidence_id,
                "output_dir": str(self._output_dir),
                "message": f"Unexpected error: {e}",
            }

    # ===== Helpers =====

    def capture_tls_certificate(self, url: str) -> Dict:
        """
        Capture the server's TLS certificate chain.
        This proves the connection was made to the authentic server —
        you can't forge a CA-signed certificate.
        """
        import ssl
        from urllib.parse import urlparse

        self._record("Capturing TLS server certificate")
        parsed = urlparse(url if "://" in url else f"https://{url}")
        hostname = parsed.hostname or ""
        port = parsed.port or 443

        if not hostname:
            return {"error": "Invalid hostname"}

        try:
            ctx = ssl.create_default_context()
            with ctx.wrap_socket(
                socket.socket(socket.AF_INET, socket.SOCK_STREAM),
                server_hostname=hostname,
            ) as sock:
                sock.settimeout(10)
                sock.connect((hostname, port))
                cert = sock.getpeercert()
                cert_der = sock.getpeercert(binary_form=True)

            # Save DER certificate to file
            cert_path = self._output_dir / "server_certificate.der"
            cert_path.write_bytes(cert_der)

            # Also save human-readable PEM
            import base64
            pem_data = (
                "-----BEGIN CERTIFICATE-----\n"
                + base64.encodebytes(cert_der).decode()
                + "-----END CERTIFICATE-----\n"
            )
            pem_path = self._output_dir / "server_certificate.pem"
            pem_path.write_text(pem_data, encoding="utf-8")

            # Extract key fields
            subject = dict(x[0] for x in cert.get("subject", ()))
            issuer = dict(x[0] for x in cert.get("issuer", ()))
            cert_info = {
                "subject_cn": subject.get("commonName", ""),
                "issuer_cn": issuer.get("commonName", ""),
                "issuer_org": issuer.get("organizationName", ""),
                "not_before": cert.get("notBefore", ""),
                "not_after": cert.get("notAfter", ""),
                "serial_number": cert.get("serialNumber", ""),
                "san": [
                    entry[1]
                    for entry in cert.get("subjectAltName", ())
                    if entry[0] == "DNS"
                ],
                "der_path": str(cert_path),
                "pem_path": str(pem_path),
                "der_sha256": hashlib.sha256(cert_der).hexdigest(),
            }

            self._record(
                f"TLS cert captured: {cert_info['subject_cn']} "
                f"(issuer: {cert_info['issuer_org']}, "
                f"serial: {cert_info['serial_number'][:16]}...)"
            )
            return cert_info

        except Exception as e:
            self._record(f"TLS certificate capture failed: {e}")
            return {"error": str(e)}

    # ===== Step 2: DNS resolution =====

    _CAPTCHA_KEYWORDS = (
        "验证码", "验证", "captcha", "verify", "challenge",
        "human verification", "robot", "机器人",
    )

    def _is_captcha_page(self, title: str, url: str) -> bool:
        """Detect if the loaded page is a CAPTCHA or anti-bot challenge."""
        text = (title or "").lower()
        url_lower = (url or "").lower()
        for kw in self._CAPTCHA_KEYWORDS:
            if kw in text or kw in url_lower:
                return True
        return False

"""
Nephele Workshop - Watermark Protection Module (local)

Embeds/extracts invisible blind watermarks using blind_watermark locally.
No network dependency — works offline.

Uses file-based I/O + bit mode for reliable embedding/extraction
(numpy array mode has known issues with blind_watermark library).

Public API unchanged:
    protect_image(image, level, copyright_info) -> Image
    extract_watermark(image) -> str | None
    save_with_watermark(image, output_path) -> bool

Developer: ArisFusion Studio
"""

import logging
import tempfile
from enum import Enum
from pathlib import Path
from typing import Optional

import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

# Watermark payload: fixed 32 bytes (256 bits)
# Enough for 10 Chinese chars (UTF-8, 3 bytes each) or 32 ASCII chars
WATERMARK_BYTES = 32
WATERMARK_BITS = WATERMARK_BYTES * 8
_WM_PASSWORD_IMG = 2024
_WM_PASSWORD_WM = 1314


class ProtectionLevel(Enum):
    NONE = "none"
    INVISIBLE = "invisible"


LEVEL_ALIASES = {"maximum": "invisible"}


def _text_to_bits(text: str) -> list[int]:
    """Convert text to fixed-length bit array via UTF-8."""
    raw = text.encode("utf-8")[:WATERMARK_BYTES]
    padded = raw.ljust(WATERMARK_BYTES, b"\x00")
    bits = []
    for byte in padded:
        for i in range(7, -1, -1):
            bits.append((byte >> i) & 1)
    return bits


def _bits_to_text(bits: list) -> str:
    """Convert bit array back to text via UTF-8."""
    raw = bytearray()
    for i in range(0, len(bits), 8):
        chunk = bits[i:i+8]
        if len(chunk) < 8:
            break
        val = 0
        for b in chunk:
            val = (val << 1) | (1 if b > 0.5 else 0)
        raw.append(val)
    return raw.rstrip(b"\x00").decode("utf-8", errors="replace")

class WatermarkEngine:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def embed(self, image: Image.Image, text: str) -> Image.Image:
        """Embed invisible watermark using file-based blind_watermark."""
        try:
            from blind_watermark import WaterMark

            rgb = image.convert("RGB")
            alpha = image.split()[3] if image.mode == "RGBA" else None

            with tempfile.TemporaryDirectory() as tmpdir:
                orig_path = str(Path(tmpdir) / "orig.png")
                wm_path = str(Path(tmpdir) / "watermarked.png")

                rgb.save(orig_path, format="PNG")

                bits = _text_to_bits(text)

                bwm = WaterMark(password_img=_WM_PASSWORD_IMG, password_wm=_WM_PASSWORD_WM)
                bwm.read_img(orig_path)
                bwm.read_wm(np.array(bits), mode="bit")
                bwm.embed(wm_path)

                result_img = Image.open(wm_path).convert("RGB")

                # Verify extraction round-trip
                extracted_bits = WaterMark(
                    password_img=_WM_PASSWORD_IMG, password_wm=_WM_PASSWORD_WM
                ).extract(wm_path, wm_shape=WATERMARK_BITS, mode="bit")
                extracted_text = _bits_to_text(extracted_bits)

                if extracted_text == text[:WATERMARK_BYTES]:
                    logger.info("[Watermark] Verified: '%s'", extracted_text)
                else:
                    logger.warning("[Watermark] Verify mismatch: '%s' -> '%s'",
                                   text[:WATERMARK_BYTES], extracted_text)

            if alpha:
                result_img = result_img.convert("RGBA")
                result_img.putalpha(alpha)

            return result_img

        except Exception as e:
            logger.error("Embed failed: %s", e, exc_info=True)
            return image

    def extract(self, image: Image.Image) -> Optional[str]:
        """Extract invisible watermark using file-based blind_watermark."""
        try:
            from blind_watermark import WaterMark

            rgb = image.convert("RGB")
            bit_len = WATERMARK_BITS

            with tempfile.TemporaryDirectory() as tmpdir:
                img_path = str(Path(tmpdir) / "check.png")
                rgb.save(img_path, format="PNG")

                extracted_bits = WaterMark(
                    password_img=_WM_PASSWORD_IMG, password_wm=_WM_PASSWORD_WM
                ).extract(img_path, wm_shape=bit_len, mode="bit")

            text = _bits_to_text(extracted_bits)
            return text.strip() if text.strip() else None
        except Exception as e:
            logger.warning("Extract failed: %s", e)
            return None

def protect_image(
    image: Image.Image,
    level: str = "none",
    copyright_info: str = "ARIS"
) -> Image.Image:
    level = LEVEL_ALIASES.get(level, level)
    if level == "invisible":
        return WatermarkEngine().embed(image, copyright_info)
    return image


def extract_watermark(image: Image.Image) -> Optional[str]:
    return WatermarkEngine().extract(image)


def save_with_watermark(image: Image.Image, output_path: str) -> bool:
    try:
        image.save(output_path, format='PNG')
        return True
    except Exception as e:
        logger.warning("Save failed: %s", e)
        return False

def pack_image(
    input_path: str,
    watermark_path: Optional[str] = None,
    output_dir: Optional[str] = None,
    watermark_mode: str = "center",
    watermark_opacity: float = 0.3,
    preview_max_size: int = 1920,
    thumbnail_max_size: int = 500,
    protection_level: str = "none",
    copyright_info: str = "© ArisFusion Studio",
    output_folder_name: str = "Delivery_Pack",
) -> dict:
    """One-click image packing: HD + preview + thumbnail."""
    try:
        src = Path(input_path)
        if not src.exists():
            return api_err(f"文件不存在: {input_path}")

        mode_map = {"center": WatermarkMode.CENTER, "tile": WatermarkMode.TILE}
        wm_mode = mode_map.get(watermark_mode, WatermarkMode.CENTER)

        wm_path = Path(watermark_path) if watermark_path else None
        if wm_path and not wm_path.exists():
            return api_err(f"水印文件不存在: {watermark_path}")

        out_dir = Path(output_dir) if output_dir else None

        packer = DeliveryPacker(
            preview_max_size=preview_max_size,
            thumbnail_max_size=thumbnail_max_size,
            watermark_opacity=watermark_opacity,
            protection_level=protection_level,
            copyright_info=copyright_info,
            output_folder_name=output_folder_name,
        )

        result_dir, results = packer.process_image(
            input_path=src,
            watermark_path=wm_path,
            output_dir=out_dir,
            watermark_mode=wm_mode,
        )

        file_info = {k: str(v) for k, v in results.items()}
        return api_ok(
            f"打包完成，共生成 {len(results)} 个文件",
            output_path=str(result_dir),
            data={"files": file_info},
        )

    except PackerError as e:
        logger.error("打包失败: %s", e)
        return api_err(str(e))
    except Exception as e:
        logger.exception("打包时发生意外错误")
        return api_err(f"意外错误: {e}")

"""Watermark Extraction Worker"""
from pathlib import Path
from PySide6.QtCore import QThread, Signal, QCoreApplication

_tr = QCoreApplication.translate

from .._utils import ensure_src_path


class WatermarkExtractWorker(QThread):
    """Worker thread for watermark extraction (rivaGan model loading is slow)."""
    finished = Signal(str)  # watermark result or empty string
    logMessage = Signal(str, str)  # (message, level)

    def __init__(self, image_path: str):
        super().__init__()
        self.image_path = image_path

    def run(self):
        """Extract watermark in background thread."""
        try:
            ensure_src_path()

            from PIL import Image
            from tools.packer.watermark_protection import extract_watermark

            img_path = Path(self.image_path)
            if not img_path.exists():
                self.logMessage.emit(_tr("WatermarkExtractWorker", "文件不存在: %s") % self.image_path, "error")
                self.finished.emit("")
                return

            self.logMessage.emit(_tr("WatermarkExtractWorker", "正在提取水印: %s") % img_path.name, "info")

            image = Image.open(img_path)
            watermark = extract_watermark(image)

            if watermark:
                self.logMessage.emit(_tr("WatermarkExtractWorker", "提取成功: %s") % watermark, "success")
                self.finished.emit(watermark)
            else:
                self.logMessage.emit("未检测到隐形水印", "warning")
                self.finished.emit("")

        except Exception as e:
            self.logMessage.emit(f"提取失败: {str(e)}", "error")
            self.finished.emit("")