first commit

.gitignore

@@ -0,0 +1,72 @@
+/target
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+.pytest_cache/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+.venv/
+env/
+bin/
+build/
+develop-eggs/
+dist/
+eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+include/
+man/
+venv/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+pip-selfcheck.json
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.cache
+nosetests.xml
+coverage.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# Rope
+.ropeproject
+
+# Django stuff:
+*.log
+*.pot
+
+.DS_Store
+
+# Sphinx documentation
+docs/_build/
+
+# PyCharm
+.idea/
+
+# VSCode
+.vscode/
+
+# Pyenv
+.python-version

pyproject.toml

@@ -0,0 +1,21 @@
+[build-system]
+build-backend = "hatchling.build"
+requires = ["hatchling"]
+
+[project]
+name = "bytedance"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "requests",
+    "pyjwt",
+    "pymerkle",
+    "jwcrypto",
+    "loguru",
+    "wget"
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/bytedance"]

src/bytedance/__init__.py

@@ -0,0 +1,4 @@
+from pkgutil import extend_path
+
+# Make "bytedance" a namespace package
+__path__ = extend_path(__path__, __name__)

src/bytedance/jeddak_secure_channel/__init__.py

@@ -0,0 +1,28 @@
+__all__ = [
+    "Client",
+    "ClientConfig",
+    "FileMode",
+    "RaRequest",
+    "RaResponse",
+    "ResponseKey",
+    "Server",
+    "ServerConfig",
+    "TksAppOptions",
+    "__version__",
+]
+
+__version__ = "0.1.7.78"
+
+from .client import Client
+from .config import ClientConfig, ServerConfig
+from .crypto import FileMode, ResponseKey
+from .ra import RaRequest, RaResponse
+# from .server import Server
+from .service import TksAppOptions
+
+
+def __getattr__(name: str) -> object:
+    if name == "Server":
+        from .server import Server
+        return Server
+    raise AttributeError(f"module {__name__} has no attribute {name}")

src/bytedance/jeddak_secure_channel/client.py

@@ -0,0 +1,300 @@
+"""
+数据发送方接口.
+"""
+
+__all__ = ["Client"]
+
+import base64
+import json
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from typing_extensions import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+if TYPE_CHECKING:
+    from _typeshed import GenericPath
+else:
+    GenericPath = str
+
+from . import error
+from .config import ClientConfig
+from .crypto import ClientSessionKey, FileMode, PrivateKey, ResponseKey, AesKey
+from .log import LogConfig, logger
+from .ra import attest_server, generate_nonce
+from .utils import RepeatTimer
+
+
+@dataclass(eq=False)
+class Client:
+    """
+    数据发送方接口.
+
+    本类的所有方法是线程安全的.
+    """
+
+    config: ClientConfig
+
+    session_key: Optional[ClientSessionKey] = field(init=False, default=None)
+
+    attest_timer: Optional[RepeatTimer] = field(init=False, default=None)
+
+    root_key_info: Dict[str, PrivateKey] = field(init=False, default_factory=dict)
+
+    last_ra_time: float = 0.0
+
+    def __post_init__(self):
+        try:
+            self._init_log()
+
+            if self.config.root_key_config:
+                self._root_key_parse()
+
+            # if self.config.attest_interval is not None:
+            #     # 初始化客户端时，以当前时间作为最近活跃时间
+            #     self.last_active_time = time.time()
+            #     # 内置最小定时间隔配置，避免因用户的不当配置造成的高负荷
+            #     self.attest_timer = timer = weak_repeat_timer(
+            #         max(self.config.attest_interval, ClientConst.MIN_ATTEST_INTERVAL),
+            #         self._attest_server_no_raise,
+            #     )
+            #     timer.start()
+
+            self._attest_server_no_raise()
+        except Exception as e:
+            logger.critical(f"__post_init__ failed, e={e}")
+
+    def _attest_server_no_raise(self):
+        logger.info("attest_server active")
+        try:
+            if not self.attest_server(None):
+                logger.error("ra attest failed, Please check your config !")
+            else:
+                self.last_ra_time = time.time()
+                logger.info("attest_server success")
+        except Exception as e:
+            logger.critical(f"Failed to attest server, {e}")
+
+    def attest_server(self, token: Optional[str] = None) -> bool:
+        """
+        对数据接收方进行远程证明.
+        """
+        if self.config.pub_key_path:
+            with open(self.config.pub_key_path) as f:
+                public_pem = f.read()
+            logger.debug(f"read pub key success, pub_key={public_pem}")
+            self.session_key = ClientSessionKey.load(public_pem)
+            return True
+        else:
+            max_try_count = 3
+            count = 0
+
+            nonce = generate_nonce()
+            while count < max_try_count:
+                count += 1
+                ret, pub_key = attest_server(token, self.config, nonce)
+                if ret:
+                    logger.info("ra attest success")
+                    if pub_key:
+                        logger.debug(f"get pub key success, pub_key={pub_key}")
+                        self.session_key = ClientSessionKey.load(pub_key)
+                        return True
+                else:
+                    if pub_key and not self.config.attest_must:
+                        logger.warning(f"ra attest failed, but get pub key success, pub_key={pub_key}")
+                        self.session_key = ClientSessionKey.load(pub_key)
+                        return True
+                    else:
+                        time.sleep(2)
+                        logger.info("ra attest failed, try again...")
+
+            return False
+
+    def _get_key(self) -> ClientSessionKey:
+        # 线程安全: 将 session_key 字段保存为本地变量, 避免访问两次
+        key = self.session_key
+        if key is None:
+            raise error.KeyMissingError()
+        return key
+
+    def encrypt(self, plaintext: Union[str, bytes]) -> str:
+        """
+        加密字符串或二进制数据.
+
+        Returns:
+            信封加密数据.
+
+        Raises:
+            KeyMissingError: 没有可用的密钥. 应通过 `attest_server` 验证数据接收方, 并获取密钥.
+        """
+        message, _ = self.encrypt_with_response(plaintext)
+        return message
+
+    def encrypt_with_response(
+            self, plaintext: Union[str, bytes], aes_key: AesKey = None
+    ) -> Tuple[str, ResponseKey]:
+        """
+        加密字符串或二进制数据, 并获取用于解密响应的密钥.
+
+        Returns:
+            信封加密数据, 及用于解密响应的密钥.
+
+        Raises:
+            KeyMissingError: 没有可用的密钥. 应通过 `attest_server` 验证数据接收方, 并获取密钥.
+        """
+        try:
+            if self.config.attest_interval is not None \
+                    and time.time() - self.last_ra_time > self.config.attest_interval:
+                self._attest_server_no_raise()
+
+            message, response_key = self._get_key().encrypt_with_response(plaintext, aes_key)
+            return message.serialize(), response_key
+        except error.KeyMissingError as e:
+            raise e
+        except Exception as e:
+            raise error.EncryptionError(
+                "Encrypt failed, Please check your config"
+            ) from e
+
+    def encrypt_file(
+            self,
+            source_path: GenericPath,
+            dest_path: GenericPath,
+            mode: FileMode,
+    ) -> str:
+        """
+        加密文件数据.
+
+        Returns:
+            信封加密元数据. 可伴随加密文件发送给数据接收方.
+        """
+        metadata, _ = self.encrypt_file_with_response(source_path, dest_path, mode)
+        return metadata
+
+    def encrypt_file_with_response(
+            self,
+            source_path: GenericPath,
+            dest_path: GenericPath,
+            mode: FileMode,
+    ) -> Tuple[str, ResponseKey]:
+        """
+        加密文件数据, 并获取用于解密响应的密钥.
+
+        Returns:
+            信封加密元数据, 及用于解密响应的密钥. 其中元数据可伴随加密文件发送给数据接收方.
+        """
+        if self.config.attest_interval is not None and time.time() - self.last_ra_time > self.config.attest_interval:
+            self._attest_server_no_raise()
+
+        if mode not in ["b", "t"]:
+            raise error.ParamError("mode", "Mode should be 'b' or 't'")
+
+        try:
+            key = self._get_key()
+            with open(source_path, f"r{mode}") as source, open(
+                    dest_path, f"w{mode}"
+            ) as dest:
+                metadata, response_key = key.encrypt_stream_with_response(
+                    source, dest, mode
+                )
+
+            return metadata.serialize(), response_key
+        except error.KeyMissingError as e:
+            raise e
+        except Exception as e:
+            raise error.EncryptionError(
+                "Encrypt file failed, Please check your config"
+            ) from e
+
+    def _root_key_parse(self):
+        try:
+            root_key_config = self.config.root_key_config
+            if not root_key_config:
+                return
+
+            if isinstance(root_key_config, str):
+                root_key_config = json.loads(root_key_config)
+
+            for app_info, pri_key_path in root_key_config.items():
+                with open(pri_key_path) as f:
+                    private_pem = f.read()
+                self.root_key_info[app_info] = PrivateKey.from_private_pem(private_pem)
+
+            logger.info("root_key_parse success")
+        except Exception as e:
+            logger.critical("root_key_parse exception={}".format(e))
+            raise error.ConfigError("RootKey conf", "RootPriKey conf error")
+
+    def gen_sign(self, app_info: str, msg: Union[str, bytes]) -> str:
+        """
+        生成签名
+        :param app_info: app_id+app_name
+        :param msg: app_id+app_name
+        :return:
+        """
+        try:
+            if not app_info or not msg:
+                logger.error("gen_sign error, app_info is NULL")
+                raise error.SignatureError("RSA signing failed") from None
+
+            if app_info not in self.root_key_info:
+                logger.info(
+                    f"gen_sign error, can not find app_info in root key conf, app_info={app_info}"
+                )
+                raise error.SignatureError(
+                    f"can not find app_info={app_info}"
+                ) from None
+
+            return base64.b64encode(self.root_key_info[app_info].sign(msg)).decode()
+
+        except Exception as e:
+            logger.critical("gen_sign exception={}".format(e))
+            raise error.SignatureError("can not find app_info") from e
+
+    # def dynamic_update_interval(self):
+    #     """
+    #     动态更新定时逻辑的时间间隔。对不活跃的客户端，增大其定时逻辑的时间间隔
+    #     """
+    #     try:
+    #         # 对于非定时RA的客户端，不做动态更新
+    #         if self.config.attest_interval is None or self.attest_timer is None:
+    #             return
+    #
+    #         # 如果超过一定时间没有活跃，就增大定时逻辑的时间间隔
+    #         if time.time() - self.last_active_time > ClientConst.INACTIVE_THRESHOLD:
+    #             self.attest_timer.interval = max(
+    #                 self.config.attest_interval, ClientConst.INACTIVE_ATTEST_INTERVAL
+    #             )
+    #         else:
+    #             # 如果活跃时间在阈值内，就恢复初始的时间间隔
+    #             self.attest_timer.interval = max(
+    #                 self.config.attest_interval, ClientConst.MIN_ATTEST_INTERVAL
+    #             )
+    #     except Exception as e:
+    #         logger.error("dynamic_update_interval exception={}".format(e))
+
+    def _init_log(self):
+        config = LogConfig()
+
+        log_config: Dict[str, Any] = {}
+        if self.config.log_config:
+            log_config = json.loads(self.config.log_config)
+
+        if dir := log_config.get("dir"):
+            config.dir = dir
+
+        if filename := log_config.get("filename"):
+            config.filename = filename
+        # else:
+        #     config.filename = "client.log"
+
+        if rotation := log_config.get("rotation"):
+            config.rotation = rotation
+
+        if retention := log_config.get("retention"):
+            config.retention = retention
+
+        if level := log_config.get("level"):
+            config.level = level
+
+        logger.init_log_config(config)

src/bytedance/jeddak_secure_channel/config.py

@@ -0,0 +1,171 @@
+"""
+安全通信配置.
+"""
+
+import json
+import os
+from dataclasses import dataclass
+from datetime import timedelta
+
+from typing_extensions import TYPE_CHECKING, Optional, Self
+
+if TYPE_CHECKING:
+    from _typeshed import GenericPath
+else:
+    GenericPath = str
+
+from .ra import RaConfig
+from .service import EpsConfig, RasConfig, TksConfig
+
+
+@dataclass
+class ClientConfig(RaConfig):
+    """
+    安全通信数据发送方配置.
+    """
+
+    attest_interval: Optional[float] = 3600
+    """
+    自动对数据接收方进行远程证明的时间间隔 (秒).
+    为空表示不自动进行远程证明 (默认).
+    """
+
+    root_key_config: Optional[str] = None
+    """
+    root_key的配置，json格式的字符串.
+    """
+
+    log_config: Optional[str] = None
+
+    pub_key_path: Optional[str] = None
+
+    @classmethod
+    def from_file(cls, path: GenericPath) -> Self:
+        """从文件读取配置对象."""
+
+        with open(path) as f:
+            return cls.from_dict(json.load(f))
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> Self:
+        """从 dict 对象创建配置对象."""
+
+        obj.pop("$schema", None)
+
+        return cls(**obj)
+
+    @classmethod
+    def from_env(cls, env_key: str) -> Self:
+        """从环境变量读取配置对象."""
+
+        config = os.getenv(env_key, "{}")
+        return cls.from_dict(json.loads(config))
+
+
+@dataclass
+class ServerConfig(EpsConfig, RasConfig, TksConfig):
+    """
+    安全通信数据接收方配置.
+    """
+
+    tks_app_id: Optional[str] = None
+    """
+    TKS app id.
+    """
+
+    tks_ring_id: Optional[str] = None
+    """
+    TKS 密钥环 id.
+    为空表示不访问 TKS 获取密钥 (默认).
+    """
+
+    tks_key_id: Optional[str] = None
+    """
+    TKS 密钥 id.
+    为空表示不访问 TKS 获取密钥 (默认).
+    """
+
+    key_file: Optional[GenericPath] = None
+    """
+    密钥文件路径.
+    为空表示不从文件读取密钥 (默认).
+    """
+
+    refresh_interval: Optional[float] = None
+    """
+    自动从 TKS 或者文件获取密钥的时间间隔 (秒).
+    为空表示不自动获取 (默认).
+    """
+
+    valid_time: float = timedelta(minutes=30).total_seconds()
+    """【未启用】密钥的有效期 (秒). 默认为 30 分钟. """
+
+    client_ra_url: Optional[str] = None
+    """
+    客户端 RA 接口地址.
+    为空表示不需要对客户端做 RA (默认).
+    """
+
+    pcc_login_url: Optional[str] = None
+    pcc_app_id: Optional[str] = None
+    pcc_password: Optional[str] = None
+    """
+    pcc_controller模块的登录信息.
+    为空表示不需要做.
+    """
+
+    bytedance_pcc_info: str = ""
+    """访问bytedance pcc_controller 网关的信息."""
+
+    bytedance_top_info: str = ""
+    """访问bytedance TOP网关的信息."""
+
+    root_key_config: Optional[str] = None
+    """
+    root_key的配置，json格式的字符串.
+    """
+
+    client_ra_config: Optional[str] = None
+    """
+    客户端RA的配置，json格式的字符串.
+    """
+
+    log_config: Optional[str] = None
+
+    need_evidence: bool = True
+
+    attest_gpu: bool = True
+
+    key_type: str = "tks"
+
+    service_id: str = ""
+
+    tks_policy_id: str = ""
+
+    tks_url: str = ""
+
+    tks_service_name: str = ""
+
+    @classmethod
+    def from_file(cls, path: GenericPath) -> Self:
+        """从文件读取配置对象."""
+
+        with open(path) as f:
+            return cls.from_dict(json.load(f))
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> Self:
+        """从 dict 对象创建配置对象."""
+
+        obj.pop("$schema", None)
+
+        return cls(**obj)
+
+    @classmethod
+    def from_env(cls, env_key: str) -> Self:
+        """从env中读取配置对象."""
+
+        conf_value = os.getenv(env_key, "{}")
+        if isinstance(conf_value, str):
+            conf_value = json.loads(conf_value)
+        return cls.from_dict(conf_value)

src/bytedance/jeddak_secure_channel/consts.py

@@ -0,0 +1,5 @@
+class ClientConst:
+    MIN_ATTEST_INTERVAL = 3600  # 内置最小的attest间隔，单位秒
+    INACTIVE_ATTEST_INTERVAL = 3600 * 3  # 非活跃客户端的attest时间间隔，单位秒
+
+    INACTIVE_THRESHOLD = 3600  # 超过一定阈值被认为是不活跃的客户端，单位秒

src/bytedance/jeddak_secure_channel/crypto/__init__.py

@@ -0,0 +1,14 @@
+__all__ = [
+    "AesKey",
+    "ClientSessionKey",
+    "EncryptedMessage",
+    "FileMode",
+    "PrivateKey",
+    "PublicKey",
+    "ResponseKey",
+    "ServerSessionKey",
+]
+
+from .aes import AesKey, FileMode
+from .envelope import ClientSessionKey, EncryptedMessage, ResponseKey, ServerSessionKey
+from .rsa import PrivateKey, PublicKey

src/bytedance/jeddak_secure_channel/crypto/aes.py

@@ -0,0 +1,174 @@
+"""
+AES 加密算法.
+"""
+
+__all__ = [
+    "AES_KEY_LEN",
+    "AES_MAC_LEN",
+    "AES_NONCE_LEN",
+    "AesKey",
+    "FileMode",
+]
+
+import base64
+import secrets
+from dataclasses import dataclass
+from io import BytesIO
+
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from typing_extensions import IO, Literal, Optional, Self, Tuple
+
+from .. import error
+
+AES_KEY_LEN = 32
+
+AES_NONCE_LEN = 12
+
+AES_MAC_LEN = 16
+
+CHUNK_SIZE = 1024
+
+FileMode = Literal["b", "t"]
+"""
+文件加密模式: b 表示二进制, t 表示文本 (按行加密, 加密数据 Base64 编码后按行保存).
+"""
+
+
+@dataclass(eq=False, frozen=True)
+class AesKey:
+    """
+    AES 密钥.
+
+    本类不含可变状态, 所有方法是线程安全的.
+    """
+
+    key: bytes
+
+    def __post_init__(self):
+        if len(self.key) != AES_KEY_LEN:
+            raise error.ParamError("key", "Invalid AES key length")
+
+    @classmethod
+    def generate(cls) -> Self:
+        """
+        生成 AES 密钥.
+        """
+        key = secrets.token_bytes(AES_KEY_LEN)
+
+        return cls(key)
+
+    def encrypt(
+        self,
+        plaintext: bytes,
+        nonce: Optional[bytes] = None,
+    ) -> Tuple[bytes, bytes, bytes]:
+        """
+        用 AES-GCM 加密数据.
+
+        Returns:
+            单次随机数, 密文, 及校验码.
+
+        Raises:
+            EncryptionError:
+        """
+        with BytesIO(plaintext) as source, BytesIO() as dest:
+            nonce, mac = self.encrypt_stream(source, dest, "b", nonce)
+            ciphertext = dest.getvalue()
+
+        return nonce, ciphertext, mac
+
+    def encrypt_stream(
+        self,
+        source: IO,
+        dest: IO,
+        mode: FileMode,
+        nonce: Optional[bytes] = None,
+    ) -> Tuple[bytes, bytes]:
+        """
+        用 AES-GCM 加密数据流.
+
+        Returns:
+            单次随机数及校验码.
+
+        Raises:
+            EncryptionError:
+        """
+        if nonce is None:
+            nonce = secrets.token_bytes(AES_NONCE_LEN)
+        elif len(nonce) != AES_NONCE_LEN:
+            raise error.ParamError("nonce", "wrong length")
+
+        try:
+            aes = Cipher(algorithms.AES(self.key), modes.GCM(nonce)).encryptor()
+
+            if mode == "b":
+                source_b: IO[bytes] = source
+                dest_b: IO[bytes] = dest
+                while plaintext := source_b.read(CHUNK_SIZE):
+                    ciphertext = aes.update(plaintext)
+                    dest_b.write(ciphertext)
+            else:
+                source_t: IO[str] = source
+                dest_t: IO[str] = dest
+                for plaintext in source_t:
+                    plaintext = plaintext.rstrip("\n")
+                    ciphertext = aes.update(plaintext.encode())
+                    dest_t.write(base64.b64encode(ciphertext).decode())
+                    dest_t.write("\n")
+
+            ciphertext = aes.finalize()
+            mac = aes.tag
+        except Exception as e:
+            raise error.EncryptionError("AES encryption failed") from e
+
+        return nonce, mac
+
+    def decrypt(self, nonce: bytes, ciphertext: bytes, mac: Optional[bytes]) -> bytes:
+        """
+        用 AES-GCM 解密并校验数据.
+
+        Raises:
+            DecryptionError:
+        """
+        with BytesIO(ciphertext) as source, BytesIO() as dest:
+            self.decrypt_stream(source, dest, nonce, mac, "b")
+            plaintext = dest.getvalue()
+
+        return plaintext
+
+    def decrypt_stream(
+        self,
+        source: IO,
+        dest: IO,
+        nonce: bytes,
+        mac: Optional[bytes],
+        mode: FileMode,
+    ) -> None:
+        """
+        用 AES-GCM 解密并校验数据流.
+
+        Raises:
+            DecryptionError:
+        """
+        try:
+            aes = Cipher(algorithms.AES(self.key), modes.GCM(nonce)).decryptor()
+
+            if mode == "b":
+                source_b: IO[bytes] = source
+                dest_b: IO[bytes] = dest
+                while ciphertext := source_b.read(CHUNK_SIZE):
+                    plaintext = aes.update(ciphertext)
+                    dest_b.write(plaintext)
+            else:
+                source_t: IO[str] = source
+                dest_t: IO[str] = dest
+                for ciphertext in source_t:
+                    ciphertext = ciphertext.rstrip("\n")
+                    plaintext = aes.update(base64.b64decode(ciphertext)).decode()
+                    dest_t.write(plaintext)
+                    dest_t.write("\n")
+
+            if mac is not None:
+                aes.finalize_with_tag(mac)
+        except Exception as e:
+            raise error.DecryptionError("AES decryption failed") from e

src/bytedance/jeddak_secure_channel/crypto/envelope.py

@@ -0,0 +1,395 @@
+"""
+信封加密算法.
+"""
+
+__all__ = [
+    "ClientSessionKey",
+    "EncryptedMessage",
+    "ResponseKey",
+    "ServerSessionKey",
+]
+
+import base64
+import dataclasses
+import json
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from io import BytesIO
+from threading import Lock
+
+from typing_extensions import IO, TYPE_CHECKING, Optional, Self, Set, Tuple, Union
+
+if TYPE_CHECKING:
+    from _typeshed import GenericPath
+else:
+    GenericPath = str
+
+
+from .. import error
+from .aes import AES_KEY_LEN, AES_MAC_LEN, AES_NONCE_LEN, AesKey, FileMode
+from .rsa import PrivateKey, PublicKey
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(eq=False, frozen=True)
+class EncryptedMessage:
+    """
+    信封加密数据.
+    可作为数据发送方的请求, 及数据接收方的响应.
+
+    本类不含可变状态, 所有方法是线程安全的.
+    """
+
+    nonce: bytes
+    """单次随机数."""
+
+    mac: bytes
+    """校验码."""
+
+    key: Optional[bytes] = None
+    """
+    用公钥加密的对称密钥.
+    在数据发送方的加密请求中使用, 在数据接收方的加密响应中为空.
+    """
+
+    ciphertext: Optional[bytes] = None
+    """
+    加密数据.
+    在字符串 / 二进制数据加密中使用, 在文件加密中为空.
+    """
+
+    def serialize(self) -> str:
+        """
+        将消息序列化为字符串.
+        """
+        return json.dumps(
+            {
+                k: base64.b64encode(v).decode()
+                for k, v in self.__dict__.items()
+                if v is not None
+            },
+            separators=(",", ":"),
+        )
+
+    @classmethod
+    def deserialize(cls, data: Union[str, bytes]) -> Self:
+        """
+        从字符串反序列化消息.
+        """
+        try:
+            return cls(**{k: base64.b64decode(v) for k, v in json.loads(data).items()})
+        except Exception:
+            raise error.DecryptionError("Invalid encoding") from None
+
+
+@dataclass(eq=False, frozen=True)
+class ResponseKey:
+    """
+    对响应加解密的密钥.
+
+    本类不含可变状态, 所有方法是线程安全的.
+    """
+
+    _key: AesKey
+
+    def encrypt(self, plaintext: Union[str, bytes]) -> str:
+        """
+        加密响应 (由服务端使用).
+        """
+        if isinstance(plaintext, str):
+            plaintext = plaintext.encode("utf-8")
+
+        nonce, ciphertext, mac = self._key.encrypt(plaintext)
+        return EncryptedMessage(nonce=nonce, mac=mac, ciphertext=ciphertext).serialize()
+
+    def decrypt(self, response: Union[str, bytes]) -> bytes:
+        """
+        解密响应 (由客户端使用).
+        """
+        message = EncryptedMessage.deserialize(response)
+
+        if message.ciphertext is None or message.key is not None:
+            raise error.ParamError("response", "Message is not server response")
+        if len(message.nonce) != AES_NONCE_LEN or len(message.mac) != AES_MAC_LEN:
+            raise error.DecryptionError("Invalid data length")
+
+        return self._key.decrypt(message.nonce, message.ciphertext, message.mac)
+
+    def encrypt_file(
+        self,
+        source_path: GenericPath,
+        dest_path: GenericPath,
+        mode: FileMode,
+        nonce: Optional[bytes] = None,
+    ) -> str:
+        """
+        加密响应文件 (由服务端使用).
+        """
+        with open(source_path, f"r{mode}") as source, open(
+            dest_path, f"w{mode}"
+        ) as dest:
+            nonce, mac = self._key.encrypt_stream(source, dest, mode, nonce)
+
+            metadata = EncryptedMessage(nonce=nonce, mac=mac)
+
+            return metadata.serialize()
+
+    def decrypt_file(
+        self,
+        metadata: str,
+        source_path: GenericPath,
+        dest_path: GenericPath,
+        mode: FileMode,
+    ):
+        """
+        解密响应文件 (由客户端使用).
+        """
+        message = EncryptedMessage.deserialize(metadata)
+
+        if message.ciphertext is not None or message.key is not None:
+            raise error.ParamError("response", "Message is not server response")
+        if len(message.nonce) != AES_NONCE_LEN or len(message.mac) != AES_MAC_LEN:
+            raise error.DecryptionError("Invalid data length")
+
+        return self.decrypt_file_with_nonce(
+            source_path, dest_path, message.nonce, message.mac, mode
+        )
+
+    def decrypt_file_with_nonce(
+        self,
+        source_path: GenericPath,
+        dest_path: GenericPath,
+        nonce: bytes,
+        mac: Optional[bytes],
+        mode: FileMode,
+    ):
+        """
+        解密响应文件 (由客户端使用).
+        """
+        with open(source_path, f"r{mode}") as source, open(
+            dest_path, f"w{mode}"
+        ) as dest:
+            self._key.decrypt_stream(source, dest, nonce, mac, mode)
+
+
+@dataclass(eq=False, frozen=True)
+class ClientSessionKey:
+    """
+    数据发送方持有的密钥.
+
+    本类不含可变状态, 所有方法是线程安全的.
+    """
+
+    server_public_key: PublicKey
+    """数据接收方公钥."""
+
+    @classmethod
+    def load(cls, public_pem: Union[str, bytes]) -> Self:
+        """
+        通过数据接收方的公钥产生发送方密钥.
+
+        Args:
+            private_pem: PEM 格式的 RSA 公钥.
+        """
+        public_key = PublicKey.from_public_pem(public_pem)
+
+        return cls(public_key)
+
+    def encrypt(
+        self, plaintext: Union[str, bytes], key: Optional[AesKey] = None
+    ) -> EncryptedMessage:
+        """
+        信封加密.
+
+        Args:
+            key: 指定加密用的对称密钥. 默认为随机生成 (一次一密).
+        Returns:
+            信封加密数据.
+        """
+        if isinstance(plaintext, str):
+            plaintext = plaintext.encode("utf-8")
+
+        with BytesIO(plaintext) as source, BytesIO() as dest:
+            message = self.encrypt_stream(source, dest, "b", key)
+            message = dataclasses.replace(message, ciphertext=dest.getvalue())
+
+        return message
+
+    def encrypt_with_response(
+        self, plaintext: Union[str, bytes], aes_key: AesKey = None
+    ) -> Tuple[EncryptedMessage, ResponseKey]:
+        """
+        信封加密, 并获取用于解密响应的密钥.
+        """
+        if aes_key is None:
+            key = AesKey.generate()
+        else:
+            key = aes_key
+
+        return self.encrypt(plaintext, key), ResponseKey(key)
+
+    def encrypt_stream(
+        self, source: IO, dest: IO, mode: FileMode, key: Optional[AesKey] = None
+    ) -> EncryptedMessage:
+        """
+        信封加密数据流.
+
+        Returns:
+            信封加密元数据.
+        """
+        symmetric_key = key or AesKey.generate()
+
+        nonce, mac = symmetric_key.encrypt_stream(source, dest, mode)
+
+        encrypted_key = self.server_public_key.encrypt(symmetric_key.key)
+
+        return EncryptedMessage(key=encrypted_key, nonce=nonce, mac=mac)
+
+    def encrypt_stream_with_response(
+        self, source: IO, dest: IO, mode: FileMode
+    ) -> Tuple[EncryptedMessage, ResponseKey]:
+        """
+        信封加密数据流, 并获取用于解密响应的密钥.
+        """
+        key = AesKey.generate()
+
+        # return self.encrypt_stream(source, dest, mode), ResponseKey(key)
+        return self.encrypt_stream(source, dest, mode, key), ResponseKey(key)
+
+
+@dataclass(eq=False)
+class ServerSessionKey:
+    """
+    数据接收方持有的密钥.
+
+    本类的所有方法是线程安全的.
+    """
+
+    private_key: PrivateKey
+    """私钥."""
+
+    expire_at: datetime
+    """【未启用】密钥对在此时间后无效."""
+
+    tks_version: Optional[int]
+    """【未启用】密钥版本, 用于判断获取的密钥是否变化. """
+
+    pem_md5: Optional[str] = None
+    """pem的MD5值."""
+
+    used_nonce: Set[bytes] = field(default_factory=set)
+    """已使用的单次随机数列表."""
+
+    used_nonce_lock: Lock = field(init=False, default_factory=Lock)
+    """线程安全: 修改 used_nonce 应该加锁."""
+
+    @classmethod
+    def load(
+        cls,
+        private_pem: Union[str, bytes],
+        expire_at: datetime,
+        tks_version: Optional[int],
+        pem_md5: Optional[str],
+    ) -> Self:
+        """
+        通过 PEM 格式的 RSA 私钥产生接收方密钥.
+
+        Args:
+            private_pem: PEM 格式的 RSA 私钥.
+            :param pem_md5:
+            :param private_pem:
+            :param expire_at:
+            :param tks_version:
+        """
+        private_key = PrivateKey.from_private_pem(private_pem)
+
+        return cls(private_key, expire_at, tks_version, pem_md5)
+
+    @classmethod
+    def generate(cls, expire_at: datetime) -> Self:
+        """
+        随机生成接收方密钥.
+        """
+        private_key = PrivateKey.generate()
+
+        return cls(private_key, expire_at, tks_version=None, pem_md5=None)
+
+    def is_expired(self) -> bool:
+        """
+        是否已超过有效期.
+        """
+        now = datetime.now(timezone.utc)
+        if now > self.expire_at:
+            logger.warning("session key maybe expired")
+
+        return False  # 服务端密钥暂时不考虑有效期
+
+    def get_public_key(self) -> str:
+        """
+        获取 PEM 格式的公钥.
+        """
+        return self.private_key.public_pem()
+
+    def decrypt(
+        self, message: Union[str, bytes, EncryptedMessage]
+    ) -> Tuple[bytes, ResponseKey]:
+        """
+        信封解密.
+
+        Returns:
+            明文数据, 及用于加密响应的密钥.
+
+        Raises:
+            DecryptionError: 加密消息不合法, 或者加密使用的密钥不是本密钥.
+        """
+        if isinstance(message, (str, bytes)):
+            message = EncryptedMessage.deserialize(message)
+        if message.ciphertext is None:
+            raise error.ParamError("message", "Message has no data")
+
+        with BytesIO(message.ciphertext) as source, BytesIO() as dest:
+            response_key = self.decrypt_stream(message, source, dest, "b")
+            plaintext = dest.getvalue()
+
+        return plaintext, response_key
+
+    def decrypt_stream(
+        self,
+        metadata: Union[str, bytes, EncryptedMessage],
+        source: IO,
+        dest: IO,
+        mode: FileMode,
+    ) -> ResponseKey:
+        """
+        信封解密数据流.
+
+        Returns:
+            用于加密响应的密钥.
+
+        Raises:
+            DecryptionError: 加密消息不合法, 或者加密使用的密钥不是本密钥.
+        """
+        if isinstance(metadata, (str, bytes)):
+            metadata = EncryptedMessage.deserialize(metadata)
+        if metadata.key is None:
+            raise error.ParamError("metadata", "Message has no key")
+
+        # with self.used_nonce_lock:
+        #     if metadata.nonce in self.used_nonce:
+        #         raise error.NonceReusedError()
+        #     self.used_nonce.add(metadata.nonce)
+
+        symmetric_key = self.private_key.decrypt(metadata.key)
+        if (
+            len(symmetric_key) != AES_KEY_LEN
+            or len(metadata.nonce) != AES_NONCE_LEN
+            or len(metadata.mac) != AES_MAC_LEN
+        ):
+            raise error.DecryptionError("Invalid data length")
+
+        symmetric_key = AesKey(symmetric_key)
+        symmetric_key.decrypt_stream(source, dest, metadata.nonce, metadata.mac, mode)
+
+        return ResponseKey(symmetric_key)

src/bytedance/jeddak_secure_channel/crypto/rsa.py

@@ -0,0 +1,164 @@
+"""
+RSA 加密和签名算法.
+"""
+
+__all__ = ["PublicKey", "PrivateKey"]
+
+from dataclasses import dataclass
+
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
+from typing_extensions import Self, Union
+
+from .. import error
+
+ENCRYPT_PADDING = padding.OAEP(
+    mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None
+)
+SIGN_PADDING = padding.PSS(
+    mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
+)
+SIGN_HASH = hashes.SHA256()
+
+
+@dataclass(eq=False, frozen=True)
+class PublicKey:
+    """
+    RSA 公钥.
+
+    本类不含可变状态, 所有方法是线程安全的.
+    """
+
+    public_key: RSAPublicKey
+
+    @classmethod
+    def from_public_pem(cls, public_pem: Union[str, bytes]) -> Self:
+        """
+        读取 PEM 格式的 RSA 公钥.
+
+        Args:
+            private_pem: PEM 格式的 RSA 公钥.
+        """
+        if isinstance(public_pem, str):
+            public_pem = public_pem.encode("utf-8")
+        try:
+            public_key = serialization.load_pem_public_key(public_pem)
+        except Exception as e:
+            raise error.ParamError("public_pem", "Invalid RSA public key") from e
+        if not isinstance(public_key, RSAPublicKey):
+            raise error.ParamError("public_pem", "Key is not RSA")
+
+        return cls(public_key)
+
+    def encrypt(self, plaintext: bytes) -> bytes:
+        """
+        RSA 加密.
+
+        Raises:
+            EncryptionError:
+        """
+        try:
+            return self.public_key.encrypt(plaintext, ENCRYPT_PADDING)
+        except Exception as e:
+            raise error.EncryptionError("RSA encryption failed") from e
+
+    def public_pem(self) -> str:
+        """
+        获取 PEM 格式的公钥.
+        """
+        return self.public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        ).decode("utf-8")
+
+    def verify(self, signature: bytes, message: Union[str, bytes]) -> bool:
+        """
+        公钥验签.
+
+        Raises:
+            SignatureError:
+        """
+        if isinstance(message, str):
+            message = message.encode("utf-8")
+        try:
+            self.public_key.verify(signature, message, SIGN_PADDING, SIGN_HASH)
+            return True
+        except Exception as e:
+            print(e)
+            return False
+
+
+@dataclass(init=False, eq=False, frozen=True)
+class PrivateKey(PublicKey):
+    """
+    RSA 密钥对.
+    """
+
+    private_key: RSAPrivateKey
+
+    def __init__(self, private_key: RSAPrivateKey):
+        # Workaround frozen class
+        object.__setattr__(self, "private_key", private_key)
+        object.__setattr__(self, "public_key", private_key.public_key())
+
+    @classmethod
+    def from_private_pem(cls, private_pem: Union[str, bytes]) -> Self:
+        """
+        读取 PEM 格式的 RSA 密钥对.
+
+        Args:
+            private_pem: PEM 格式的 RSA 私钥.
+        """
+        if isinstance(private_pem, str):
+            private_pem = private_pem.encode("utf-8")
+        private_key = serialization.load_pem_private_key(private_pem, password=None)
+        if not isinstance(private_key, RSAPrivateKey):
+            raise error.ParamError("private_pem", "Key is not RSA")
+
+        return cls(private_key)
+
+    @classmethod
+    def generate(cls) -> Self:
+        """
+        随机生成密钥对.
+        """
+        private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096)
+        return cls(private_key)
+
+    def private_pem(self) -> str:
+        """
+        获取 PEM 格式的私钥.
+        """
+        return self.private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        ).decode("utf-8")
+
+    def decrypt(self, plaintext: bytes) -> bytes:
+        """
+        RSA 解密.
+
+        Raises:
+            DecryptionError:
+        """
+        try:
+            return self.private_key.decrypt(plaintext, ENCRYPT_PADDING)
+        except Exception as e:
+            raise error.DecryptionError("RSA decryption failed") from e
+
+    def sign(self, message: Union[str, bytes]) -> bytes:
+        """
+        私钥签名.
+
+        Raises:
+            SignatureError:
+        """
+        if isinstance(message, str):
+            message = message.encode("utf-8")
+        try:
+            signature = self.private_key.sign(message, SIGN_PADDING, SIGN_HASH)
+            return signature
+        except Exception as e:
+            raise error.SignatureError("RSA signing failed") from e

src/bytedance/jeddak_secure_channel/error.py

@@ -0,0 +1,131 @@
+"""
+安全通信错误类型.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from typing_extensions import Literal, Optional
+
+
+class ErrorKind(Enum):
+    CONFIG = "config"
+    PARAM = "param"
+
+    NETWORK = "network"
+    SERVICE = "service"
+
+    ENCRYPTION = "encryption"
+    DECRYPTION = "decryption"
+    SIGNATURE = "signature"
+    KEY_MISSING = "key_missing"
+    KEY_EXPIRED = "key_expired"
+    NONCE_REUSED = "nonce_reused"
+
+
+class SecureChannelError(Exception):
+    kind: ErrorKind
+
+
+@dataclass
+class ConfigError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.CONFIG)
+
+    config_object: object
+    message: str
+
+    def __str__(self) -> str:
+        return f"{self.config_object}: {self.message}"
+
+
+@dataclass
+class ParamError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.PARAM)
+
+    param: str
+    message: str
+
+    def __str__(self) -> str:
+        return f"{self.param}: {self.message}"
+
+
+@dataclass
+class NetworkError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.NETWORK)
+
+    service_name: Literal["EPS", "RA", "RAS", "TKS"]
+    service: str
+    endpoint: Optional[str]
+
+    def _service(self) -> str:
+        if self.endpoint is None:
+            return f"{self.service_name} service {self.service}"
+        else:
+            return f"endpoint {self.endpoint} of {self.service_name} service {self.service}"
+
+    def __str__(self) -> str:
+        return f"Cannot access {self._service()}"
+
+
+@dataclass
+class ServiceError(NetworkError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.SERVICE)
+
+    message: str
+
+    def __str__(self) -> str:
+        return f"{self._service().capitalize()} returned error: {self.message}"
+
+
+@dataclass
+class EncryptionError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.ENCRYPTION)
+
+    message: str
+
+    def __str__(self) -> str:
+        return self.message
+
+
+@dataclass
+class DecryptionError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.DECRYPTION)
+
+    message: str
+
+    def __str__(self) -> str:
+        return self.message
+
+
+@dataclass
+class SignatureError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.SIGNATURE)
+
+    message: str
+
+    def __str__(self) -> str:
+        return self.message
+
+
+@dataclass
+class KeyMissingError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.KEY_MISSING)
+
+    def __str__(self) -> str:
+        return "Can not get key, Please check your config !"
+
+
+@dataclass
+class KeyExpiredError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.KEY_EXPIRED)
+
+    def __str__(self) -> str:
+        return "Key has expired"
+
+
+@dataclass
+class NonceReusedError(SecureChannelError):
+    kind: ErrorKind = field(init=False, repr=False, default=ErrorKind.NONCE_REUSED)
+
+    def __str__(self) -> str:
+        return "Nonce is reused"

src/bytedance/jeddak_secure_channel/log/__init__.py

@@ -0,0 +1,24 @@
+__all__ = [
+    "DEBUG",
+    "INFO",
+    "WARNING",
+    "ERROR",
+    "CRITICAL",
+    "FORMAT",
+    "LogConfig",
+    "debug",
+    "info",
+    "warning",
+    "error",
+    "critical",
+    "init_log_config"
+]
+
+from .config import DEBUG, INFO, WARNING, ERROR, CRITICAL, FORMAT, LogConfig
+from .logger import init_log_config, info, debug, warning, error, critical
+
+_initialized = False
+if not _initialized:
+    config = LogConfig()
+    init_log_config(config)
+    _initialized = True

src/bytedance/jeddak_secure_channel/log/config.py

@@ -0,0 +1,31 @@
+__all__ = [
+    "DEBUG",
+    "INFO",
+    "WARNING",
+    "ERROR",
+    "CRITICAL",
+    "FORMAT",
+    "LogConfig"
+]
+
+# Log level constants
+DEBUG = "DEBUG"
+INFO = "INFO"
+WARNING = "WARNING"
+ERROR = "ERROR"
+CRITICAL = "CRITICAL"
+# Log output format
+FORMAT = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {file.path}:{line} | {message}"
+
+
+class LogConfig:
+    """Log configuration class."""
+
+    def __init__(self):
+        self.rotation = 100  # Size limit for log files in MB
+        self.retention = 7  # Days to retain log files
+        self.compression = "zip"
+        self.dir = "jsc_log"
+        self.filename = "jsc.log"
+        self.format = FORMAT
+        self.level = INFO

src/bytedance/jeddak_secure_channel/log/logger.py

@@ -0,0 +1,88 @@
+__all__ = [
+    "debug",
+    "info",
+    "warning",
+    "error",
+    "critical",
+    "init_log_config"
+]
+
+import os.path
+import sys
+from typing import Any
+
+from loguru import logger
+
+from . import LogConfig, FORMAT
+
+
+def init_log_config(*configs: LogConfig):
+    """Initialize logging configuration with given configs."""
+    logger.remove()
+
+    log_level = "INFO"
+    if os.environ.get("log_level"):
+        log_level = os.environ.get("log_level")
+    if os.environ.get("LOG_LEVEL"):
+        log_level = os.environ.get("LOG_LEVEL")
+
+    logger.add(sink=sys.stdout, format=FORMAT, level=log_level)
+
+    for config in configs:
+        if not config.dir:
+            config.dir = os.environ.get("LOG_DIR", "aicc_log")
+
+        if not config.filename:
+            config.filename = "jsc.log"
+
+        os.makedirs(config.dir, exist_ok=True)
+        fullpath = os.path.join(config.dir, config.filename)
+
+        if config.rotation <= 0:
+            config.rotation = 100
+        rotation = f"{config.rotation} MB"
+
+        if config.retention <= 0:
+            config.retention = 7
+        retention = f"{config.retention} days"
+
+        logger.add(
+            sink=fullpath,
+            rotation=rotation,
+            retention=retention,
+            compression=config.compression,
+            format=config.format,
+            level=config.level
+        )
+
+
+def msg_proc(message: str) -> str:
+    if message:
+        return f"[jsc]:{message}"
+    else:
+        return message
+
+
+def debug(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log debug message."""
+    logger.opt(depth=1).debug(msg_proc(message), *args, **kwargs)
+
+
+def info(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log info message."""
+    logger.opt(depth=1).info(msg_proc(message), *args, **kwargs)
+
+
+def warning(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log warning message."""
+    logger.opt(depth=1).warning(msg_proc(message), *args, **kwargs)
+
+
+def error(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log error message."""
+    logger.opt(depth=1).error(msg_proc(message), *args, **kwargs)
+
+
+def critical(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log critical message."""
+    logger.opt(depth=1).critical(msg_proc(message), *args, **kwargs)

src/bytedance/jeddak_secure_channel/logger_factory.py

@@ -0,0 +1,134 @@
+import logging
+import multiprocessing
+import os
+import sys
+
+mutex = multiprocessing.Lock()
+
+
+class Factory(object):
+    @staticmethod
+    def init(*args):
+        """
+        Initialize a global instance
+        :param args:
+        :return:
+        """
+        pass
+
+    @staticmethod
+    def get_global_instance():
+        """
+        Get the global instance
+        :return:
+        """
+        pass
+
+    @staticmethod
+    def get_instance(*args):
+        """
+        Get an instance
+        :param args:
+        :return:
+        """
+        pass
+
+    @staticmethod
+    def _raise_value_error(invalid_type, invalid_value):
+        """
+
+        :param invalid_type: str
+        :param invalid_value: str
+        :return:
+        """
+        raise ValueError("invalid {}: {}".format(invalid_type, invalid_value))
+
+
+class LoggerFactory(Factory):
+    logger = None
+    root_dir = ""
+    task_logger_level = ""
+    local_logger = {}
+
+    @staticmethod
+    def init():
+        top_logger_level = os.getenv("TOP_LOGGER_LEVEL", "DEBUG").strip('"')
+        log_path = os.getenv("LOGDIR", f"./jsc_log").strip('"')
+        task_logger_level = os.getenv("TASK_LOGGER_LEVEL", "DEBUG").strip('"')
+
+        LoggerFactory.task_logger_level = task_logger_level
+
+        try:
+            import bytedlogger
+
+            bytedlogger.config_default()
+            LoggerFactory.logger = logging.getLogger("top")
+        except ImportError:
+            LoggerFactory.logger = logging.getLogger("top")
+        LoggerFactory.logger.setLevel(top_logger_level)
+
+        if log_path is None or log_path == "":
+            log_path = "/home/jeddak/log"
+            print(f"Root logger path is not set, use log_path {log_path}")
+
+        print(f"log_path: {log_path}")
+        LoggerFactory.root_dir = log_path
+
+        if not os.path.exists(log_path):
+            os.makedirs(log_path)
+
+        top_logger_path = os.path.join(log_path, "top.log")
+
+        fh = logging.FileHandler(top_logger_path, mode="a")
+        fh.setLevel(top_logger_level)
+        formatter = logging.Formatter(
+            "%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s"
+        )
+        fh.setFormatter(formatter)
+        LoggerFactory.logger.addHandler(fh)
+        # LoggerFactory.logger.addHandler(logging.StreamHandler(sys.stdout))
+
+    @staticmethod
+    def get_global_instance():
+        if LoggerFactory.logger is None:
+            LoggerFactory.init()
+        return LoggerFactory.logger
+
+    @staticmethod
+    def attach_logger(logger, log_path, log_file_name, logger_level, _formatter):
+        with mutex:
+            if not os.path.exists(log_path):
+                os.makedirs(log_path)
+        logger.setLevel(logger_level)
+        _fh = logging.FileHandler(log_file_name, mode="a")
+        _fh.setLevel(logger_level)
+        # _formatter = logging.Formatter(
+        #     f'%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - {logger_str} - %(message)s')
+        _fh.setFormatter(_formatter)
+
+        logger.addHandler(_fh)
+        logger.addHandler(logging.StreamHandler(sys.stdout))
+        return logger
+
+    @staticmethod
+    def get_local_instance(log_id="jsc_log"):
+        if LoggerFactory.logger is None:
+            LoggerFactory.init()
+
+        logger_name = log_id
+        if logger_name in LoggerFactory.local_logger:
+            return LoggerFactory.local_logger[logger_name]
+
+        log_path = os.path.join(LoggerFactory.root_dir, log_id)
+        log_file_name = os.path.join(log_path, logger_name + ".log")
+
+        _logger = logging.getLogger(log_id)
+
+        _formatter = logging.Formatter(
+            f"%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - {logger_name} - %(message)s"
+        )
+        _logger = LoggerFactory.attach_logger(
+            _logger, log_path, log_file_name, LoggerFactory.task_logger_level, _formatter
+        )
+        LoggerFactory.local_logger[logger_name] = _logger
+        return _logger

src/bytedance/jeddak_secure_channel/ra.py

@@ -0,0 +1,768 @@
+"""
+安全通信远程证明功能.
+"""
+
+__all__ = [
+    "RaConfig",
+    "RaPod",
+    "RaRequest",
+    "RaResponse",
+    "RaResponseMetadata",
+    "RaResponsePod",
+    "RaResponsePods",
+    "attest_server",
+    "prepare_ra_request",
+    "validata_ra_request",
+    "attest_client",
+    "generate_nonce",
+]
+
+import time
+import uuid
+import base64
+import json
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Tuple
+import jwt
+import hashlib
+from pymerkle import verify_inclusion, MerkleProof
+from jwcrypto import jwk
+from jwt.exceptions import ExpiredSignatureError, InvalidIssuerError, InvalidSignatureError
+from volcenginesdkcore.rest import ApiException
+import requests
+from typing_extensions import (
+    Dict,
+    List,
+    Literal,
+    NotRequired,
+    Optional,
+    TypedDict,
+    Union,
+)
+
+from . import error
+from .crypto import PublicKey
+from .log import logger
+from .utils import request_bytedance_gateway
+
+from cryptography.hazmat.primitives.asymmetric import rsa, padding
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.backends import default_backend
+
+
+RA_NONCE_LEN = 12
+RA_TYPE_LOCAL = "local"
+RA_TYPE_TCA = "tca"
+
+
+class RaPod(TypedDict):
+    cluster_id: str
+
+    namespace: str
+
+    deployment: str
+
+    pod: NotRequired[Literal["random", "all"]]
+    """默认为 random, 如果想要全部则为 all."""
+
+
+@dataclass
+class RaConfig:
+    """
+    访问远程证明功能的配置.
+    """
+
+    ra_url: str = ""
+    """远程证明端点的 HTTP 地址."""
+
+    ra_type: str = RA_TYPE_TCA
+    """ra的类型，local：本地RA， tca：TCA的RAS_adapter接口(默认)."""
+
+    ra_service_name: str = ""
+    """数据接收端（服务端）的服务名. ,用于验证通过机密容器部署的服务"""
+
+    ra_pods_info: str = ""
+    """
+        待验证的节点的pod信息，用于验证非通过机密容器部署的服务, 
+        格式: {cluster_id:***, namespace:***, deployment:***}
+        ra_service_name 和 ra_pods_info 二选一
+    """
+
+    ra_uid: str = ""
+    """用户的UID"""
+
+    ra_key_negotiation: bool = True
+    """是否在验证的时候需要同步进行协商密钥."""
+
+    ra_need_token: bool = True  # RA 功能稳定后要改成 False
+    """是否要返回验证 token. 默认为是."""
+
+    attest_must: bool = False
+    """是否一定要返回验证报告才能采用公钥."""
+
+    ra_policy_id: str = ""
+    """进行验证所使用的策略, 如果 ra_need_token 为 True, 则必填."""
+
+    bytedance_top_info: str = ""
+    """访问bytedance TOP网关的信息."""
+
+    ra_attested_pods: List[RaPod] = field(default_factory=list)
+    """old 需要验证的节点信息."""
+
+    ra_policy_ids: List[str] = field(default_factory=list)
+    """old 进行验证所使用的策略, 如果 ra_need_token 为 True, 则必填."""
+
+    auth_token: str = ""
+    """用来进行鉴权的凭证"""
+
+
+class RaRequest(TypedDict):
+    """数据发送方发出的远程证明请求."""
+
+    key_negotiation: bool
+    """是否在验证的时候需要同步进行协商密钥."""
+
+    token: bool
+    """是否要返回验证 token."""
+
+    policy_ids: List[str]
+    """进行验证所使用的策略."""
+
+    attested_pods: List[RaPod]
+    """需要验证的节点信息."""
+
+
+class RaResponseMetadata(TypedDict):
+    Code: int
+    Message: str
+
+
+class RaResponsePod(TypedDict):
+    evidence: str
+    token: str
+    key_info: dict
+
+
+RaResponsePods = Dict[str, RaResponsePod]
+
+
+class RaResponse(TypedDict):
+    """数据接收方发出的远程证明响应."""
+
+    ResponseMetadata: RaResponseMetadata
+
+    ResponseResult: RaResponsePods
+
+
+def prepare_ra_request(config: RaConfig) -> RaRequest:
+    """
+    准备远程证明请求.
+
+    Returns:
+        远程证明请求. 可作为 JSON 发送给数据接收方远程证明端点.
+    """
+    return RaRequest(
+        key_negotiation=config.ra_key_negotiation,
+        token=config.ra_need_token,
+        policy_ids=config.ra_policy_ids,
+        attested_pods=config.ra_attested_pods,
+    )
+
+
+def validata_ra_request(request: RaRequest) -> None:
+    """
+    验证远程证明请求有效性.
+
+    Raises:
+        ParamError: 请求的 nonce 或 timestamp 字段无效.
+    """
+    nonce = request.get("hex_runtime_data")
+    if nonce:
+        nonce = base64.b64decode(nonce)
+        if len(nonce) != RA_NONCE_LEN:
+            raise error.ParamError("nonce", "Invalid RA request")
+
+    # time = datetime.fromtimestamp(request["timestamp"], timezone.utc)
+    # now = datetime.now(timezone.utc)
+    # if time > now + timedelta(minutes=1) or time < now - timedelta(minutes=2):
+    #     raise error.ParamError("timestamp", "Invalid RA request")
+
+
+def generate_nonce():
+    return ''.join(str(uuid.uuid4()).split('-'))[-16:]
+
+
+def request_tca(config: RaConfig, nonce: str = None):
+    try:
+        ra_url = config.ra_url
+        ra_service_name = config.ra_service_name
+        ra_pods_info = config.ra_pods_info
+        ra_key_negotiation = config.ra_key_negotiation
+        ra_policy_id = config.ra_policy_id
+        ra_uid = config.ra_uid
+
+        top_info = config.bytedance_top_info
+        if top_info and isinstance(top_info, str):
+            top_info = json.loads(top_info)
+        # assert isinstance(top_info, dict)
+
+        headers = {"UID": ra_uid}
+
+        if not nonce:
+            nonce = generate_nonce()
+
+        if ra_service_name:
+            body = {
+                "Nonce": nonce,
+                "PolicyID": ra_policy_id,
+                "ServiceName": ra_service_name,
+                "KeyNegotiation": ra_key_negotiation,
+                "Token": True,
+                "Version": 2
+            }
+        elif ra_pods_info:
+            if isinstance(ra_pods_info, str):
+                ra_pods_info = json.loads(ra_pods_info)
+            body = {
+                "Nonce": nonce,
+                "PolicyID": ra_policy_id,
+                "AttestedPods": [ra_pods_info],
+                "KeyNegotiation": ra_key_negotiation,
+                "Token": True,
+                "Version": 2
+            }
+        else:
+            raise error.ServiceError("RA", config.ra_url, None, "call TCA error")
+
+        if top_info.get("target_uid") and not top_info.get("aicc_saas_trn"):
+            from bytedance.jeddak_secure_channel.utils import top_request
+            return top_request(ra_url, body, headers, top_info)
+        else:
+            return request_bytedance_gateway(ra_url, body, headers, top_info)
+    except Exception as e:
+        logger.critical(f"Response is not JSON: service={config.ra_url}")
+        raise error.ServiceError("RA", config.ra_url, None, "call TCA error") from e
+
+
+def attest_server_local(token: Optional[str], config: RaConfig, nonce=None):
+    """
+    对数据接收方进行远程证明(非TCA方式).
+    """
+    timestamp = str(int(datetime.now(timezone.utc).timestamp()))
+    auth_token = token
+    if not auth_token:
+        auth_token = config.auth_token
+
+    headers = {
+        "Content-Type": "application/json",
+        "timestamp": f"{timestamp}",
+        "Authorization": f"Bearer {auth_token}"
+    }
+
+    if not nonce:
+        nonce = generate_nonce()
+
+    request = {"Nonce": nonce}
+    try:
+        logger.info(
+            f"request ra, url={config.ra_url}, header={headers}, params={request}"
+        )
+        response = requests.post(config.ra_url, headers=headers, json=request)
+        return response
+
+    except Exception as e:
+        logger.critical(f"RA Network error: url={config.ra_url}")
+        raise error.NetworkError("RA", config.ra_url, endpoint=None) from e
+
+    # if response.status_code != 200:
+    #     logger.error(
+    #         f"Service error: service={config.ra_url} code={response.status_code}"
+    #     )
+    #     raise error.ServiceError(
+    #         "RA", config.ra_url, None, f"code={response.status_code}"
+    #     )
+
+    # try:
+    #     response_json: RaResponse = response.json()
+    #     logger.info(f"ra response={response_json}")
+    #
+    #     res_meta_data = response_json.get("ResponseMetadata")
+    #     res_result = response_json.get("ResponseResult")
+    #
+    #     if not res_meta_data or not res_result:
+    #         raise error.ServiceError("RA", config.ra_url, None, "wrong format")
+    #
+    #     if res_meta_data.get("Code") != 0:
+    #         raise error.ServiceError(
+    #             "RA", config.ra_url, None, res_meta_data.get("Message", "")
+    #         )
+    #
+    #     for pod_name, pod_ra_info in res_result.items():
+    #         # evidence = pod_ra_info.get("evidence")
+    #         # token = pod_ra_info.get("token")
+    #         key_info = pod_ra_info.get("key_info")
+    #         if not key_info:
+    #             raise error.ServiceError("RA", config.ra_url, "RSA key", "null")
+    #
+    #         return True, key_info["pub_key_info"]
+    # except Exception as e:
+    #     logger.critical(f"Response is not JSON: service={config.ra_url}")
+    #     raise error.ServiceError("RA", config.ra_url, None, "not JSON") from e
+    #
+    # return False, ""
+
+
+def verify_jwt_token(token):
+    try:
+        tmp_payload = token.split(".")[1]
+        jwk_data = json.loads(base64.urlsafe_b64decode(tmp_payload + "=" * (4 - len(tmp_payload) % 4))).get('jwk')
+
+        public_key = jwk.JWK(**jwk_data)
+
+        # 转换为 PEM 格式，以便 PyJWT 使用
+        pem_public_key = public_key.export_to_pem()
+
+        header = jwt.get_unverified_header(token)
+
+        time.sleep(2)
+
+        # 验证 Token
+        payload = jwt.decode(token, pem_public_key, algorithms=[header.get('alg')],
+                             issuer="Bytedance-Remote-Attestation-Service",
+                             options={"verify_signature": True})
+
+        logger.debug(f"Token is valid. Payload:{payload}")
+        return True
+    except ExpiredSignatureError:
+        logger.warning("verify_ra_token: Token has expired.")
+        return False
+    except InvalidIssuerError:
+        logger.warning("verify_ra_token: Invalid issuer.")
+        return False
+    except InvalidSignatureError:
+        logger.warning("verify_ra_token: Invalid signature.")
+        return False
+    except Exception as e:
+        logger.warning(f"verify_ra_token: {e}")
+        return False
+
+
+def base64url_decode(encoded_str):
+    padded_str = encoded_str + '=' * ((4 - len(encoded_str) % 4) % 4)  # 补齐padding
+    return base64.urlsafe_b64decode(padded_str)
+
+
+def jwk_to_rsa_public_key(jwk_info: dict):
+    """将JWK格式的公钥转换为RSA公钥对象"""
+    # 从JWK中提取必要的参数
+    n = int.from_bytes(base64url_decode(jwk_info['n']), byteorder='big')
+    e = int.from_bytes(base64url_decode(jwk_info['e']), byteorder='big')
+    public_key = rsa.RSAPublicNumbers(e, n).public_key(default_backend())
+    return public_key
+
+
+def verify_jwt_token_new(token):
+    """验证JWT token，不依赖JWT库"""
+    try:
+        # 分割JWT：header.payload.signature
+        parts = token.split(".")
+        if len(parts) != 3:
+            logger.error("verify_ra_token: Invalid token format")
+            return False
+
+        header_encoded, payload_encoded, signature_encoded = parts
+
+        header = json.loads(base64url_decode(header_encoded).decode('utf-8'))
+        payload = json.loads(base64url_decode(payload_encoded).decode('utf-8'))
+
+        # 验证token是否过期
+        if 'exp' in payload:
+            if payload['exp'] < int(time.time()):
+                logger.warning("verify_ra_token: Token has expired.")
+
+        # 验证issuer
+        if payload.get('iss') != "Bytedance-Remote-Attestation-Service":
+            logger.error("verify_ra_token: Invalid issuer.")
+            return False
+
+        # 签名验证
+        if 'jwk' in header:
+            public_key = jwk_to_rsa_public_key(header['jwk'])
+            alg = header.get("alg")
+        elif 'jwk' in payload:
+            public_key = jwk_to_rsa_public_key(payload['jwk'])
+            alg = payload.get('jwk').get("alg")
+        else:
+            public_key = None
+            alg = None
+            logger.warning("can not find pub_key from JWT token")
+
+        if public_key:
+            message = f"{header_encoded}.{payload_encoded}".encode('utf-8')
+            signature = base64url_decode(signature_encoded)
+            try:
+                if alg == "RS384":
+                    public_key.verify(signature, message, padding.PKCS1v15(), hashes.SHA384())
+                else:
+                    logger.warning(f"no support JWT sign algorithms: [{alg}]")
+                    return False
+            except Exception as e:
+                logger.error(f"JWT Token Signature verification failed: {e}")
+                return False
+
+        return True
+
+    except Exception as e:
+        logger.error(f"verify_jwt_token Exception={e}")
+        return False
+
+
+def gen_report_data(nonce_up, nonce_down, dh_param, expiration_time):
+    try:
+        nonce_up_raw = bytes.fromhex(nonce_up)
+        nonce_down_raw = base64.b16decode(nonce_down.upper())
+        dh_param_bytes = base64.b64decode(dh_param)
+        expiration_time_bytes = str(expiration_time).encode()
+        report_data_input = nonce_up_raw + nonce_down_raw + dh_param_bytes + expiration_time_bytes
+        h = hashlib.sha256()
+        h.update(report_data_input)
+        report_data = h.hexdigest()
+
+        h = hashlib.sha512()
+        h.update(report_data.encode())
+        report_data_bytes = h.digest()
+
+        return report_data_bytes
+
+    except Exception as e:
+        logger.error(f"gen report data,msg={str(e)}")
+        raise Exception(f"gen report data,msg={str(e)}")
+
+
+def verify_nonce(proof, report_data, nonce, pub_key_info):
+    try:
+        # 前64位是nonce的sha256哈希，后64位是公钥的sha256哈希
+        nonce_hash = hashlib.sha256(nonce.encode()).hexdigest()
+
+        # 计算pub_key_info的SHA256哈希
+        if pub_key_info:
+            pub_key_hash = hashlib.sha256(pub_key_info.encode()).hexdigest()
+        else:
+            # 如果公钥为空，用64个0替换
+            pub_key_hash = "0000000000000000000000000000000000000000000000000000000000000000"
+
+        # 拼接两个哈希值
+        base = nonce_hash + pub_key_hash
+        logger.debug(f'base={base}')
+
+        if not proof:
+            logger.debug(f'report_data={report_data}')
+            return report_data == base
+
+        h = hashlib.sha512()
+        h.update(b'\x00' + base.encode())
+        base = h.hexdigest()
+        root = report_data
+
+        logger.debug(f"new base={base}")
+
+        proof = MerkleProof.deserialize(proof)
+        verify_inclusion(bytes.fromhex(base), bytes.fromhex(root), proof)
+        return True
+    except Exception as e:
+        logger.debug(f"verify nonce failed!!!, msg={str(e)}")
+        return False
+
+
+def check_pubkey_valid(pk_s: str, pk_r: str, timestamp: str, version: str, service_name: str, sign: str) -> bool:
+    """
+    对pk_s的正确性进行检验
+    :param pk_s:服务公钥
+    :param pk_r:rar公钥
+    :param timestamp:时间戳
+    :param version:version
+    :param service_name:version
+    :param sign: sign(pk_s+timestamp, pr_r)
+    :return: True/False
+    """
+    if not pk_s or not pk_r or not sign:
+        return False
+
+    rsa_pub_r = PublicKey.from_public_pem(pk_r)
+
+    if isinstance(sign, str):
+        try:
+            sign = base64.b16decode(sign.upper().encode())
+        except Exception as e:
+            logger.warning(f"ra: check_pubkey_valid b16decode sign failed, Exception={e}")
+            sign = base64.b64decode(sign.encode())
+
+    msg = f"{pk_s}{timestamp}{service_name}".encode()
+    if version:
+        version_prefix = int(version).to_bytes(length=8, byteorder="big", signed=False)
+        msg = version_prefix + msg
+
+    ret = rsa_pub_r.verify(sign, msg)
+
+    logger.debug(f"check_pubkey_valid: pk_s={pk_s}, pk_r={pk_r}, version={version}, timestamp={timestamp},"
+                 f"service_name={service_name}, msg={msg}, sign={sign}, ret={ret}")
+    return ret
+
+
+def attest_server(token: Optional[str], config: RaConfig, nonce: str = None) -> Tuple[bool, str]:
+    """
+    对数据接收方进行远程证明.
+    :return (ra_status, pub_key)
+    """
+    try:
+        if config.bytedance_top_info:
+            response = request_tca(config, nonce)
+        else:
+            response = attest_server_local(token, config, nonce)
+    except ApiException as e:
+        logger.error(f"attest_server error: config={config} ApiException={e}")
+        raise error.ServiceError("RA", config.ra_url, None, "NULL")
+    except Exception as e:
+        logger.error(f"attest_server error: config={config} Exception={e}")
+        raise error.ServiceError("RA", config.ra_url, None, "NULL")
+
+    if response is None:
+        logger.error(f"RA Service error: config={config} response=NULL")
+        raise error.ServiceError("RA", config.ra_url, None, "NULL")
+
+    if isinstance(response, requests.Response):
+        if response.status_code != 200:
+            logger.error(
+                f"RA Service error: code={response.status_code}, content={response.content}, config={config}"
+            )
+            raise error.ServiceError(
+                "RA", config.ra_url, None, f"code={response.status_code}"
+            )
+        else:
+            response_text = response.text
+            ra_res = json.loads(response_text).get("Result")
+            if not ra_res:
+                ra_res = json.loads(response_text).get("ResponseResult")
+    else:
+        ra_res = response
+
+    if not ra_res:
+        logger.error(f"RA response is NULL, config={config}")
+        return False, ""
+
+    try:
+        attestation_status = False
+        public_key_info = ""
+        pub_r = ""
+
+        # 取第一个节点验证就可以
+        for cluster_info, res_info in ra_res.items():
+            # raw_evidence = res_info.get("evidence")
+            # if raw_evidence:
+            #     tmp = json.loads(raw_evidence)
+            #     raw_evidence_json = json.dumps(tmp, indent=4)
+            #     logger.debug(f"get ra evidence:{raw_evidence_json}")
+            if config.ra_key_negotiation:
+                public_key_info = res_info.get("key_info").get("pub_key_info")
+            else:
+                public_key_info = ""
+
+            raw_token = res_info.get("token")
+            if not raw_token:
+                logger.critical("RA response token is NULL, conf={config}")
+                return attestation_status, public_key_info
+            else:
+                evidence_info = res_info.get("evidence")
+                evidence_info_json = json.loads(evidence_info)
+                proof = evidence_info_json.get('proof')
+                logger.debug(f"proof={proof}")
+
+                if config.ra_key_negotiation:
+                    key_info = res_info.get("key_info")
+                    pub_r = key_info.get("rar_pub_key")
+                    timestamp = key_info.get("timestamp")
+                    version = key_info.get("version")
+                    sign = key_info.get("rar_signature")
+                    service_name = config.ra_service_name.strip()
+
+                    if sign:
+                        if not check_pubkey_valid(public_key_info, pub_r, timestamp, version, service_name, sign):
+                            logger.critical("verify sign failed")
+                            return attestation_status, public_key_info
+                        else:
+                            logger.info("ra: get pub_key, verify sign success")
+
+                attestation_status = verify_ra_result(raw_token, public_key_info, pub_r, proof, nonce)
+
+            break
+        return attestation_status, public_key_info
+    except Exception as e:
+        logger.critical(f"Response is not JSON: conf={config}, Exception={e}")
+        raise error.ServiceError("RA", config.ra_url, None, "not JSON") from e
+
+
+def verify_ra_result(raw_token, pub_key_info, pub_r, proof, nonce, challenge_report_data: str = None) -> bool:
+    try:
+        report_data = ""
+        token_parts = raw_token.split(".")
+        if len(token_parts) >= 2:
+            token_data = token_parts[1]
+            try:
+                # 使用base64.urlsafe_b64decode进行URL安全模式的解码
+                decoded_data = base64.urlsafe_b64decode(
+                    token_data + "=" * (4 - len(token_data) % 4)
+                ).decode("utf-8")
+                token_dict = json.loads(decoded_data)
+                report_data = token_dict.get('tdx').get('report_data')
+            except Exception as e:
+                logger.error(f"解码出现错误: {e}")
+                raise Exception("decode token error")
+        logger.info(f"in ra token,report_data={report_data}")
+
+        try:
+            if challenge_report_data:
+                if report_data != challenge_report_data:
+                    return False
+            else:
+                nonce_ret = verify_nonce(proof, report_data, nonce, pub_r) or \
+                            verify_nonce(proof, report_data, nonce, pub_key_info)
+                if not nonce_ret:
+                    # logger.error("verify nonce failed")
+                    # return False
+                    logger.debug("verify nonce failed")
+                    return True
+
+        except Exception as e:
+            logger.error(f"verify nonce failed!!!,msg={str(e)}")
+            return False
+
+        if not verify_jwt_token(raw_token):
+            logger.warning("token verify failed")
+            return False
+        else:
+            token_dict = {}
+            token_parts = raw_token.split(".")
+            if len(token_parts) >= 2:
+                token_data = token_parts[1]
+                try:
+                    # 使用base64.urlsafe_b64decode进行URL安全模式的解码
+                    decoded_data = base64.urlsafe_b64decode(
+                        token_data + "=" * (4 - len(token_data) % 4)
+                    ).decode("utf-8")
+
+                    token_dict = json.loads(decoded_data)
+                except Exception as e:
+                    logger.critical(f"token 解码出现错误: {e}")
+                    return False
+
+            attestation_status = len(token_dict.get("policies_matched")) > 0
+        return attestation_status
+    except Exception as e:
+        logger.critical("verify_ra_result error,msg={}".format(str(e)))
+        raise Exception("verify_ra_result error,msg={}".format(str(e)))
+
+
+def attest_client(client_ra_config: dict, client_bytedance_top_info: str, nonce: str = None) -> bool:
+    """
+    对数据发送方进行远程证明.
+    """
+    ra_service_name: Optional[str] = client_ra_config.get("ra_service_name")
+    ra_policy_id: Optional[str] = client_ra_config.get("ra_policy_id")
+    ra_pods_info = client_ra_config.get("ra_pods_info")
+    ra_key_negotiation: Optional[bool] = client_ra_config.get("ra_key_negotiation")
+
+    top_info = client_bytedance_top_info
+    if top_info and isinstance(top_info, str):
+        top_info = json.loads(top_info)
+
+    ra_url = client_ra_config.get("ra_url", "open.volcengineapi.com")
+    ra_uid = client_ra_config.get("ra_uid")
+    if ra_uid:
+        headers = {"UID": ra_uid}
+    else:
+        headers = {}
+
+    if not nonce:
+        nonce = generate_nonce()
+
+    if ra_service_name:
+        body = {
+            "Nonce": nonce,
+            "PolicyID": ra_policy_id,
+            "ServiceName": ra_service_name,
+            "KeyNegotiation": ra_key_negotiation,
+            "Token": True,
+        }
+    elif ra_pods_info:
+        if isinstance(ra_pods_info, str):
+            ra_pods_info = json.loads(ra_pods_info)
+        body = {
+            "Nonce": nonce,
+            "PolicyID": ra_policy_id,
+            "AttestedPods": [ra_pods_info],
+            "KeyNegotiation": ra_key_negotiation,
+            "Token": True,
+        }
+    else:
+        raise error.ServiceError("RA", ra_service_name, None, "call TCA error")
+
+    response = request_bytedance_gateway(ra_url, body, headers, top_info)
+    if not response:
+        logger.error("Client RA error: response NULL")
+        raise error.ServiceError("RA", ra_service_name, None, "NULL")
+
+    if response.status_code != 200:
+        logger.error(
+            f"client ra error: ra_service_name={ra_service_name} code={response.status_code}"
+        )
+        raise error.ServiceError(
+            "RA", ra_service_name, None, f"code={response.status_code}"
+        )
+
+    try:
+        ra_res = json.loads(response.text).get("Result")
+        if not ra_res:
+            logger.error(
+                f"attest_client failed, verify_result is NULL, ra interface response={response.text}"
+            )
+            return False
+
+        attestation_status = False
+        for cluster_info, res_info in ra_res.items():
+            # raw_evidence = res_info.get("evidence")
+            # tmp = json.loads(raw_evidence)
+            # raw_evidence_json = json.dumps(tmp, indent=4)
+            # logger.debug(f"get ra evidence:{raw_evidence_json}")
+
+            raw_token = res_info.get("token")
+            if not raw_token:
+                logger.critical("token is None")
+                return False
+
+            if not verify_jwt_token(raw_token):
+                logger.critical("token verify failed")
+                return False
+
+            token_dict = {}
+            token_parts = raw_token.split(".")
+            if len(token_parts) >= 2:
+                token_data = token_parts[1]
+                try:
+                    # 使用base64.urlsafe_b64decode进行URL安全模式的解码
+                    decoded_data = base64.urlsafe_b64decode(
+                        token_data + "=" * (4 - len(token_data) % 4)
+                    ).decode("utf-8")
+
+                    token_dict = json.loads(decoded_data)
+                except Exception as e:
+                    logger.critical(f"token 解码出现错误: {e}")
+                    return False
+
+            attestation_status = len(token_dict.get("policies_matched")) > 0
+            break
+
+        return attestation_status
+    except Exception as e:
+        logger.critical(f"Response is not JSON: client_service_name={ra_service_name}")
+        raise error.ServiceError("RA", ra_service_name, None, "not JSON") from e

src/bytedance/jeddak_secure_channel/server.py

@@ -0,0 +1,734 @@
+"""
+数据接收方接口.
+"""
+
+__all__ = ["Server"]
+
+import base64
+import hashlib
+import json
+import os
+import random
+import time
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from threading import Lock
+
+import requests
+from typing_extensions import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+if TYPE_CHECKING:
+    from _typeshed import GenericPath
+else:
+    GenericPath = str
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
+
+from bytedance.tks.client import TKSClient, TKSConfig, utc8_timestamp
+from . import error
+from .config import ServerConfig
+from .crypto import FileMode, PublicKey, ResponseKey, ServerSessionKey
+from .log import LogConfig, logger
+from .ra import (
+    RaRequest,
+    RaResponse,
+    RaResponseMetadata,
+    RaResponsePod,
+    RaResponsePods,
+    attest_client,
+    validata_ra_request,
+)
+from .service import EpsClient, RasClient
+from .utils import RepeatTimer, request_bytedance_gateway, weak_repeat_timer
+
+
+@dataclass(eq=False)
+class Server:
+    """
+    数据接收方接口.
+
+    本类的所有方法是线程安全的.
+    """
+
+    config: ServerConfig
+
+    ras_client: RasClient = field(init=False, repr=False)
+
+    eps_client: EpsClient = field(init=False, repr=False)
+
+    refresh_timer: Optional[RepeatTimer] = field(init=False, default=None)
+
+    session_key: Optional[ServerSessionKey] = field(init=False, default=None)
+
+    old_session_key: Optional[ServerSessionKey] = field(init=False, default=None)
+
+    session_key_lock: Lock = field(init=False, default_factory=Lock)
+    """
+    线程安全: 修改 session_key 应该加锁.
+    """
+
+    root_key_info: Dict[str, PublicKey] = field(init=False, default_factory=dict)
+
+    def __post_init__(self):
+        try:
+            # 处理环境变量
+            # need_evidence = os.environ.get("NEED_EVIDENCE")
+            # if need_evidence != None:
+            #     if str(need_evidence).lower() == "false":
+            #         self.config.need_evidence = False
+            #
+            # attest_gpu = os.environ.get("ATTEST_GPU")
+            # if attest_gpu != None:
+            #     if str(attest_gpu).lower() == "false":
+            #         self.config.attest_gpu = False
+            self._init_log()
+
+            self.ras_client = RasClient(self.config)
+            self.eps_client = EpsClient(self.config)
+
+            if self.config.root_key_config:
+                self._root_key_parse()
+
+            if self.config.refresh_interval is not None:
+                refresh_interval = self.config.refresh_interval
+
+                # 避免多实例集中访问TKS，所以这里加了一个随机数
+                refresh_interval = refresh_interval + random.randint(1, 60)
+                logger.info(
+                    f"start auto update key, refresh_interval={refresh_interval}"
+                )
+
+                self.refresh_timer = timer = weak_repeat_timer(
+                    refresh_interval, self._auto_update
+                )
+                timer.start()
+
+            self._auto_update()
+        except Exception as e:
+            logger.critical(f"__post_init__ failed, e={e}")
+
+    def _auto_update(self):
+        """
+        :return:
+        """
+        try:
+            self._auto_load_key()
+        except Exception as e:
+            logger.critical(f"load tks key failed, e={e}")
+
+        try:
+            if self.config.client_ra_config is not None:
+                if not self.attest_client():
+                    logger.error("attest client failed, Please check your config !")
+
+        except Exception as e:
+            logger.critical(f"attest client failed, e={e}")
+
+    def _auto_load_key(self):
+        try:
+            logger.info("start load_key")
+            if self.config.key_file:
+                if not self.load_file_key():
+                    logger.error("load key from local key file failed, Please check your config !")
+            elif self.config.key_type == "tks":
+                if not self.load_tks_key():
+                    logger.error("load key from tks failed, Please check your config !")
+            elif self.config.key_type == "random":
+                if not self.generate_key():
+                    logger.error("generate key failed !")
+            else:
+                logger.critical("key config error, Both TKS keyid or key file is empty")
+
+        except Exception as e:
+            logger.critical(f"load key failed, e={e}")
+
+    def _add_key(self, private_pem: Union[str, bytes], tks_version: Optional[int]):
+        """
+        添加密钥对.
+        """
+        # 判断是否需要更新密钥
+        if isinstance(private_pem, str):
+            new_md5 = hashlib.md5(private_pem.encode()).hexdigest()
+        else:
+            new_md5 = hashlib.md5(private_pem).hexdigest()
+
+        if self.session_key:
+            old_md5 = self.session_key.pem_md5
+        else:
+            old_md5 = ""
+
+        if old_md5 != new_md5:
+            expire_at = datetime.now(timezone.utc) + timedelta(
+                seconds=self.config.valid_time
+            )
+            session_key = ServerSessionKey.load(
+                private_pem, expire_at, tks_version, new_md5
+            )
+            with self.session_key_lock:
+                self.old_session_key = self.session_key
+                self.session_key = session_key
+
+    def _pcc_login(self) -> Dict[str, str]:
+        """
+        访问pcc_controller模块，获取cookie
+        :return:
+        """
+
+        if self.config.pcc_login_url is None:
+            raise error.ConfigError("pcc_login_url", "PCC login url is missing")
+        if self.config.pcc_password is None:
+            raise error.ConfigError("pcc_password", "PCC password is missing")
+
+        ts = str(utc8_timestamp())
+        req_body = {
+            "Username": self.config.pcc_app_id,
+            "Password": hashlib.md5(
+                self.config.pcc_password.encode() + ts.encode()
+            ).hexdigest(),
+        }
+
+        response = requests.post(
+            self.config.pcc_login_url, headers={"TimeStamp": ts}, json=req_body
+        )
+        cookies = response.cookies
+        jeddak_pcc_session = cookies.get("jeddak_pcc_session")
+        assert jeddak_pcc_session is not None
+        cookies = {"jeddak_pcc_session": jeddak_pcc_session}
+        return cookies
+
+    def _get_pairkey_from_pem(self, private_key: bytes) -> Optional[Dict[str, str]]:
+        """
+        从私钥pem中产生公私钥对 for peixuan
+        :param private_key:
+        :return:
+        """
+        private_key_pem = serialization.load_pem_private_key(private_key, password=None)
+        if not isinstance(private_key_pem, RSAPrivateKey):
+            logger.error(f"not RSAPrivateKey, private_key={private_key}")
+            return None
+
+        pub_key = private_key_pem.public_key()
+        pub_key_str = pub_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        ).decode("utf-8")
+
+        ret = dict()
+        ret["pub_key"] = pub_key_str
+        ret["pri_key"] = private_key.decode()
+
+        return ret
+
+    def _get_tks_info_from_pcc(self, service_id: str, app_id: str):
+        try:
+            if not service_id:
+                return None
+
+            top_conf = json.loads(self.config.bytedance_top_info)
+            top_conf["action"] = "ListJscKey"
+
+            url = top_conf.get("url", "open.volcengineapi.com")
+            body = {"ServiceId": service_id, "AccountId": app_id}
+            response = request_bytedance_gateway(url, body, {}, top_conf)
+            if not response or response.status_code != 200:
+                logger.error("get tks info from pcc_controller failed")
+                return None
+
+            ra_res = json.loads(response.text).get("Result")
+            if not ra_res:
+                logger.error(f"get key info from pcc_controller error: ServiceId={service_id}")
+                return None
+
+            key_list = json.loads(response.text).get("Result").get("KeyList")
+            if len(key_list) == 1:
+                ring_id = key_list[0].get("RingID")
+                key_id = key_list[0].get("KeyID")
+                return ring_id, key_id
+            else:
+                logger.error(f"get key info from pcc_controller error: key_list={key_list}")
+                return None
+
+        except Exception as e:
+            logger.critical("_get_tks_info_from_pcc exception={}".format(e))
+            return None
+
+    def load_tks_key(self) -> bool:
+        """
+        从 TKS 获取或更新密钥对.
+        """
+        logger.info("start load TKS key")
+
+        app_id, ring_id, key_id = (
+            self.config.tks_app_id,
+            self.config.tks_ring_id,
+            self.config.tks_key_id,
+        )
+
+        pcc_conf = None
+        top_conf = None
+        if self.config.bytedance_pcc_info:
+            pcc_conf = json.loads(self.config.bytedance_pcc_info)
+        elif self.config.bytedance_top_info:
+            top_conf = json.loads(self.config.bytedance_top_info)
+
+        enable_tls = True
+        if enable_tls_env := os.environ.get("enable_tls"):
+            enable_tls = enable_tls_env.lower() == "true"
+        elif enable_tls_env := os.environ.get("ENABLE_TLS"):
+            enable_tls = enable_tls_env.lower() == "true"
+
+        cfg = TKSConfig(
+            addr=self.config.tks_url,
+            enable_tls=enable_tls,
+            pcc_config=pcc_conf,
+            top_config=top_conf,
+        )
+
+        tks_client = TKSClient(app_id, cfg)
+
+        need_evidence_conf = self.config.need_evidence
+        need_evidence = self.config.need_evidence
+        need_evidence_env = os.environ.get("NEED_EVIDENCE")
+        if need_evidence_env:
+            need_evidence = not (need_evidence_env.lower() == "false")
+        logger.info(
+            f"need_evidence: conf={need_evidence_conf}, env={need_evidence_env}, use={need_evidence}"
+        )
+
+        attest_gpu_conf = self.config.attest_gpu
+        attest_gpu = attest_gpu_conf
+        attest_gpu_env = os.environ.get("ATTEST_GPU")
+        if attest_gpu_env:
+            attest_gpu = not (attest_gpu_env.lower() == "false")
+        logger.info(
+            f"attest_gpu: conf={attest_gpu_conf}, env={attest_gpu_env}, use={attest_gpu}"
+        )
+
+        max_try_count = 3
+        for _ in range(max_try_count):
+            try:
+                service_id = self.config.service_id
+                if not service_id:
+                    service_id = os.getenv("JPCC_JCK_SERVICE_ID")
+
+                key_info = self._get_tks_info_from_pcc(service_id, self.config.tks_app_id)
+                if key_info:
+                    ring_id = key_info[0]
+                    key_id = key_info[1]
+
+                if self.config.pcc_login_url:
+                    cookies = self._pcc_login()
+                    private_key = tks_client.get_key(
+                        ring_id=ring_id,
+                        key_id=key_id,
+                        need_evidence=need_evidence,
+                        attest_gpu=attest_gpu,
+                        cookies=cookies,
+                    )
+                else:
+                    private_key = tks_client.get_key(
+                        ring_id=ring_id,
+                        key_id=key_id,
+                        need_evidence=need_evidence,
+                        attest_gpu=attest_gpu,
+                        policy_id=self.config.tks_policy_id,
+                        tks_url=self.config.tks_url,
+                        tks_uid=self.config.tks_app_id,
+                        tks_service_name=self.config.tks_service_name,
+                    )
+
+                pair_key_info = self._get_pairkey_from_pem(private_key)
+                logger.info("get key from TKS success")
+                logger.debug(f"TKS key info={pair_key_info}")
+
+                self._add_key(private_key, None)
+
+                return True
+            except Exception as e:
+                logger.critical(f"load_tks_key exception={e}")
+                logger.critical("load_tks_key failed, try again ....")
+                time.sleep(5)
+
+        return False
+
+    def load_file_key(self, path: Optional[GenericPath] = None):
+        """
+        从文件读取 RSA 密钥.
+
+        Args:
+            path: 密钥文件路径. 默认为配置中的 `key_file`.
+        """
+
+        path = path or self.config.key_file
+        if path is None:
+            raise error.ParamError("path", "Key file is not specified")
+
+        with open(path) as f:
+            self.load_key(f.read())
+
+        return True
+
+    def load_key(self, private_key: Union[str, bytes]):
+        """
+        导入 RSA 密钥.
+
+        Args:
+            private_key: PEM 格式的 RSA 私钥.
+        """
+        self._add_key(private_key, tks_version=None)
+
+    def generate_key(self):
+        """
+        随机生成密钥对.
+        """
+        expire_at = datetime.now(timezone.utc) + timedelta(
+            seconds=self.config.valid_time
+        )
+
+        session_key = ServerSessionKey.generate(expire_at)
+        # 修改时加锁
+        with self.session_key_lock:
+            self.old_session_key = self.session_key
+            self.session_key = session_key
+
+        return True
+
+    def remove_expired_keys(self):
+        """
+        删除已过期的密钥对.
+        """
+        # 修改时加锁
+        with self.session_key_lock:
+            if self.session_key is not None and self.session_key.is_expired():
+                self.session_key = None
+
+            if self.old_session_key is not None and self.old_session_key.is_expired():
+                self.old_session_key = None
+
+    def _get_public_key(self) -> str:
+        """
+        获取已保存的最新的公钥.
+
+        Raises:
+            KeyMissingError: 没有可用的密钥. 应通过 `load_tks_key`, `load_key`, 或
+                `generate_key` 添加密钥.
+        """
+        with self.session_key_lock:
+            current_key = self.session_key
+
+        if not current_key:
+            raise error.KeyMissingError() from None
+
+        return current_key.get_public_key()
+
+    def handle_ra_request(self, request: RaRequest) -> RaResponse:
+        """
+        处理来自数据发送方的远程证明请求.
+        应当在调用本函数前, 验证请求中 `token` 身份令牌的有效性.
+
+        Returns:
+            远程证明响应. 可作为 JSON 响应给数据发送方.
+        Raises:
+            NetworkError: 访问证明生成服务 (EPS) 或远程证明服务 (RAS) 时出现问题.
+            KeyMissingError: 没有可用的密钥. 应通过 `load_tks_key`, `load_key`, 或
+                `generate_key` 添加密钥.
+        """
+        validata_ra_request(request)
+
+        # quote = self.eps_client.get_quote(request["nonce"])
+        # token = self.ras_client.get_attestation_evaluation("TDX", quote)
+
+        res_meta_data = RaResponseMetadata(Code=0, Message="success")
+        res_result: RaResponsePods = {
+            "pod_1": RaResponsePod(
+                evidence="evidence",
+                token="token",
+                key_info={"pub_key_info": self._get_public_key()},
+            ),
+        }
+
+        return RaResponse(ResponseMetadata=res_meta_data, ResponseResult=res_result)
+
+    def decrypt(
+            self,
+            message: Union[str, bytes],
+            app_info: Optional[str] = None,
+            sign: Optional[str] = None,
+    ) -> bytes:
+        """
+        解密二进制数据.
+
+        Raises:
+            DecryptionError: 加密消息不合法, 或者加密使用的密钥不在当前可用的密钥中.
+            KeyMissingError: 没有可用的密钥. 应通过 `load_tks_key`, `load_key`, 或
+                `generate_key` 添加密钥.
+        """
+        if app_info and sign:
+            if not self.verify_sign(app_info, message, sign):
+                raise error.SignatureError("verify sign failed")
+
+        plaintext, _ = self.decrypt_with_response(message)
+        return plaintext
+
+    def decrypt_with_response(
+            self,
+            message: Union[str, bytes],
+            app_info: Optional[str] = None,
+            sign: Optional[str] = None,
+    ) -> Tuple[bytes, ResponseKey]:
+        """
+        解密二进制数据, 并获取用于加密响应的密钥.
+
+        Raises:
+            DecryptionError: 加密消息不合法, 或者加密使用的密钥不在当前可用的密钥中.
+            KeyMissingError: 没有可用的密钥. 应通过 `load_tks_key`, `load_key`, 或
+                `generate_key` 添加密钥.
+        """
+        last_error: Optional[Exception] = None
+
+        if app_info and sign:
+            if not self.verify_sign(app_info, message, sign):
+                raise error.SignatureError("verify sign failed")
+
+        with self.session_key_lock:
+            current_key = self.session_key
+            old_key = self.old_session_key
+
+        for key in [current_key, old_key]:
+            if key is None:
+                continue
+            try:
+                response_key = key.decrypt(message)
+                return response_key
+            except Exception as e:
+                last_error = e
+                continue
+
+        if last_error:
+            # 尝试更新一下密钥，因为上面的失败可能是本地密钥版本晚于客户端的密钥版本，等TKS支持密钥版本后可用密钥版本号来比较
+            self._auto_load_key()
+            if self.session_key is not None:
+                try:
+                    response_key = self.session_key.decrypt(message)
+                    return response_key
+                except Exception as e:
+                    last_error = e
+
+        raise last_error or error.KeyMissingError()
+
+    def decrypt_file(
+            self,
+            metadata: Union[str, bytes],
+            source_path: GenericPath,
+            dest_path: GenericPath,
+            mode: FileMode,
+            app_info: Optional[str] = None,
+            sign: Optional[str] = None,
+    ) -> ResponseKey:
+        """
+        解密文件数据.
+
+        Returns:
+            用于加密响应的密钥.
+        """
+        if mode not in ["b", "t"]:
+            raise error.ParamError("mode", "Mode should be 'b' or 't'")
+
+        if app_info and sign:
+            # TODO: fill in message
+            if not self.verify_sign(app_info, metadata, sign):
+                raise error.SignatureError("verify sign failed")
+
+        with self.session_key_lock:
+            current_key = self.session_key
+            old_key = self.old_session_key
+
+        last_error: Optional[Exception] = None
+        for key in [current_key, old_key]:
+            if key is None:
+                continue
+            try:
+                with open(source_path, f"r{mode}") as source, open(
+                        dest_path, f"w{mode}"
+                ) as dest:
+                    response_key = key.decrypt_stream(metadata, source, dest, mode)
+                    return response_key
+            except Exception as e:
+                last_error = e
+                continue
+
+        if last_error:
+            # 尝试更新一下密钥，因为上面的失败可能是本地密钥版本晚于客户端的密钥版本，等TKS支持密钥版本后可用密钥版本号来比较
+            self._auto_load_key()
+            if self.session_key is not None:
+                try:
+                    with open(source_path, f"r{mode}") as source, open(
+                            dest_path, f"w{mode}"
+                    ) as dest:
+                        response_key = self.session_key.decrypt_stream(
+                            metadata, source, dest, mode
+                        )
+                        return response_key
+                except Exception as e:
+                    last_error = e
+
+        raise last_error or error.KeyMissingError()
+
+    def attest_client(self) -> bool:
+        """
+        对客户端进行远程证明, 判断客户端是否可信.
+
+        Returns:
+            True 为可信, False 为不可信.
+        """
+        client_ra_config = self.config.client_ra_config
+        if not client_ra_config:
+            return True
+
+        client_ra_config = json.loads(client_ra_config)
+        if not client_ra_config.get("open"):
+            return True
+
+        logger.info("start attest client")
+        nonce = ''.join(str(uuid.uuid4()).split('-'))[-16:]
+        max_try_count = 3
+        for _ in range(max_try_count):
+            ret = attest_client(client_ra_config, self.config.bytedance_top_info, nonce)
+            if ret:
+                logger.info("attest client success")
+                return True
+            else:
+                logger.info("attest client failed, try again...")
+
+        return False
+
+    def _root_key_parse(self):
+        """
+        [{
+            "app_info": "lenovo_xiao_tian",
+            "pub_key": "******"
+        }]
+        :return:
+        """
+        try:
+            root_key_config = self.config.root_key_config
+            if not root_key_config:
+                return
+
+            if isinstance(root_key_config, str):
+                root_key_config = json.loads(root_key_config)
+
+            for tmp_info in root_key_config:
+                app_info = tmp_info.get("app_info")
+                pub_key = tmp_info.get("pub_key")
+                if pub_key:
+                    with open(pub_key) as f:
+                        public_pem = f.read()
+                    self.root_key_info[app_info] = PublicKey.from_public_pem(public_pem)
+                else:
+                    # 从pcc_controller中读取
+                    public_pem = self._get_app_key_from_pcc(app_info)
+                    if public_pem:
+                        self.root_key_info[app_info] = PublicKey.from_public_pem(
+                            public_pem
+                        )
+                    else:
+                        logger.error("get app pubkey from pcc_controller failed")
+
+            logger.info("root_key_parse success")
+        except Exception as e:
+            logger.critical("root_key_parse exception={}".format(e))
+            raise error.ConfigError("RootKey conf", "RootPubKey conf error")
+
+    def verify_sign(
+            self, app_info: str, msg: Union[str, bytes], sign: Union[str, bytes]
+    ) -> bool:
+        """
+        签名验证
+        :param app_info: app_id+app_name
+        :param msg:
+        :param sign:
+        :return:
+        """
+        try:
+            if not app_info or not msg or not sign:
+                logger.info(f"verify_sign error, app_info={app_info}, sign={sign}")
+                raise error.SignatureError("RSA signing failed") from None
+
+            if app_info not in self.root_key_info:
+                logger.info(
+                    f"verify_sign error, can not find app_info in root key conf, app_info={app_info}"
+                )
+                raise error.SignatureError(
+                    f"can not find app_info={app_info}"
+                ) from None
+
+            if isinstance(sign, str):
+                sign = base64.b64decode(sign.encode())
+
+            return self.root_key_info[app_info].verify(sign, msg)
+
+        except Exception as e:
+            logger.critical("verify_sign exception={}".format(e))
+            raise error.SignatureError("can not find app_info") from e
+
+    def _get_app_key_from_pcc(self, app_info):
+        try:
+            if not app_info:
+                return None
+
+            top_conf = json.loads(self.config.bytedance_top_info)
+            top_conf["action"] = "ListJscAppInfo"
+
+            url = top_conf.get("top_url", "open.volcengineapi.com")
+            body = {"PageSize": -1}
+            headers = {}
+            response = request_bytedance_gateway(url, body, headers, top_conf)
+            if not response or response.status_code != 200:
+                logger.error("get app key from pcc_controller failed")
+                return None
+
+            ra_res = json.loads(response.text).get("Result")
+            if not ra_res:
+                logger.error(
+                    f"get app key from pcc_controller error: app_info={app_info}"
+                )
+                return None
+
+            app_list = json.loads(response.text).get("Result").get("AppList")
+            for app in app_list:
+                if app_info == app.get("AppInfo"):
+                    return app.get("AppKey")
+
+            return None
+        except Exception as e:
+            logger.critical("get_app_key_from_pcc exception={}".format(e))
+            return None
+
+    def _init_log(self):
+        config = LogConfig()
+
+        log_config = dict()
+        if self.config.log_config:
+            log_config = json.loads(self.config.log_config)
+
+        if config_dir := log_config.get("dir"):
+            config.dir = config_dir
+
+        if config_filename := log_config.get("filename"):
+            config.filename = config_filename
+        # else:
+        #     config.filename = "server.log"
+
+        if config_rotation := log_config.get("rotation"):
+            config.rotation = config_rotation
+
+        if config_retention := log_config.get("retention"):
+            config.retention = config_retention
+
+        if config_level := log_config.get("level"):
+            config.level = config_level
+
+        logger.init_log_config(config)

src/bytedance/jeddak_secure_channel/service/__init__.py

@@ -0,0 +1,16 @@
+__all__ = [
+    "EpsClient",
+    "EpsClient",
+    "EpsConfig",
+    "EpsConfig",
+    "RasClient",
+    "RasConfig",
+    "TksAppOptions",
+    "TksClient",
+    "TksConfig",
+]
+
+from .eps import EpsClient, EpsConfig
+from .ras import RasClient, RasConfig
+from .tks import TksClient, TksConfig
+from .tks_models import AppOptions as TksAppOptions

src/bytedance/jeddak_secure_channel/service/eps.py

@@ -0,0 +1,133 @@
+"""
+证明生成服务 (Evidence Provision Service, EPS) 的客户端.
+"""
+
+__all__ = ["EpsClient", "EpsConfig"]
+
+import logging
+from dataclasses import dataclass
+
+from typing_extensions import Any, Mapping
+
+from .. import error
+from .eps_models import Evidence, ExtendRuntimeMeasurementParams, NonceData
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EpsConfig:
+    """访问证明生成服务所需的配置."""
+
+    eps_url: str = "http://localhost:8006"
+    """EPS 服务 Attestation Agent 的 HTTP 地址."""
+
+
+def aa_request(endpoint: str, params: Mapping[str, Any], config: EpsConfig) -> str:
+    """
+    向 Attestation Agent 的指定端点发送 HTTP GET 请求.
+
+    Returns:
+        响应体的字符串.
+    """
+    import requests
+
+    try:
+        response = requests.get(config.eps_url + endpoint, params)
+    except Exception as e:
+        logger.exception(f"Network error: service={config.eps_url} {endpoint=}")
+        raise error.NetworkError("EPS", config.eps_url, endpoint) from e
+
+    if response.status_code != 200:
+        logger.error(
+            f"Service error: service={config.eps_url} {endpoint=} code={response.status_code}"
+        )
+        raise error.ServiceError("EPS", config.eps_url, endpoint, f"code={response.status_code}")
+
+    content_type = response.headers.get("content-type", "")
+    if not (content_type.startswith("application/octet-stream") or content_type.startswith("application/json")
+            or content_type.startswith("text/plain")):
+        logger.error(f"Service error: service={config.eps_url} {endpoint=} {content_type=}")
+        raise error.ServiceError("EPS", config.eps_url, endpoint, f"{content_type=}")
+
+    try:
+        return str(response.content, encoding="utf-8")
+    except UnicodeDecodeError as e:
+        logger.exception(f"Service error: service={config.eps_url} {endpoint=}")
+        raise error.ServiceError("EPS", config.eps_url, endpoint, "cannot decode text") from e
+
+
+@dataclass(eq=False)
+class EpsClient:
+    """
+    证明生成服务 (Evidence Provision Service, EPS) 的客户端.
+
+    本类的所有方法是线程安全的.
+    """
+
+    config: EpsConfig
+    """访问证明生成服务所需的配置."""
+
+    def get_token(self, token_type: str) -> str:
+        """
+        访问 Attestation Agent 的 `/aa/token` 端点, 获取令牌.
+
+        Returns:
+            令牌.
+        """
+        return aa_request("/aa/token", {"token_type": token_type}, self.config)
+
+    def get_evidence_raw(self, nonce: str) -> str:
+        """
+        访问 Attestation Agent 的 `/aa/evidence` 端点, 获取证明.
+
+        Returns:
+            JSON 编码的证明. 包含 `cc_eventlog`, `quote`, `aa_eventlog` 字段.
+        """
+        nonce_data = NonceData(runtime_data="", nonce=nonce, user_data="")
+
+        return aa_request("/aa/evidence", nonce_data, self.config)
+
+    def get_evidence(self, nonce: str) -> Evidence:
+        """
+        访问 Attestation Agent 的 `/aa/evidence` 端点, 获取证明.
+
+        Returns:
+            证明内容, 包含 `cc_eventlog`, `quote`, `aa_eventlog` 字段.
+        """
+        import json
+
+        raw = self.get_evidence_raw(nonce)
+
+        return json.loads(raw)
+
+    def get_eventlog(self, nonce: str) -> str:
+        """
+        访问 Attestation Agent 的 `/aa/eventlog` 端点, 获取事件日志.
+
+        Returns:
+            Base64 编码的事件日志. 事件日志长度较大 (~1.5MB).
+        """
+        nonce_data = NonceData(runtime_data="", nonce=nonce, user_data="")
+
+        return aa_request("/aa/eventlog", nonce_data, self.config)
+
+    def get_quote(self, nonce: str) -> str:
+        """
+        访问 Attestation Agent 的 `/aa/quote` 端点, 获取度量值.
+
+        Returns:
+            Base64 编码的度量值.
+        """
+        nonce_data = NonceData(runtime_data="", nonce=nonce, user_data="")
+
+        return aa_request("/aa/quote", nonce_data, self.config)
+
+    def get_extended_runtime_measurement(self, params: ExtendRuntimeMeasurementParams) -> str:
+        """
+        访问 Attestation Agent 的 `/aa/extend_runtime_measurement` 端点, 获取运行时度量事件日志.
+
+        Returns:
+            明文的运行时度量事件日志.
+        """
+        return aa_request("/aa/extend_runtime_measurement", params, self.config)

src/bytedance/jeddak_secure_channel/service/eps_models.py

@@ -0,0 +1,47 @@
+"""
+证明生成服务 (EPS) 定义的数据类型.
+"""
+
+from typing_extensions import NotRequired, TypedDict
+
+
+class NonceData(TypedDict):
+    """
+    Attestation Agent 的 `/aa/evidence`, `/aa/quote`, `/aa/eventlog` 端点的请求体.
+    目前 Attestation Agent 会将三个字段值直接连接.
+    """
+
+    runtime_data: str
+
+    nonce: str
+
+    user_data: str
+
+
+class ExtendRuntimeMeasurementParams(TypedDict):
+    """
+    Attestation Agent 的 `/aa/extend_runtime_measurement` 端点的请求体.
+    """
+
+    domain: str
+
+    operation: str
+
+    context: str
+
+    register_index: NotRequired[int]
+
+
+class Evidence(TypedDict):
+    """
+    Attestation Agent 的 `/aa/evidence` 端点的响应体.
+    """
+
+    cc_eventlog: str
+    """Base64 编码的事件日志. 事件日志长度较大 (~1.5MB)."""
+
+    quote: str
+    """Base64 编码的度量值."""
+
+    aa_eventlog: str
+    """明文的运行时度量事件日志."""

src/bytedance/jeddak_secure_channel/service/ras.py

@@ -0,0 +1,88 @@
+"""
+远程证明服务 (Remote Attestation Service, RAS) 的客户端.
+"""
+
+__all__ = ["RasClient", "RasConfig"]
+
+import base64
+import json
+import logging
+from dataclasses import dataclass
+
+from typing_extensions import Any
+
+from .. import error
+
+# try:
+#     import grpc
+#
+#     from .ras_models import attestation_pb2, attestation_pb2_grpc
+# except ModuleNotFoundError:
+#     grpc = None
+#     attestation_pb2_grpc = None
+
+
+logger = logging.getLogger(__name__)
+
+KNOWN_TEE = ["SEV", "SGX", "SNP", "TDX", "Sample", "AzSnpVtpm", "CSV", "CCA", "AzTdxVtpm"]
+
+
+@dataclass
+class RasConfig:
+    """访问远程证明服务所需的配置."""
+
+    ras_address: str = "localhost:50004"
+    """远程证明服务的 GRPC 地址."""
+
+
+def encode_base64_no_pad(data: Any) -> str:
+    return base64.urlsafe_b64encode(json.dumps(data).encode("utf-8")).decode("utf-8").rstrip("=")
+
+
+@dataclass(eq=False)
+class RasClient:
+    """
+    远程证明服务 (Remote Attestation Service, RAS) 的客户端.
+
+    本类的所有方法是线程安全的.
+    """
+
+    config: RasConfig
+    """访问远程证明服务所需的配置."""
+
+    def get_attestation_evaluation(self, tee: str, quote: str) -> str:
+        """
+        访问远程证明服务的 AttestationRequest 端点来验证度量值.
+
+        Args:
+            tee: 所属 TEE 环境类型.
+            quote: Base64 编码的度量值. 参考 `EpsClient#get_quote()`.
+        Returns:
+            Base64 编码的验证令牌.
+        """
+        pass
+
+        # if grpc is None or attestation_pb2_grpc is None:
+        #     raise error.ConfigError(self, "RAS depends on grpcio which is not installed")
+        #
+        # # An early sanity check.
+        # if tee not in KNOWN_TEE:
+        #     raise error.ParamError("tee", f"Unknown TEE type '{tee}'")
+        #
+        # encoded_evidence = encode_base64_no_pad({"quote": quote})
+        # request = attestation_pb2.AttestationRequest(  # pyright: ignore
+        #     tee=tee, evidence=encoded_evidence, policy_ids=[]
+        # )
+        #
+        # try:
+        #     channel = grpc.insecure_channel(self.config.ras_address)
+        #     client = attestation_pb2_grpc.AttestationServiceStub(channel)
+        #
+        #     response = client.AttestationEvaluate(request)
+        # except Exception as e:
+        #     logger.exception(
+        #         f"Network error service={self.config.ras_address} endpoint=AttestationEvaluate"
+        #     )
+        #     raise error.NetworkError("RAS", self.config.ras_address, "AttestationEvaluate") from e
+        #
+        # return response.attestation_token

src/bytedance/jeddak_secure_channel/service/tks.py

@@ -0,0 +1,51 @@
+"""
+密钥管理服务 (Trusted Key Service, TKS) 的客户端.
+TCA采用背调模式
+"""
+
+__all__ = ["TksClient", "TksConfig"]
+
+from dataclasses import dataclass, field
+from typing_extensions import Any, Optional, Tuple
+
+from .eps import EpsClient
+from .ras import RasClient
+
+
+@dataclass
+class TksConfig:
+    """访问密钥管理服务所需的配置."""
+
+    tks_url: str = "http://localhost:6789"
+    """TKS 服务的 HTTP 地址."""
+
+    tks_timezone: float = 8.0
+    """TKS 服务使用的时区. 一般不需要修改."""
+
+    tks_enable_server_auth: bool = True
+    """
+    是否对 TKS 服务提供的证明报告进行验证.
+    默认为是.
+    """
+
+    tks_enable_bi_auth: bool = True
+    """
+    是否将客户端的证明报告发送给 TKS 服务.
+    默认为是.
+    """
+
+
+@dataclass(eq=False)
+class TksClient:
+    """
+    密钥管理服务 (Trusted Key Service, TKS) 的客户端.
+
+    本类的所有方法是线程安全的.
+    """
+
+    config: TksConfig
+    """访问密钥管理服务所需的配置."""
+
+    ras_client: Optional[RasClient] = field(repr=False, default=None)
+
+    eps_client: Optional[EpsClient] = field(repr=False, default=None)

src/bytedance/jeddak_secure_channel/service/tks_models.py

@@ -0,0 +1,334 @@
+"""
+密钥管理服务 (TKS) 定义的数据类型.
+"""
+
+from dataclasses import dataclass
+
+from typing_extensions import Generic, Literal, NotRequired, Optional, TypedDict, TypeVar
+
+
+class TksAuthHeaders(TypedDict):
+    AppID: str
+    """请求TKS的客户端的身份ID"""
+
+    Timestamp: str
+    """
+    Unix时间戳（UTC+8）
+    Example: 1722253648
+    """
+
+    Signature: NotRequired[str]
+    """
+    对请求Body的签名，仅当账号注册时的认证方式带有cert时（cert | cert+zti）
+    """
+
+    Token: NotRequired[str]
+    """
+    认证Token，仅当账号注册时的认证方式带有zti时（zti | cert+zti）
+    """
+
+
+class ResponseError(TypedDict):
+    Code: int
+    """
+    错误码
+    Example: 69905
+    """
+
+    Message: str
+    """
+    请求处理错误原因，正常请求为空
+    Example: error occured while processing your request
+    """
+
+
+class ResponseMeta(TypedDict):
+    Method: str
+    """
+    请求方法
+    Example: KeyGet
+    """
+
+    Error: NotRequired[ResponseError]
+    """错误信息，在请求正常响应情况下为空"""
+
+
+T = TypeVar("T")
+
+
+class TksResponse(Generic[T], TypedDict):
+    ResponseMetadata: ResponseMeta
+    """描述响应状态的元信息"""
+
+    Result: NotRequired[T]
+    """返回结果，当请求出错时为空"""
+
+
+class ClientChallenge(TypedDict):
+    NonceDown: str
+    """
+    经过hex编码后，长度为256 bits的随机挑战值
+    Example: 5779d18b2b8718cef025282df2a313f154ad4f7b862b2948ecf6cf401de17634
+    """
+
+    RxpirationTime: int
+    """
+    时区为UTC+8的Unix时间戳，表示随机挑战的过期时间
+    Example: 1721972547
+    """
+
+    ChallSig: int
+    """
+    经过Based64编码后的，TKS针对challenge自身的签名。使用TKS RA阶段提供的cert对应的私钥签名（ECDSA-SECP384R1）
+    Example: yk63SiKLaBB7E0OVQs3Flw8NsTRNEYJIUos8s9F14sg=
+    """
+
+
+class SecurityAttestation(TypedDict):
+    NonceUp: str
+    """
+    随机字符串，针对TSK进行远程证明的挑战值（32随机bytes，经过Hex编码）
+    Example: a314b43502077cd89d21378b29954207aaf5b3848d4d7fa9f16682ba1f6f3317
+    """
+
+    BiAuth: NotRequired[bool]
+    """
+    是否要求TKS挑战请求方
+    Example: true
+    """
+
+
+class AttestationResult(TypedDict):
+    RAType: str
+    """
+    TKS提供的远程证明报告的类型
+    Example: sgx | tdx
+    """
+
+    Report: str
+    """
+    经过base64编码后的远程证明报告，需要根据raType来进行验证
+    Example: MTIzNA==
+    """
+
+    Cert: str
+    """
+    经过base64编码后的，PEM形式的证书（ECDH-SECP384R1-SHA256）
+    Example: LS0tLS1CRUdJTiBQVU==
+    """
+
+    DHParam: str
+    """
+    经过base64编码后的，密钥协商公共参数（ECDH-SECP384R1-SHA256）
+    Example: LS0tLS1CRUdJTiBQVUJMSUMg
+    """
+
+    Challenge: NotRequired[ClientChallenge]
+    """当biAuth参数为true时，返回的针对请求方的挑战内容"""
+
+
+class KMSInfo(TypedDict):
+    Type: Literal["tks"]
+    """托管密钥服务的类型，当前仅支持tks"""
+
+
+class RAEvidence(TypedDict):
+    TEEType: str
+    """指定请求方的TEE类型"""
+
+    Report: str
+    """
+    经过Base64编码的，请求方TEE生成的远程证明报告，需要包含clientChall中提供的随机挑战nonceDown
+    Example: OUeRWgdMoQahtmzTA0/36qA5Lh8wp8nIK/ySOe0wQeA=
+    """
+
+
+class ZTIEvidence(TypedDict):
+    Token: str
+    """
+    请求方的ZTI JWT Token
+    """
+
+
+class CertEvidence(TypedDict):
+    Signature: str
+    """
+    经过Base64编码的，请求方的私钥签名的nonceDown
+    Example: RZeo3OjJM/Gam+aiun+hNzD1mMnnJmpiWckLnVfqqBE=
+    """
+
+
+class KeyIndex(TypedDict):
+    AppID: str
+    """
+    指定要获取的密钥Owner的AppID （PSM）
+    Example: security.research.tks_owner
+    """
+
+    RingID: str
+    """
+    密钥环ID
+    Example: 640a7c2a-c061-4964-b7d8-951e9272da21
+    """
+
+    KeyID: str
+    """
+    指定要获取的密钥的ID
+    Example: d880cf63-7c3e-45e1-83e5-ab62ba3d3441
+    """
+
+
+class KeyIndexWithVersion(KeyIndex):
+    Version: int
+    """
+    返回Key的版本，从1开始
+    Example: 1
+    """
+
+
+class ChallengeResponse(TypedDict):
+    ClientChall: NotRequired[ClientChallenge]
+    """
+    在远程阶段获取到的，由TKS返回的挑战信息。
+    """
+
+    RAEvidence: NotRequired[RAEvidence]
+    """
+    提供基于TEE远程证明相关的证明信息，当访问控制策略中包含RA类型时需要提供
+    """
+
+    ZTIEvidence: NotRequired[ZTIEvidence]
+    """
+    提供基于ZTI相关的证明信息，当访问控制策略中包含ZTI类型时需要提供
+    """
+
+    CertEvidence: NotRequired[CertEvidence]
+    """
+    提供基于证书相关的证明信息，当访问控制策略中包含Cert类型时需要提供
+    """
+
+
+class KeyExportation(KeyIndex, ChallengeResponse):
+    DHParam: str
+    """
+    经过Base64编码的，请求方的DH公共参数
+    Example: KGIW9H8TCIi+ZNwikN8cc4V2cTNzOi/VfLNF7oYBwvE=
+    """
+
+    KMSInfo: NotRequired[KMSInfo]
+    """
+    描述密钥所在的KMS位置，当前仅支持TKS
+    """
+
+    Passport: NotRequired[str]
+    """
+    JWT Token。由TKS签发的访问Token，获取Token的参与方在验证通过可以直接获取密钥。
+    """
+
+
+class ExportationResult(KeyIndexWithVersion):
+    Key: str
+    """
+    经过Base64编码后的Key的密文，通过请求参数密钥协商出来的密钥加密（AES-256-GCM）
+    Example: aq1xcWeuyyKTDY8kXqgqk9TmBjhqFhwZB7+p0GgK60I=
+    """
+
+
+class KeyCreation(TypedDict):
+    RingID: str
+    """
+    密钥环ID（为空情况下，放置到用户的一个default密钥环下）
+    Example: 640a7c2a-c061-4964-b7d8-951e9272da21
+    """
+
+    KeyID: NotRequired[str]
+    """
+    密钥ID，置空则由TKS随机分配一个UUID
+    Example: d880cf63-7c3e-45e1-83e5-ab62ba3d3441
+    """
+
+    KeyName: str
+    """
+    密钥名称
+    Example: test-demo-key
+    """
+
+    KeyAlgo: str
+    """
+    密钥算法（当前仅支持基于对称加密算法）
+    Example: SYMMTRIC_128 | SYMMTRIC_256
+    """
+
+    Source: Literal["external", "internal"]
+    """
+    密钥来源
+    - external: 外部导入
+    - internal: 由TKS生成
+    Example: internal
+    """
+
+    Description: NotRequired[str]
+    """
+    密钥描述信息
+    Example: Key for test
+    """
+
+    Lifetime: NotRequired[int]
+    """
+    密钥生命周期，Unix时间戳（UTC+8）。默认情况为-1，表示无限期
+    Example: 1722254231
+    """
+
+    KMSInfo: NotRequired[KMSInfo]
+    """
+    描述密钥托管位置
+    """
+
+    ACLRules: NotRequired[str]
+    """
+    描述密钥的访问控制策略，JSON字符串，默认情况下允许所有人访问，即
+    {"default": "accept","rules": [],"range": "keyID","issueToken": "false"}
+    Example: {"type": "RA", "tee": "tdx", "measurements":
+              {"tdx.rtmr0": "561bf3eee22f", "tdx.rtmr1": "d4247b6bd6eb"}}
+    """
+
+
+class KeyImportation(TypedDict):
+    RingID: str
+    """
+    密钥环ID（为空情况下，放置到用户的一个default密钥环下
+    Example: 640a7c2a-c061-4964-b7d8-951e9272da21
+    """
+
+    KeyID: str
+    """
+    密钥ID，置空则由TKS随机分配一个UUID
+    Example: d880cf63-7c3e-45e1-83e5-ab62ba3d3441
+    """
+
+    Key: str
+    """
+    密钥密文。经过Base64编码后的密钥密文，加密密钥应通过AES-GCM-256加密，
+    AES密钥基于TKS公共参数协商得出（ECDH-SECP384R1-SHA256）。可为空，置空后密钥状态为empty
+    Example: wjBXVnzCFOc/NhQ6vXZhllGa1Kirs9VvE+GVSuYCEZM=
+    """
+
+    DHParam: str
+    """
+    经过Base64编码后的请求方侧的密钥协商参数，帮助TKS派生出通信密钥（ECDH-SECP384R1-SHA256）
+    Example: wjBXVnzCFOc/NhQ6vXZhllGa1Kirs9VvE+GVSuYCEZM=
+    """
+
+
+@dataclass
+class AppOptions:
+    """特定于 App 的认证选项."""
+
+    id: str
+    """App ID."""
+
+    signing_key: Optional[str] = None
+    """PEM 格式的 App 私钥. 当 TKS 中 authType 设置为 cert 时使用."""
+
+    zti_token: Optional[str] = None
+    """ZTI 令牌. 当 TKS 中 authType 设置为 zti 时使用."""

src/bytedance/jeddak_secure_channel/utils.py

@@ -0,0 +1,469 @@
+__all__ = [
+    "RepeatTimer",
+    "WeakMethod",
+    "weak_repeat_timer",
+    "request_bytedance_gateway",
+]
+
+import hashlib
+import hmac
+import json
+import os
+import urllib.parse
+import weakref
+from datetime import datetime, timezone
+from threading import Timer
+from typing import Dict
+
+import requests
+import volcenginesdkcore
+import volcenginesdksts
+from typing_extensions import (
+    Any,
+    Callable,
+    Concatenate,
+    Generic,
+    Iterable,
+    Mapping,
+    Optional,
+    ParamSpec,
+    Protocol,
+    TypeVar,
+)
+from volcenginesdkcore.rest import ApiException
+
+from . import error
+from .log import logger
+
+
+class RepeatTimer(Timer):
+    """
+    定时重复执行函数.
+    """
+
+    def run(self):
+        while not self.finished.wait(self.interval):
+            self.function(*self.args, **self.kwargs)
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+R = TypeVar("R")
+R_co = TypeVar("R_co", covariant=True)
+
+
+class MethodProtocol(Protocol[T, P, R]):
+    __self__: T
+
+    __func__: Callable[Concatenate[T, P], R]
+
+    __call__: Callable[P, R]
+
+
+class WeakMethod(Generic[T, P, R_co]):
+    """
+    对方法 (bound method) 的弱引用.
+    """
+
+    # instance: weakref.ReferenceType[T]
+    # func: weakref.ReferenceType[Callable[Concatenate[T, P], R_co]]
+
+    def __init__(self, method: Callable[P, R_co]):
+        # Type checkers do not differentiate methods from callables, so
+        # we have to ignore the casting error
+        _method: MethodProtocol[T, P, R_co] = method  # pyright: ignore
+
+        self.instance = weakref.ref(_method.__self__)
+        self.func = weakref.ref(_method.__func__)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Optional[R_co]:
+        instance = self.instance()
+        func = self.func()
+
+        if instance is None or func is None:
+            logger.info("Weak method called on freed object")
+            return
+
+        return func(instance, *args, **kwargs)
+
+
+def weak_repeat_timer(
+    interval: float,
+    method: Callable[P, None],
+    args: Optional[Iterable[Any]] = None,
+    kwargs: Optional[Mapping[str, Any]] = None,
+) -> RepeatTimer:
+    """
+    定时重复执行对象上的方法. 在对象被垃圾回收时自动结束线程.
+    """
+    # Due to limitations of 'ParamSpec', we cannot annotate 'args' and 'kwargs' precisely
+    timer = RepeatTimer(interval, WeakMethod(method), args, kwargs)
+    weakref.finalize(method.__self__, timer.cancel)  # pyright: ignore
+
+    return timer
+
+
+class TopRequester:
+    def __init__(
+        self,
+        ak,
+        sk,
+        token=None,
+        host="open.volcengineapi.com",
+        region="cn-beijing",
+        use_tls=True,
+        service="pcc",
+        version="2024-12-24",
+        debug=False,
+    ):
+        volc_config = volcenginesdkcore.Configuration()
+        volc_config.host = host
+        volc_config.ak = ak
+        volc_config.sk = sk
+        volc_config.session_token = token
+        volc_config.region = region
+        volc_config.debug = debug
+        if use_tls:
+            volc_config.scheme = "https"
+        else:
+            volc_config.scheme = "http"
+        self._api_client = volcenginesdkcore.ApiClient(configuration=volc_config)
+        self._service = service
+        self._version = version
+
+    def _construct_resource_path(
+        self,
+        action: str,
+        version: str,
+        service: str,
+        method: str = "post",
+        content_type: str = "application_json",
+    ) -> str:
+        return f"/{action}/{version}/{service}/{method.lower()}/{content_type}/"
+
+    def request(
+        self,
+        method,
+        query,
+        header,
+        param,
+        action,
+        body,
+        service=None,
+        version=None,
+        content_type="application/json",
+    ):
+        resp = self._api_client.call_api(
+            self._construct_resource_path(
+                action=action,
+                version=self._version if not version else version,
+                service=self._service if not service else service,
+                method=method,
+                content_type=content_type.replace("/", "_"),
+            ),
+            method,
+            body=body,
+            auth_settings=["volcengineSign"],
+            header_params={
+                "Accept": "application/json",
+                "Content-Type": content_type,
+                **header,
+            },
+            path_params=param,
+            query_params=query,
+            post_params=[],
+            files={},
+            response_type=bytes,
+            collection_formats={},
+            _return_http_data_only=False,
+            _preload_content=True,
+        )
+
+        status, headers, data = resp[1], resp[2], resp[0]
+        logger.info(f"Top request response,status={status}, headers: {headers}")
+
+        return data
+
+    def request_with_assume_role(
+        self,
+        method,
+        query,
+        header,
+        param,
+        service,
+        action,
+        body,
+        version,
+        target_uid,
+        target_resource=None,
+        region="cn-beijing",
+        content_type="application/json",
+        use_tls=True,
+        host="open.volcengineapi.com",
+    ):
+        data = {
+            "TargetAccount": target_uid,
+            "TargetResource": os.getenv("PCC_SERVICE", "pcc")
+            if not target_resource
+            else target_resource,
+        }
+        access_info = self.request("POST", [], {}, {}, "UpdateUserAccessInfo", data)
+        ak = access_info["Ak"]
+        sk = access_info["Sk"]
+        sts_token = access_info["StsToken"]
+
+        customer_requester = TopRequester(
+            ak=ak, sk=sk, token=sts_token, region=region, use_tls=use_tls, host=host
+        )
+        return customer_requester.request(
+            method, query, header, param, action, body, service, version, content_type
+        )
+
+
+def _hmac_sha256(key: bytes, content: str):
+    return hmac.new(key, content.encode(), hashlib.sha256).digest()
+
+
+def _hash_sha256(content: str):
+    return hashlib.sha256(content.encode()).hexdigest()
+
+
+def _format_query(params: Dict[str, str]) -> str:
+    return "&".join(
+        urllib.parse.quote(key, safe="-_.~")
+        + "="
+        + urllib.parse.quote(value, safe="-_.~")
+        for key, value in params.items()
+    )
+
+
+def top_request(
+    ra_url: str,
+    body: Any,
+    additional_headers: Optional[Dict[str, str]],
+    top_config: Dict[str, str],
+) -> requests.Response:
+    top_ak = top_config["ak"]
+    top_sk = top_config["sk"]
+    top_service = top_config["service"]
+    top_region = top_config.get("region", "cn-beijing")
+    top_method = top_config.get("method", "POST").upper()
+    top_action = top_config.get("action", "GetAttestationBackend")
+    top_version = top_config.get("version", "2024-12-24")
+    use_tls: bool = top_config.get("use_tls", True)  # type: ignore
+    sts_token = top_config.get("sts_token", None)
+    target_uid = top_config.get("target_uid", None)
+    target_resource = top_config.get("target_resource", "pcc")
+
+    if ra_url:
+        top_host = ra_url
+    else:
+        top_host = top_config.get("url", "open.volcengineapi.com")
+
+    requester = TopRequester(
+        ak=top_ak,
+        sk=top_sk,
+        token=sts_token,
+        region=top_region,
+        use_tls=use_tls,
+        host=top_host,
+        service=top_service,
+        version=top_version,
+    )
+
+    resp = None
+    if not target_uid:
+        try:
+            resp = requester.request(
+                top_method, [], additional_headers, {}, top_action, body
+            )
+        except ApiException as e:
+            logger.error(
+                "status: {}, reason: {}, body: {}, headers: {}".format(
+                    e.status, e.reason, e.body, e.headers
+                )
+            )
+            raise e
+        except Exception as e:
+            logger.error("request error: {}".format(e))
+            raise e
+    else:
+        try:
+            resp = requester.request_with_assume_role(
+                top_method,
+                [],
+                additional_headers,
+                {},
+                top_service,
+                top_action,
+                body,
+                top_version,
+                target_uid,
+                target_resource,
+                top_region,
+                use_tls=use_tls,
+                host=top_host
+            )
+        except ApiException as e:
+            logger.error(
+                "status: {}, reason: {}, body: {}, headers: {}".format(
+                    e.status, e.reason, e.body, e.headers
+                )
+            )
+            raise e
+        except Exception as e:
+            logger.error("request error: {}".format(e))
+            raise e
+
+    return resp
+
+
+def get_top_access_key(top_config: Dict[str, str]) -> tuple:
+    try:
+        configuration = volcenginesdkcore.Configuration()
+        configuration.ak = top_config.get("ak")
+        configuration.sk = top_config.get("sk")
+        configuration.region = top_config.get("region", "cn-beijing")
+
+        duration_seconds = int(top_config.get("duration_seconds", "3600"))
+        role_session_name = top_config.get("role_session_name", "temp_")
+        role_trn = top_config.get("aicc_saas_trn")
+
+        volcenginesdkcore.Configuration.set_default(configuration)
+
+        api_instance = volcenginesdksts.STSApi()
+        assume_role_request = volcenginesdksts.AssumeRoleRequest(
+            duration_seconds=duration_seconds,
+            role_session_name=role_session_name,
+            role_trn=role_trn,
+        )
+
+        ret = api_instance.assume_role(assume_role_request)
+        if ret and ret.credentials:
+            ak = ret.credentials.access_key_id
+            sk = ret.credentials.secret_access_key
+            session_token = ret.credentials.session_token
+
+            return ak, sk, session_token
+        else:
+            return None, None, None
+
+    except Exception as e:
+        logger.error("get_top_access_key: {}".format(e))
+        raise e
+
+
+# The first parameter was url, and is now unused
+def request_bytedance_gateway(
+    url: str,
+    body: Any,
+    additional_headers: Optional[Dict[str, str]],
+    top_config: Dict[str, str],
+) -> requests.Response:
+    top_ak = top_config["ak"]
+    top_sk = top_config["sk"]
+
+    if not additional_headers:
+        additional_headers = dict()
+
+    if top_config.get("aicc_saas_trn"):
+        ak, sk, session_token = get_top_access_key(top_config)
+        if ak and sk and session_token:
+            top_ak = ak
+            top_sk = sk
+            additional_headers["X-Security-Token"] = session_token
+            additional_headers["sts_token"] = session_token
+        if top_config.get("target_uid"):
+            additional_headers["UID"] = top_config.get("target_uid")
+
+    top_service = top_config["service"]
+    top_region = top_config.get("region", "cn-beijing")
+    top_method = top_config.get("method", "POST").upper()
+    top_action = top_config.get("action", "GetAttestationBackend")
+    top_version = top_config.get("version", "2024-12-24")
+    url_rewrite = top_config.get("url_rewrite", "")
+    http_scheme = top_config.get("http_scheme", "https")
+
+    top_url = top_config.get("url", "") if top_config.get("url", "") else url
+    if not top_url:
+        top_url = url_rewrite
+    if not top_url:
+        logger.error("ra request error url=None")
+        raise error.ParamError("top_url", "top_url is None")
+
+    if top_url.startswith("http"):
+        url_parsed = urllib.parse.urlparse(top_url)
+        top_host = url_parsed.hostname
+    else:
+        top_host = top_url
+
+    params = {"Action": top_action, "Version": top_version}
+    query_str = _format_query(params)
+
+    date_now = datetime.now(timezone.utc)
+    body_str = json.dumps(body, separators=(",", ":"))
+
+    x_date = date_now.strftime("%Y%m%dT%H%M%SZ")
+    short_x_date = x_date[:8]
+    x_content_sha256 = _hash_sha256(body_str)
+
+    headers = {
+        "Content-Type": "application/json",
+        "Host": top_host,
+        "X-Content-Sha256": x_content_sha256,
+        "X-Date": x_date,
+    }
+
+    signed_headers = ";".join(key.lower() for key in headers.keys())
+
+    canonical_request = "\n".join(
+        [
+            top_method,
+            "/",
+            query_str,
+            *[key.lower() + ":" + value for key, value in headers.items()],
+            "",
+            signed_headers,
+            x_content_sha256,
+        ]
+    )
+
+    canonical_request_hash = _hash_sha256(canonical_request)
+    credential_scope = "/".join([short_x_date, top_region, top_service, "request"])
+    string_to_sign = "\n".join(
+        ["HMAC-SHA256", x_date, credential_scope, canonical_request_hash]
+    )
+
+    logger.debug(
+        f"ra request={canonical_request}, hash={canonical_request_hash}, sign={string_to_sign}"
+    )
+
+    k_date = _hmac_sha256(top_sk.encode(), short_x_date)
+    k_region = _hmac_sha256(k_date, top_region)
+    k_service = _hmac_sha256(k_region, top_service)
+    k_signing = _hmac_sha256(k_service, "request")
+    signature = _hmac_sha256(k_signing, string_to_sign).hex()
+
+    headers["Authorization"] = (
+        f"HMAC-SHA256 Credential={top_ak}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}"
+    )
+
+    if url_rewrite:
+        url_parsed = urllib.parse.urlparse(url_rewrite)
+        headers["Host"] = url_parsed.hostname or ""
+        url_final = url_rewrite + "?" + query_str
+    else:
+        if top_url.startswith("http"):
+            url_final = top_url + "?" + query_str
+        else:
+            url_final = f"{http_scheme}://{top_url}/?{query_str}"
+
+    if additional_headers:
+        headers = {**headers, **additional_headers}
+
+    r = requests.request(
+        method=top_method, url=url_final, headers=headers, data=body_str, timeout=30
+    )
+    logger.debug(f"call ra: url={url_final}, header={headers}, body={body_str}, res={json.loads(r.text).get('Result')}")
+    logger.info(f"ra response headers: X-Tt-Logid={r.headers.get('X-Tt-Logid')}")
+    return r

src/bytedance/tks/attester.py

@@ -0,0 +1,41 @@
+__all__ = [
+    "Attester",
+    "EPSAttester",
+]
+
+import os
+from abc import ABC, abstractmethod
+
+import requests
+
+from . import exceptions as exp
+
+
+class Attester(ABC):
+    @abstractmethod
+    def get_evidence(self, report_data: str) -> bytes:
+        raise NotImplementedError
+
+
+class EPSAttester(Attester):
+    def __init__(self) -> None:
+        pass
+
+    def get_evidence(self, report_data: str, pod_name: str = "", attest_gpu: bool = True) -> bytes:
+        host_ip = os.environ.get("JPCC_HOST_IP", "localhost")
+        eps_addr = f"{host_ip}:8006"
+        if not pod_name:
+            pod_name = os.environ.get("JPCC_POD_NAME")
+        if not pod_name:
+            pod_name = os.environ.get("HOSTNAME", "")
+
+        url = "http://{}/aa/evidence?runtime_data={}&pod_name={}&attest_gpu={}&static_report=true".format(
+            eps_addr, report_data, pod_name, str(attest_gpu).lower())
+
+        try:
+            resp = requests.get(url=url)
+            resp.raise_for_status()
+        except requests.RequestException as e:
+            raise exp.RAError("Bad connection to EPS '{}': {}".format(url, e)) from e
+
+        return resp.content

src/bytedance/tks/client.py

@@ -0,0 +1,628 @@
+import base64
+import dataclasses
+import hashlib
+import os
+import secrets
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+import requests
+from typing_extensions import IO, TYPE_CHECKING, Any, Dict, Optional, Union
+from volcenginesdkcore.rest import ApiException
+from bytedance.jeddak_secure_channel.ra import verify_ra_result, gen_report_data
+
+from . import exceptions as exp
+from .attester import EPSAttester
+from .crypto import aes, ec
+
+if TYPE_CHECKING:
+    from _typeshed import GenericPath
+else:
+    GenericPath = str
+
+__all__ = [
+    "TKSClient",
+    "TKSConfig",
+]
+
+PATH_ATTEST = "/api/tks/v1/security/attest"
+PATH_RING_CREATE = "/api/tks/v1/ring/create"
+PATH_KEY_EXPORT = "/api/tks/v1/key/export"
+PATH_KEY_CREATE = "/api/tks/v1/key/create"
+PATH_KEY_IMPORT = "/api/tks/v1/key/import"
+PATH_POLICY_SET = "/api/tks/v1/policy/set"
+PATH_KEY_LIST = "/api/tks/v1/key/list"
+PATH_RING_LIST = "/api/tks/v1/ring/list"
+PATH_KEY_INFO = "/api/tks/v1/key/getinfo"
+PATH_RAS_VERIFY = "/api/ras/v1/verify"
+
+
+def utc8_timestamp() -> int:
+    return int(datetime.now(timezone(timedelta(hours=8))).timestamp())
+
+
+@dataclass
+class TKSConfig:
+    addr: str = "jeddakchain.bytedance.com"
+    pod_name: str = ""
+    enable_tls: bool = False
+    pcc_config: dict = dataclasses.field(default_factory=dict)
+    top_config: dict = dataclasses.field(default_factory=dict)
+
+
+class TKSClient:
+    def __init__(self, app_id: str, config: TKSConfig) -> None:
+        self._app_id = app_id
+        self._config = config
+        self.attester = EPSAttester()
+        self.CHUNK_SIZE = 1024
+
+    def _request_tks(
+            self,
+            url: str,
+            body: Any,
+            headers: Optional[Dict[str, str]] = None,
+            cookies: Optional[Dict[str, str]] = None,
+    ) -> dict:
+        try:
+            if self._config.top_config:
+                if self._config.top_config.get("url_rewrite", ""):
+                    from bytedance.jeddak_secure_channel.utils import (
+                        request_bytedance_gateway,
+                    )
+
+                    resp = request_bytedance_gateway(
+                        url, body, headers, self._config.top_config
+                    )
+                else:
+                    from bytedance.jeddak_secure_channel.utils import top_request
+
+                    resp = top_request(url, body, headers, self._config.top_config)
+            else:
+                resp = requests.post(
+                    url=url, json=body, headers=headers, cookies=cookies
+                )
+        except ApiException as e:
+            raise exp.TKSError(
+                "Status '{}' from TKS, URL: '{}', Content: {}".format(
+                    e.status, url, e.body
+                ),
+                http_status=e.status
+            ) from e
+        except Exception as e:
+            raise exp.TKSError("Network error") from e
+        try:
+            if isinstance(resp, requests.Response):
+                resp.raise_for_status()
+                resp = resp.json()
+        except requests.HTTPError as e:
+            raise exp.TKSError(
+                "Status '{}' from TKS, URL: '{}', Content: {}".format(
+                    resp.status_code, url, resp.text
+                ),
+                http_status=resp.status_code
+            ) from e
+        except requests.RequestException as e:
+            raise exp.TKSError("Bad connection to TKS: {}".format(e)) from e
+
+        # if "Result" not in resp:
+        #     raise exp.TKSError("Unexpected response: {}".format(resp))
+        res = resp["Result"] if "Result" in resp else resp
+
+        return res
+
+    def challenge_server(self, nonce: str) -> str:
+        return self._do_server_ra(nonce)["Report"]
+
+    def create_ring(self, ring_name: str, desc: str = "", **kwargs) -> dict:
+        try:
+            if self._config.top_config:
+                url = self._config.addr
+                self._config.top_config["action"] = "CreateTksRing"
+            else:
+                url = self._get_tks_url(PATH_RING_CREATE)
+
+            result = self._request_tks(
+                url=url,
+                body={"RingName": ring_name, "Description": desc},
+                headers=self._get_headers(),
+                **kwargs,
+            )
+            return result
+        except requests.RequestException as e:
+            raise exp.TKSError("create_ring to TKS: {}".format(e)) from e
+
+    def create_key(
+            self,
+            ring_id: str,
+            algo: str,
+            key_name: str = "",
+            desc: str = "",
+            **kwargs,
+    ) -> dict:
+        key_meta = self._do_key_creation(
+            ring_id, algo=algo, key_name=key_name, src="internal", desc=desc, **kwargs
+        )
+        return key_meta
+
+    def import_key(
+            self,
+            ring_id: str,
+            algo: str,
+            key: bytes,
+            key_name: str = "",
+            desc: str = "",
+            usage_scenario="ModelEncryption",
+            **kwargs,
+    ) -> dict:
+        # Create an empty key record
+        key_meta = self._do_key_creation(
+            ring_id,
+            algo=algo,
+            key_name=key_name,
+            src="external",
+            desc=desc,
+            usage_scenario=usage_scenario,
+            **kwargs,
+        )
+
+        # Locally generate key pair
+        sk_local, pk_local = ec.generate_key_pair()
+
+        # Challenge server
+        ra_result = self._do_server_ra(bi_auth=False)
+        pk_server = base64.b64decode(ra_result["DHParam"].encode())
+
+        # Encrypt key
+        ek = ec.diffie_hellman_key_exchange(sk_local, pk_server)
+        encrypted_key = aes.gcm_encrypt(ek, key)
+
+        if self._config.top_config:
+            url = self._config.addr
+            self._config.top_config["action"] = "ImportTksKey"
+        else:
+            url = self._get_tks_url(PATH_KEY_IMPORT)
+
+        _ = self._request_tks(
+            url=url,
+            body={
+                "RingID": ring_id,
+                "KeyID": key_meta["KeyID"],
+                "Key": base64.b64encode(encrypted_key).decode(),
+                "DHParam": base64.b64encode(pk_local).decode(),
+            },
+            headers=self._get_headers(),
+            **kwargs,
+        )
+
+        return key_meta
+
+    def get_key_info(
+            self,
+            ring_id: str,
+            key_id: str,
+    ) -> dict:
+        if self._config.top_config:
+            url = self._config.addr
+            self._config.top_config["action"] = "GetTksKeyinfo"
+        else:
+            url = self._get_tks_url(PATH_KEY_INFO)
+        result = self._request_tks(
+            url=url,
+            body={
+                "RingID": ring_id,
+                "KeyID": key_id,
+            },
+            headers=self._get_headers(),
+        )
+
+        return result
+
+    def get_key(
+            self,
+            ring_id: str,
+            key_id: str,
+            need_evidence: bool = True,
+            attest_gpu: bool = True,
+            policy_id='',
+            tks_url='',
+            tks_uid='',
+            tks_service_name='',
+            **kwargs,
+    ) -> bytes:
+        print(f'tks get_key: policy_id={policy_id}, url={tks_url}, uid={tks_uid}, ra_service_name={tks_service_name}')
+        nonce = secrets.token_hex(32)
+        # 获取TKS的证明报告
+        ra_result = self._do_server_ra(nonce=nonce, bi_auth=True)
+        # 调用verify接口对TKS进行远程证明
+        if policy_id:
+            if self._config.top_config:
+                url = self._config.addr
+                self._config.top_config["action"] = "VerifyRasEvidence"
+            else:
+                url = self._get_tks_url(PATH_RAS_VERIFY)
+            report_data_bytes = gen_report_data(nonce, ra_result["Challenge"]["NonceDown"], ra_result["DHParam"],
+                                                ra_result['Challenge']['ExpirationTime'])
+
+            urlsafe_base64 = base64.urlsafe_b64encode(report_data_bytes)
+            report_data = urlsafe_base64.rstrip(b'=').decode('utf-8')
+            req_body = {
+                "Evidence": ra_result.get("Report").rstrip('='),
+                "PolicyId": policy_id,
+                "ReportData": {"Raw": report_data}
+            }
+            res = self._request_tks(
+                url=url, body=req_body, headers=self._get_headers(), **kwargs
+            )
+            verify_res = verify_ra_result(res.get("AttestationResult"), None, "", None, nonce,
+                                          report_data_bytes.hex())
+            if not verify_res:
+                raise Exception("verify tks failed!!!")
+            print("verify tks success!!!")
+
+        # get pub_key from server
+        pk_server = base64.b64decode(ra_result["DHParam"].encode())
+
+        sk_local, pk_local = ec.generate_key_pair()
+        req_body = {
+            "AppID": self._app_id,
+            "RingID": ring_id,
+            "KeyID": key_id,
+            "DHParam": base64.b64encode(pk_local).decode(),
+            "ClientChall": ra_result["Challenge"],
+            # "RAEvidence": {"TEEType": "coco", "Report": base64.b64encode(evidence).decode()},
+        }
+
+        if need_evidence:
+            evidence = self.attester.get_evidence(
+                ra_result["Challenge"]["NonceDown"][:64],
+                self._config.pod_name,
+                attest_gpu,
+            )
+
+            req_body["RAEvidence"] = {
+                "TEEType": "coco",
+                "Report": base64.b64encode(evidence).decode(),
+            }
+        if self._config.top_config:
+            url = self._config.addr
+            self._config.top_config["action"] = "ExportTksKey"
+        else:
+            url = self._get_tks_url(PATH_KEY_EXPORT)
+
+        res = self._request_tks(
+            url=url, body=req_body, headers=self._get_headers(), **kwargs
+        )
+
+        encrypted_key = res["Key"]
+
+        ek = ec.diffie_hellman_key_exchange(local_sk=sk_local, peer_pk=pk_server)
+
+        data_key = aes.gcm_decrypt(ek, base64.b64decode(encrypted_key))
+
+        return data_key
+
+    def set_key_policy(self, ring_id: str, key_id: str, policy: str, **kwargs) -> dict:
+        if self._config.top_config:
+            url = self._config.addr
+            self._config.top_config["action"] = "ImportTksKey"
+        else:
+            url = self._get_tks_url(PATH_POLICY_SET)
+
+        result = self._request_tks(
+            url=url,
+            headers=self._get_headers(),
+            body={
+                "ID": key_id,
+                "Range": "key",
+                "Rules": policy,
+            },
+            **kwargs,
+        )
+        return result
+
+    def _do_key_creation(
+            self,
+            ring_id: str,
+            algo: str,
+            key_name: str = "",
+            src: str = "internal",
+            desc: str = "",
+            usage_scenario: str = "ModelEncryption",
+            **kwargs,
+    ) -> dict:
+        if self._config.top_config:
+            url = self._config.addr
+            self._config.top_config["action"] = "CreateTksKey"
+        else:
+            url = self._get_tks_url(PATH_KEY_CREATE)
+        result = self._request_tks(
+            url=url,
+            body={
+                "RingID": ring_id,
+                "KeyName": key_name,
+                "Algo": algo,
+                "Source": src,
+                "Description": desc,
+                "UsageScenario": usage_scenario,
+            },
+            headers=self._get_headers(),
+            **kwargs,
+        )
+
+        return result
+
+    def list_ring_id(
+            self,
+            page_number: int = 0,
+            page_size: int = 10,
+            show_invisible: bool = False,
+            **kwargs,
+    ):
+        """
+        :param page_size:
+        :param page_number:
+        :param show_invisible:
+        :param kwargs:
+        :return:
+        """
+        try:
+            if self._config.top_config:
+                url = self._config.addr
+                self._config.top_config["action"] = "ListTksRing"
+            else:
+                url = self._get_tks_url(PATH_RING_LIST)
+
+            body = {
+                "PageSize": page_size,
+                "PageNumber": page_number,
+                "ShowInvisible": show_invisible,
+            }
+
+            result = self._request_tks(
+                url=url, body=body, headers=self._get_headers(), **kwargs
+            )
+            return result
+        except requests.RequestException as e:
+            raise exp.TKSError("list ring id error: {}".format(e)) from e
+
+    def list_key(
+            self,
+            ring_id: Optional[str] = None,
+            page_number: int = 0,
+            page_size: int = 10,
+            status: Optional[str] = None,
+            **kwargs,
+    ):
+        """
+        :param ring_id:
+        :param page_size:
+        :param page_number:
+        :param status:
+        :param kwargs:
+        :return:
+        """
+        try:
+            if self._config.top_config:
+                url = self._config.addr
+                self._config.top_config["action"] = "ListTksKey"
+            else:
+                url = self._get_tks_url(PATH_KEY_LIST)
+
+            body: dict = {"PageSize": page_size, "PageNumber": page_number}
+            if ring_id:
+                body["RingID"] = ring_id
+            if status:
+                body["Status"] = status
+
+            result = self._request_tks(
+                url=url, body=body, headers=self._get_headers(), **kwargs
+            )
+            return result
+        except requests.RequestException as e:
+            raise exp.TKSError("list key error: {}".format(e)) from e
+
+    def _do_server_ra(
+            self, nonce: Optional[str] = None, bi_auth: bool = True, **kwargs
+    ) -> dict:
+        if nonce is None:
+            nonce = secrets.token_hex(32)
+        if self._config.top_config:
+            url = self._config.addr
+            self._config.top_config["action"] = "AttestTksSecurity"
+        else:
+            url = self._get_tks_url(PATH_ATTEST)
+
+        resp = self._request_tks(
+            url=url,
+            body={"NonceUp": nonce, "BiAuth": bi_auth},
+            headers=self._get_headers(),
+            **kwargs,
+        )
+
+        return resp
+
+    def _get_tks_url(self, path: str) -> str:
+        if self._config.addr.startswith("http"):
+            return "{}{}".format(self._config.addr, path)
+        else:
+            scheme = "https" if self._config.enable_tls else "http"
+            return "{}://{}{}".format(scheme, self._config.addr, path)
+
+    def _generate_pcc_token(self, timestamp: str) -> str:
+        """
+        生成访问bytedance pcc_controller的token
+        :param timestamp:
+        :return:
+        """
+        ak = self._config.pcc_config["ak"]
+        sk = self._config.pcc_config["sk"]
+
+        h = hashlib.sha256()
+        h.update(sk.encode())
+        sk = h.hexdigest()
+        h = hashlib.sha256()
+        h.update((ak + sk + str(timestamp)).encode())
+        token = h.hexdigest()
+        return token
+
+    def _get_headers(self) -> dict:
+        timestamp = str(int(datetime.now(timezone.utc).timestamp()))
+        headers = {
+            "AppID": self._app_id,
+            # "Timestamp": str(utc8_timestamp()),
+            "Timestamp": timestamp,
+            # "nonce": "a314b43502077cd89d21378b29954207aaf5b3848d4d7fa9f16682ba1f6f3317"
+        }
+
+        if self._config.pcc_config:
+            headers["Content-Type"] = "application/json"
+            headers["J-Internal-Ak"] = self._config.pcc_config["ak"]
+            headers["J-Internal-Token"] = self._generate_pcc_token(timestamp)
+
+        return headers
+
+    def aes_encrypt(self, aes_key: Union[bytes, str], data: Union[bytes, str]) -> bytes:
+        if isinstance(aes_key, str):
+            aes_key = aes_key.encode()
+        if isinstance(data, str):
+            data = data.encode()
+        return aes.gcm_encrypt(aes_key, data)
+
+    def aes_decrypt(self, aes_key: Union[bytes, str], data: Union[bytes, str]) -> bytes:
+        if isinstance(aes_key, str):
+            aes_key = aes_key.encode()
+        if isinstance(data, str):
+            data = data.encode()
+        return aes.gcm_decrypt(aes_key, data)
+
+    def encrypt_file(
+            self,
+            aes_key: Union[bytes, str],
+            source_path: str,
+            dest_path: str,
+            mode: str = "b",
+    ) -> bool:
+        if mode not in ["b", "t"]:
+            raise Exception("mode argument error")
+
+        if isinstance(aes_key, str):
+            aes_key = aes_key.encode()
+
+        with open(source_path, f"r{mode}") as source, open(
+                dest_path, f"w{mode}"
+        ) as dest:
+            if mode == "b":
+                source_b: IO[bytes] = source
+                dest_b: IO[bytes] = dest
+                plaintext = source_b.read()
+                ciphertext = self.aes_encrypt(aes_key, plaintext)
+                dest_b.write(ciphertext)
+            else:
+                source_t: IO[str] = source
+                dest_t: IO[str] = dest
+                for plaintext in source_t:
+                    plaintext = plaintext.rstrip("\n")
+                    ciphertext = self.aes_encrypt(aes_key, plaintext)
+                    dest_t.write(base64.b64encode(ciphertext).decode() + "\n")
+
+        return True
+
+    def decrypt_file(
+            self,
+            aes_key: Union[bytes, str],
+            source_path: str,
+            dest_path: str,
+            mode: str = "b",
+    ) -> bool:
+        if mode not in ["b", "t"]:
+            raise Exception("mode argument error")
+
+        if isinstance(aes_key, str):
+            aes_key = aes_key.encode()
+
+        with open(source_path, f"r{mode}") as source, open(
+                dest_path, f"w{mode}"
+        ) as dest:
+            if mode == "b":
+                source_b: IO[bytes] = source
+                dest_b: IO[bytes] = dest
+                ciphertext = source_b.read()
+                plaintext = self.aes_decrypt(aes_key, ciphertext)
+                dest_b.write(plaintext)
+            else:
+                source_t: IO[str] = source
+                dest_t: IO[str] = dest
+                for ciphertext in source_t:
+                    ciphertext = ciphertext.rstrip("\n")
+                    ciphertext = base64.b64decode(ciphertext.encode())
+                    plaintext = self.aes_decrypt(aes_key, ciphertext)
+                    dest_t.write(plaintext.decode() + "\n")
+                    # dest_t.write(base64.b64encode(plaintext).decode() + "\n")
+        return True
+
+
+default_config = TKSConfig()
+
+if __name__ == "__main__":
+    def md5_hash(data: bytes, salt: str = "") -> str:
+        return hashlib.md5(data + salt.encode()).hexdigest()
+
+
+    ts = str(utc8_timestamp())
+    req_body = {"Username": "alice", "Password": md5_hash("alice0".encode(), ts)}
+    response = requests.post(
+        "https://jeddakchain.bytedance.com/api/user/login",
+        headers={"TimeStamp": ts},
+        json=req_body,
+    )
+    cookies = response.cookies
+    jeddak_pcc_session = cookies.get("jeddak_pcc_session")
+
+    cookies = {"jeddak_pcc_session": jeddak_pcc_session}
+
+    cfg = TKSConfig(addr="jeddakchain.bytedance.com", enable_tls=True)
+    client = TKSClient("alice", cfg)
+    # print(
+    #     "Key Get",
+    #     client.get_key(
+    #         "2ee6348f-fafa-4e97-a824-0b6ed0e4d07e", "f532b84c-8f80-4705-b5a7-6217f4729924"
+    #     ),
+    # )
+    #
+    # print("Challenge Server", client.challenge_server("1234"))
+    # print(
+    #     "Key Creation",
+    #     client.create_key(
+    #         "4894a298-759a-4b94-a292-f43b2f8deea6",
+    #         "SYMMETRIC_256",
+    #         "test-encrypt-sdk",
+    #         desc="测试本地加密SDK",
+    #         cookies=cookies,
+    #     ),
+    # )
+    #                                         )
+    print(
+        "Key Insert",
+        client.import_key(
+            "4894a298-759a-4b94-a292-f43b2f8deea6",
+            "SYMMETRIC_256",
+            os.urandom(32),
+            "test-py-import",
+            cookies=cookies,
+        ),
+    )
+    # print("Key Set Policy", client.set_key_policy(
+    #     "b3bf3cd7-2a68-427f-865f-8688dc96c7d8",
+    #     "1110fb2f-8dca-4e27-9a37-04a60585fd29",
+    #     json.dumps([
+    #         [
+    #             {
+    #                 "version": "v1",
+    #                 "type": "RAPolicy",
+    #                 "optional": False,
+    #                 "inversion": False,
+    #                 "policyID": "78485ba8-4e67-46a5-a34e-1a9b7930df02"
+    #             }
+    #         ]
+    #     ])
+    # ))

src/bytedance/tks/consts.py

@@ -0,0 +1,122 @@
+import os
+
+
+class Conf:
+    # RELATIVE_CONST_INDEX_PATH = './templates/html/main/index.html'
+    # CONST_INDEX_PATH = './html/main/index.html'
+
+    # cookie key名称，用以保存登陆态的session token
+    SESSION_COOKIE_KEY = "session_token"
+    # 合法时间戳的最大误差，单位秒
+    TIMESTAMP_THRESHOLD = 60
+    # session token的过期时间，单位秒
+    SESSION_TOKEN_EXPIRE = 86400
+    # redis配置
+    REDIS_HOST = "localhost"
+    REDIS_PORT = 6379
+    REDIS_USERNAME = "default"
+    REDIS_PASSWORD = ""
+    # pcc后端的分流配置
+    BACKEND_MAP = {
+        "/api/tks/": ["http://localhost:10000", "appID", "AK", "SK"],
+        "/api/bgs/": ["http://localhost:10001", "UID", "AK", "SK"],
+        "/api/ras/": ["http://localhost:10002", "UID", "AK", "SK"],
+        "/api/rag/": ["http://localhost:10003", "UID", "AK", "SK"],
+    }
+    STREAM_PATH = "/api/rag/collection/query_and_generate"
+
+    UPLOAD_FOLDER = os.path.abspath("./upload/rag/doc")
+    DOWNLOAD_FOLDER = os.path.abspath("./download/attest/report")
+    ALLOWED_EXTENSIONS = {".txt", ".pdf", ".pptx", ".docx"}
+
+    # pcc后端转发配置
+    PCC_UID = "jeddak_team"
+    PCC_PASSWORD = "some_password"
+    PCC_AK = "some_AK"
+    PCC_SK = "some_SK"
+    RING_ID = "id_1234"
+    POLICY_ID = "policy_1234"
+    KEY_TYPE = "SYMMETRIC_256"
+
+    RAG_SERVICE_NAME = "rag-deployment"
+    HANDLER_SERVICE_NAME = "ai-handler"
+    # RAG_SERVICE = {
+    #     "ServiceInfo": {
+    #         "_key": "rag",
+    #         "_val": "rag-deployment"
+    #     }
+    # }
+    # HANDLER_SERVICE = {
+    #     "ServiceInfo": {
+    #         "_key": "handler",
+    #         "_val": "ai-handler"
+    #     }
+    # }
+    # SERVICE_INFO_LIST = {"rag": RAG_SERVICE, "handler": HANDLER_SERVICE}
+    SERVICE_INFO = {"ServiceInfo": {"rag": RAG_SERVICE_NAME, "handler": HANDLER_SERVICE_NAME}}
+
+    PCC_URL = "https://jeddakchain.bytedance.com"
+    PATH_MAP = {
+        "/api_mobile/doc/update": "/api/rag/collection",
+        "/api_mobile/rag/query_and_generate": "/api/rag/collection/query_and_generate",
+        # security使用同一个接口，如果改成不同，ENV处理逻辑要改
+        "/api_mobile/security/status": "/api/ras/v1/attestation_user",
+        "/api_mobile/security/status/download": "/api/ras/v1/attestation_user",
+        # TODO: 确定接口
+        "/api_mobile/key/set": "/api/tks/",
+    }
+
+    # RAG_ADDR = 'http://localhost:10003'
+    # RAG_COLLECTION_CREATE_ENDPOINT = '/api/rag/collection/create'
+    # RAS_SHIM_ADDR = "http://localhost:10002"
+
+    DOC_STATUS_MAP = {"creating": "处理中", "ready": "处理完成", "failed": "处理失败"}
+
+    SALT = "ai-handler-salt"
+
+    def update_from_env(self):
+        for k in self.BACKEND_MAP:
+            env_key = f"ENV_{k.split('/')[-2].upper()}_URL"
+            # ENV_TKS_URL, ...
+            if env_key in os.environ:
+                self.BACKEND_MAP[k][0] = os.getenv(env_key)
+        for k in self.PATH_MAP:
+            env_key = f"HANDLER_{k.split('/')[2].upper()}_URL"
+            if env_key in os.environ:
+                self.PATH_MAP[k] = os.getenv(env_key)
+        self.REDIS_USERNAME = (
+            os.getenv("ENV_REDIS_USERNAME") if "ENV_REDIS_USERNAME" in os.environ else None
+        )
+        self.REDIS_PASSWORD = (
+            os.getenv("ENV_REDIS_PASSWORD") if "ENV_REDIS_PASSWORD" in os.environ else None
+        )
+        if "ENV_REDIS_HOST" in os.environ:
+            self.REDIS_HOST = os.getenv("ENV_REDIS_HOST")
+        if "ENV_REDIS_PORT" in os.environ:
+            self.REDIS_PORT = int(os.getenv("ENV_REDIS_PORT"))
+        if "ENV_SESSION_TOKEN_EXPIRE" in os.environ:
+            self.SESSION_TOKEN_EXPIRE = int(os.getenv("ENV_SESSION_TOKEN_EXPIRE"))
+        if "RAG_SERVICE_NAME" in os.environ:
+            self.RAG_SERVICE_NAME = os.getenv("RAG_SERVICE_NAME")
+        if "HANDLER_SERVICE_NAME" in os.environ:
+            self.HANDLER_SERVICE_NAME = os.getenv("HANDLER_SERVICE_NAME")
+        if "RING_ID" in os.environ:
+            self.RING_ID = os.getenv("RING_ID")
+        if "POLICY_ID" in os.environ:
+            self.POLICY_ID = os.getenv("POLICY_ID")
+        if "KEY_TYPE" in os.environ:
+            self.KEY_TYPE = os.getenv("KEY_TYPE")
+        if "PCC_UID" in os.environ:
+            self.PCC_UID = os.getenv("PCC_UID")
+        if "PCC_PASSWORD" in os.environ:
+            self.PCC_PASSWORD = os.getenv("PCC_PASSWORD")
+        if "PCC_AK" in os.environ:
+            self.PCC_AK = os.getenv("PCC_AK")
+        if "PCC_SK" in os.environ:
+            self.PCC_SK = os.getenv("PCC_SK")
+        os.makedirs(conf.UPLOAD_FOLDER, exist_ok=True)
+        os.makedirs(conf.DOWNLOAD_FOLDER, exist_ok=True)
+        print(f"final backend_map: {self.BACKEND_MAP}")
+
+
+conf = Conf()

src/bytedance/tks/crypto/aes.py

@@ -0,0 +1,47 @@
+__all__ = [
+    "gcm_encrypt",
+    "gcm_decrypt",
+]
+
+import secrets
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+from ..consts import conf
+
+
+def gcm_decrypt(key: bytes, ciphertext: bytes) -> bytes:
+    iv, ciphertext = ciphertext[:12], ciphertext[12:]
+    ciphertext, tag = ciphertext[:-16], ciphertext[-16:]
+    decryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv, tag),
+    ).decryptor()
+
+    return decryptor.update(ciphertext) + decryptor.finalize()
+
+
+def gcm_encrypt(key: bytes, plaintext: bytes) -> bytes:
+    iv = secrets.token_bytes(12)
+    encryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv),
+    ).encryptor()
+
+    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
+    return iv + ciphertext + encryptor.tag
+
+
+def generate_key_from_string(seed_string: str) -> bytes:
+    backend = default_backend()
+    # 将字符串种子转换为字节
+    password = seed_string.encode()
+    salt = conf.SALT.encode()
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=backend
+    )
+    key = kdf.derive(password)
+    return key

src/bytedance/tks/crypto/ec.py

@@ -0,0 +1,39 @@
+from typing import Tuple
+
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+
+__all__ = [
+    "diffie_hellman_key_exchange",
+    "generate_key_pair",
+]
+
+
+def diffie_hellman_key_exchange(local_sk: bytes, peer_pk: bytes) -> bytes:
+    local_sk_key = serialization.load_pem_private_key(local_sk, password=None)
+    peer_pk_key = serialization.load_pem_public_key(peer_pk)
+    assert isinstance(local_sk_key, ec.EllipticCurvePrivateKey)
+    assert isinstance(peer_pk_key, ec.EllipticCurvePublicKey)
+    shared_key = local_sk_key.exchange(ec.ECDH(), peer_pk_key)
+    derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=None).derive(
+        shared_key
+    )
+    return derived_key
+
+
+def generate_key_pair() -> Tuple[bytes, bytes]:
+    private_key = ec.generate_private_key(ec.SECP384R1())
+    public_key = private_key.public_key()
+
+    private_key_pem = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    public_key_pem = public_key.public_bytes(
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
+    )
+
+    return private_key_pem, public_key_pem

src/bytedance/tks/exceptions.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+__all__ = [
+    "TKSError",
+    "RAError",
+]
+
+
+class TKSError(Exception):
+    def __init__(self, message: str, http_status: Optional[int] = None):
+        super().__init__(message)
+        self.http_status = http_status
+
+    def __str__(self):
+        if self.http_status:
+            return f"{self.args[0]} (HTTP Status: {self.http_status})"
+        return self.args[0]
+
+
+class RAError(TKSError):
+    def __init__(self, message: str, http_status: Optional[int] = None):
+        super().__init__(message, http_status)

src/bytedance/volcengine_aicc_sdk/__init__.py

@@ -0,0 +1,14 @@
+__all__ = [
+    "AICCClient",
+    "AICCModelType",
+    "__version__",
+]
+
+__version__ = "0.0.31"
+
+from .aicc_client import AICCModelType
+from .aicc_client import AICCClient
+
+
+def __getattr__(name: str) -> object:
+    raise AttributeError(f"module {__name__} has no attribute {name}")

src/bytedance/volcengine_aicc_sdk/aicc_client.py

@@ -0,0 +1,225 @@
+from bytedance.volcengine_aicc_sdk.seedance_inference_client import SeedanceClient
+from bytedance.volcengine_aicc_sdk.seedream_inference_client import SeedreamClient
+"""
+AICC访问客户端
+"""
+
+__all__ = [
+    "AICCClient",
+    "AICCModelType"
+]
+
+from bytedance.volcengine_aicc_sdk.cu_inference_client import CUClient
+from bytedance.volcengine_aicc_sdk.auc_inference_client import AUCClient
+from bytedance.volcengine_aicc_sdk.ark_inference_client import ARKClient
+from bytedance.volcengine_aicc_sdk.deepseek_inference_client import DeepSeekClient
+from bytedance.volcengine_aicc_sdk.asr_websocket_client import AsrWebSocketClient
+from .asr_inference_client import ASRClient
+from volcenginesdkarkruntime import Ark
+from .common.error import AICCException
+from .common.log import LogConfig, logger
+from bytedance.volcengine_aicc_sdk.secure_http_client import AICCConfig
+
+
+class AICCModelType:
+    DOUBAO = "DOUBAO"  # 机密豆包
+    AICC_CU = "AICC-CU"   # AICC CU版模型
+    OSLM = "OSLM"  # 开源模型
+    ASR = "ASR"
+    ASR_WS = "ASR-WS"
+    SEEDANCE = "SEEDANCE"
+    SEEDREAM = "SEEDREAM"
+    AUC = "AUC"
+
+
+class AICCClient:
+    """
+    AICC机密推理访问客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = "",
+                 api_key: str = "",
+                 endpoint: str = "",
+                 aicc_config: AICCConfig = None,
+                 **kwargs
+                 ):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.aicc_config = aicc_config
+        self.endpoint = endpoint
+        self.kwargs = kwargs
+
+        if not self.base_url or not self.api_key:
+            raise AICCException(f"New AICC client failed, base_url={self.base_url}, api_key={self.api_key}")
+
+        if not self.endpoint:
+            self.endpoint = self.kwargs.get("model_name")
+
+        self.config = dict()
+        self._init_config()
+        self._init_log()
+
+    def get_llm_instance(self, model_type):
+        model_type = model_type.upper()
+        if AICCModelType.SEEDANCE in model_type:
+            return SeedanceClient(base_url=self.base_url,
+                                  api_key=self.api_key,
+                                  endpoint=self.endpoint,
+                                  aicc_config=self.aicc_config,
+                                  **self.kwargs)
+        elif AICCModelType.AICC_CU in model_type:
+            return CUClient(base_url=self.base_url,
+                            aicc_service_name=self.aicc_config.aicc_llm_service_name,
+                            policy_id=self.aicc_config.policy_id,
+                            is_openai=self.kwargs.get("is_openai"),
+                            is_secure=self.kwargs.get("is_secure"),
+                            ak=self.aicc_config.ak,
+                            sk=self.aicc_config.sk,
+                            region=self.aicc_config.region,
+                            **self.kwargs)
+        elif AICCModelType.ASR == model_type:
+            return ASRClient(base_url=self.base_url,
+                             aicc_service_name=self.aicc_config.aicc_llm_service_name,
+                             policy_id=self.aicc_config.policy_id,
+                             is_openai=self.kwargs.get("is_openai"),
+                             is_secure=self.kwargs.get("is_secure"),
+                             ak=self.aicc_config.ak,
+                             sk=self.aicc_config.sk,
+                             region=self.aicc_config.region,
+                             **self.kwargs)
+        elif AICCModelType.OSLM in model_type:
+            return DeepSeekClient(base_url=self.base_url,
+                                  api_key=self.api_key,
+                                  ep=self.endpoint,
+                                  aicc_service_name=self.aicc_config.aicc_llm_service_name,
+                                  policy_id=self.aicc_config.policy_id,
+                                  is_openai=self.kwargs.get("is_openai"),
+                                  is_secure=self.kwargs.get("is_secure"),
+                                  ak=self.aicc_config.ak,
+                                  sk=self.aicc_config.sk,
+                                  region=self.aicc_config.region,
+                                  **self.kwargs)
+        elif AICCModelType.DOUBAO in model_type:
+            return ARKClient(base_url=self.base_url,
+                             api_key=self.api_key,
+                             ep=self.endpoint,
+                             is_openai=self.kwargs.get("is_openai"),
+                             is_secure=self.kwargs.get("is_secure"),
+                             ak=self.aicc_config.ak,
+                             sk=self.aicc_config.sk,
+                             region=self.aicc_config.region,
+                             **self.kwargs)
+        elif AICCModelType.ASR_WS == model_type:
+            return AsrWebSocketClient(base_url=self.base_url,
+                                      api_key=self.api_key,
+                                      endpoint=self.endpoint,
+                                      aicc_config=self.aicc_config,
+                                      **self.kwargs)
+        elif AICCModelType.AUC == model_type:
+            return AUCClient(base_url=self.base_url,
+                             api_key=self.api_key,
+                             endpoint=self.endpoint,
+                             aicc_config=self.aicc_config,
+                             **self.kwargs)
+        elif AICCModelType.SEEDREAM == model_type:
+            return SeedreamClient(base_url=self.base_url,
+                                  api_key=self.api_key,
+                                  endpoint=self.endpoint,
+                                  aicc_config=self.aicc_config,
+                                  **self.kwargs)
+        else:
+            raise AICCException(f"New AICC client failed, model_type={model_type}")
+
+    @staticmethod
+    def get_inference_instance(model_type: str,
+                               base_url: str = None,
+                               api_key: str = None,
+                               ep: str = None,
+                               aicc_service_name: str = None,
+                               policy_id: str = None,
+                               is_openai: bool = False,
+                               is_secure: bool = True,
+                               ak: str = None,
+                               sk: str = None,
+                               region: str = "cn-beijing",
+                               **kwargs):
+        if model_type.upper() == AICCModelType.DOUBAO:
+            return ARKClient(base_url=base_url,
+                             api_key=api_key,
+                             ep=ep,
+                             is_openai=is_openai,
+                             is_secure=is_secure,
+                             ak=ak,
+                             sk=sk,
+                             region=region,
+                             **kwargs)
+        elif model_type.upper() == AICCModelType.AICC_CU:
+            return CUClient(base_url=base_url,
+                            aicc_service_name=aicc_service_name,
+                            policy_id=policy_id,
+                            is_openai=is_openai,
+                            is_secure=is_secure,
+                            ak=ak,
+                            sk=sk,
+                            region=region,
+                            **kwargs)
+        elif model_type.upper() == AICCModelType.OSLM:
+            return DeepSeekClient(base_url=base_url,
+                                  api_key=api_key,
+                                  ep=ep,
+                                  aicc_service_name=aicc_service_name,
+                                  policy_id=policy_id,
+                                  is_openai=is_openai,
+                                  is_secure=is_secure,
+                                  ak=ak,
+                                  sk=sk,
+                                  region=region,
+                                  **kwargs)
+        elif model_type.upper() == AICCModelType.ASR:
+            return ASRClient(base_url=base_url,
+                             aicc_service_name=aicc_service_name,
+                             policy_id=policy_id,
+                             is_openai=is_openai,
+                             is_secure=is_secure,
+                             ak=ak,
+                             sk=sk,
+                             region=region,
+                             **kwargs)
+        else:
+            err_str = f"can not support model type:{model_type}"
+            logger.error(err_str)
+            raise AICCException(err_str)
+
+    @staticmethod
+    def get_volcengine_ark(base_url: str = None, api_key: str = None, **kwargs):
+        return Ark(api_key=api_key, base_url=base_url, **kwargs)
+
+    @staticmethod
+    def hello() -> str:
+        info = "hello volcengine_aicc_sdk"
+        print(info)
+        return info
+
+    def _init_config(self):
+        pass
+
+    def _init_log(self):
+        config = LogConfig()
+
+        if log_dir := self.config.get("dir", "aicc_log"):
+            config.dir = log_dir
+
+        if filename := self.config.get("filename", "aicc.log"):
+            config.filename = filename
+
+        if rotation := self.config.get("rotation", 100):
+            config.rotation = rotation
+
+        if retention := self.config.get("retention", 30):
+            config.retention = retention
+
+        if level := self.config.get("level"):
+            config.level = level
+
+        logger.init_log_config(config)

src/bytedance/volcengine_aicc_sdk/aicc_inference.py

@@ -0,0 +1,64 @@
+"""
+AICC机密推理客户端BASE
+"""
+__all__ = ["AICCInference"]
+
+from abc import ABC, abstractmethod
+from typing import IO, Union
+
+
+class AICCInference(ABC):
+    """
+    AICC机密推理访问客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 is_openai: bool = True,
+                 is_secure: bool = True,
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing"
+                 ):
+        self.base_url = base_url
+        self.is_openai = is_openai
+        self.is_secure = is_secure
+        self.ak = ak
+        self.sk = sk
+        self.region = region
+
+    @staticmethod
+    def hello() -> str:
+        info = "hello aicc inference"
+        print(info)
+        return info
+
+    @abstractmethod
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param output_stream:是否启用流式响应
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def function_call(self, messages: list, tools: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param tools:
+        :param output_stream: 是否开启流式
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def asr_recognize(self, file_name: str) -> Union[str, None]:
+        """
+        ASR转文字
+        :param file_name:
+        :return:
+        """
+        pass

src/bytedance/volcengine_aicc_sdk/ark_inference_client.py

@@ -0,0 +1,343 @@
+"""
+机密豆包客户端
+"""
+__all__ = ["ARKClient"]
+
+import copy
+import json
+import os
+from typing import IO, Union
+
+from volcenginesdkarkruntime import Ark
+from volcenginesdkarkruntime._utils._key_agreement import (
+    aes_gcm_encrypt_base64_string, 
+    aes_gcm_decrypt_base64_string, 
+    key_agreement_client
+)
+from openai import OpenAI
+from volcenginesdkarkruntime.types.chat import ChatCompletion, ChatCompletionChunk
+
+from .aicc_inference import AICCInference
+from .common.error import AICCException
+from .common.log import logger
+
+
+class ARKClient(AICCInference):
+    """
+    机密豆包客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 api_key: str = None,
+                 ep: str = None,
+                 is_openai: bool = False,
+                 is_secure: bool = True,
+                 thinking_type: str = "disabled",
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing",
+                 **kwargs):
+        if not base_url or not api_key or not ep:
+            raise AICCException(f"New ARK client failed, base_url={base_url}, api_key={api_key}, ep={ep}")
+
+        super().__init__(base_url, is_openai, is_secure, ak, sk, region)
+
+        os.environ["VOLC_ARK_ENCRYPTION"] = "AICC"
+
+        self.base_url = base_url
+        self.api_key = api_key
+        self.ep = ep
+        self.is_openai = is_openai
+        self.is_secure = is_secure
+        self.thinking_type = thinking_type
+
+        self.kwargs = kwargs
+
+        self.client = None
+        try:
+            self.client = Ark(api_key=api_key, base_url=base_url, ak=ak, sk=sk, region=region)
+        except Exception as e:
+            logger.error(f"New ARK client failed: {e}")
+            raise AICCException("New ARK client failed")
+
+    def _chat_completion_to_str(self, chat_completion: ChatCompletion) -> str:
+        ret = dict()
+        ret["id"] = chat_completion.id
+        ret["object"] = chat_completion.object
+        ret["created"] = chat_completion.created
+        ret["model"] = chat_completion.model
+        ret["service_tier"] = chat_completion.service_tier
+
+        if chat_completion.usage:
+            ret["usage"] = json.loads(chat_completion.usage.json())
+        else:
+            ret["usage"] = None
+
+        ret["choices"] = list()
+        for choice in chat_completion.choices:
+            ret["choices"].append(json.loads(choice.json()))
+
+        return json.dumps(ret)
+
+    def _chat_completion_chunk_to_str(self, chat_completion_chunk: ChatCompletionChunk) -> str:
+        ret = dict()
+        ret["id"] = chat_completion_chunk.id
+        ret["object"] = chat_completion_chunk.object
+        ret["created"] = chat_completion_chunk.created
+        ret["model"] = chat_completion_chunk.model
+        ret["service_tier"] = chat_completion_chunk.service_tier
+        if chat_completion_chunk.usage:
+            ret["usage"] = json.loads(chat_completion_chunk.usage.json())
+        else:
+            ret["usage"] = None
+        ret["choices"] = list()
+        for choice in chat_completion_chunk.choices:
+            ret["choices"].append(json.loads(choice.json()))
+
+        return json.dumps(ret)
+
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        :param messages:
+        :param output_stream:
+        :return:
+        """
+        if not messages:
+            raise AICCException("ARK completions, messages is NULL")
+
+        try:
+            thinking = {
+                "type": self.thinking_type
+            }
+
+            include_usage = True if self.kwargs.get("include_usage", "True").lower() == "true" and output_stream \
+                else False
+
+            if self.is_openai:
+                openai_messages = copy.deepcopy(messages)
+
+                client = OpenAI(base_url=self.base_url, api_key=self.api_key)
+                ark_client = Ark(base_url=self.base_url, api_key=self.api_key)
+
+                ark_certificate = ark_client._get_endpoint_certificate(self.ep)
+                if isinstance(ark_certificate, tuple):
+                    cipher_client = ark_certificate[0]
+                    ring_id = ark_certificate[1]
+                    key_id = ark_certificate[2]
+                    key, nonce, token = cipher_client.generate_ecies_key_pair()
+                elif isinstance(ark_certificate, key_agreement_client):
+                    ring_id, key_id = ark_certificate.get_cert_ring_key_id()
+                    key, nonce, token = ark_certificate.generate_ecies_key_pair()
+                else:
+                    logger.error(f"get ark_certificate ERROR")
+                    raise AICCException("get ark_certificate ERROR")
+
+                info = {
+                    'Version': 'AICCv0.1',
+                    'RingID': ring_id,
+                    'KeyID': key_id,
+                }
+
+                if self.is_secure:
+                    extra_headers = {
+                        "X-Session-Token": token,
+                        "x-ark-moderation-scene": "aicc-skip",
+                        "X-Encrypt-Info": json.dumps(info),
+                    }
+                    for i in range(len(openai_messages)):
+                        content = openai_messages[i].get("content")
+
+                        if isinstance(content, str):
+                            openai_messages[i]["content"] = aes_gcm_encrypt_base64_string(key, nonce, content)
+                        elif isinstance(content, list):
+                            encrypted_content = []
+                            for item in content:
+                                if isinstance(item, dict) and "type" in item and item["type"] == "text":
+                                    # 加密文本
+                                    encrypted_item = item.copy()
+                                    encrypted_item["text"] = aes_gcm_encrypt_base64_string(key, nonce, item["text"])
+                                    encrypted_content.append(encrypted_item)
+                                elif isinstance(item, dict) and "type" in item and item["type"] == "image_url":
+                                    # 加密图片
+                                    encrypted_item = item.copy()
+                                    encrypted_item["image_url"]["url"] = aes_gcm_encrypt_base64_string(
+                                        key, nonce, item["image_url"]["url"])
+                                    encrypted_content.append(encrypted_item)
+                                elif isinstance(item, dict) and "type" in item and item["type"] == "video_url":
+                                    # 加密视频
+                                    error_info = f"Sorry, 密文豆包不支持openai格式"
+                                    logger.error(error_info)
+                                    raise AICCException(error_info)
+
+                                    # encrypted_item = item.copy()
+                                    # if isinstance(encrypted_item["video_url"], dict):
+                                    #     encrypted_item["video_url"]["url"] = aes_gcm_encrypt_base64_string(
+                                    #         key, nonce, item["video_url"]["url"])
+                                    # else:
+                                    #     encrypted_item["video_url"] = aes_gcm_encrypt_base64_string(
+                                    #         key, nonce, item["video_url"])
+                                    # encrypted_content.append(encrypted_item)
+                                else:
+                                    encrypted_content.append(item)
+                            openai_messages[i]["content"] = encrypted_content
+                        else:
+                            raise AICCException("no support this content type")
+                else:
+                    extra_headers = {
+                    }
+
+                completion = client.chat.completions.create(
+                    model=self.ep,
+                    messages=openai_messages,
+                    stream_options={
+                        "include_usage": include_usage
+                    },
+                    extra_headers=extra_headers,
+                    stream=True if output_stream else False,
+                )
+                if output_stream:
+                    for chunk in completion:
+                        if self.is_secure and chunk.choices and chunk.choices[0].delta.content:
+                            chunk.choices[0].delta.content = aes_gcm_decrypt_base64_string(
+                                key, nonce, chunk.choices[0].delta.content)
+                        output_stream.write(f"{self._chat_completion_chunk_to_str(chunk)}\n")
+                else:
+                    if self.is_secure and completion.choices and completion.choices[0].message.content:
+                        completion.choices[0].message.content = \
+                            aes_gcm_decrypt_base64_string(key, nonce, completion.choices[0].message.content)
+                    cont = self._chat_completion_to_str(completion)
+                    return cont
+            else:
+                resp = self.client.chat.completions.create(
+                    model=self.ep,
+                    messages=messages,
+                    thinking=thinking,
+                    stream=True if output_stream else False,
+                    stream_options={
+                        "include_usage": include_usage
+                    },
+                    extra_headers={"x-is-encrypted": str(self.is_secure).lower(), "x-ark-moderation-scene": "aicc-skip"}
+                )
+
+                if output_stream:
+                    for chunk in resp:
+                        output_stream.write(f"{self._chat_completion_chunk_to_str(chunk)}\n")
+                else:
+                    return self._chat_completion_to_str(resp)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"ARK completions ERROR, e={e}")
+            raise AICCException("ARK completions ERROR")
+
+    def chat_completions_stream(self, messages: list, output_stream: IO) -> Union[str, None]:
+        return self.chat_completions(messages, output_stream)
+
+    def chat_completions_stream_yield(self, messages: list):
+        """
+        :param messages:
+        :return:
+        """
+        if not messages:
+            raise AICCException("ARK completions, messages is NULL")
+
+        try:
+            thinking = {
+                "type": self.thinking_type
+            }
+
+            include_usage = True if self.kwargs.get("include_usage", "True").lower() == "true" else False
+
+            if self.is_openai:
+                openai_messages = copy.deepcopy(messages)
+
+                client = OpenAI(base_url=self.base_url, api_key=self.api_key)
+                ark_client = Ark(base_url=self.base_url, api_key=self.api_key)
+                res_a = ark_client._get_endpoint_certificate(self.ep)
+                cipher_client = res_a[0]
+                ring_id = res_a[1]
+                key_id = res_a[2]
+                key, nonce, token = cipher_client.generate_ecies_key_pair()
+                info = {
+                    'Version': 'AICCv0.1',
+                    'RingID': ring_id,
+                    'KeyID': key_id,
+                }
+
+                if self.is_secure:
+                    extra_headers = {
+                        "X-Session-Token": token,
+                        "x-ark-moderation-scene": "aicc-skip",
+                        "X-Encrypt-Info": json.dumps(info),
+                    }
+                    for i in range(len(openai_messages)):
+                        openai_messages[i]["content"] = aes_gcm_encrypt_base64_string(
+                            key, nonce, openai_messages[i].get("content"))
+                else:
+                    extra_headers = {
+                    }
+
+                completion = client.chat.completions.create(
+                    model=self.ep,
+                    messages=openai_messages,
+                    stream_options={
+                        "include_usage": include_usage
+                    },
+                    extra_headers=extra_headers,
+                    stream=True,
+                )
+
+                for chunk in completion:
+                    if self.is_secure and chunk.choices and chunk.choices[0].delta.content:
+                        chunk.choices[0].delta.content = aes_gcm_decrypt_base64_string(
+                            key, nonce, chunk.choices[0].delta.content)
+                    yield self._chat_completion_chunk_to_str(chunk)
+            else:
+                resp = self.client.chat.completions.create(
+                    model=self.ep,
+                    messages=messages,
+                    thinking=thinking,
+                    stream=True,
+                    stream_options={
+                        "include_usage": include_usage
+                    },
+                    extra_headers={"x-is-encrypted": str(self.is_secure).lower(), "x-ark-moderation-scene": "aicc-skip"}
+                )
+
+                for chunk in resp:
+                    yield self._chat_completion_chunk_to_str(chunk)
+        except Exception as e:
+            logger.error(f"ARK completions ERROR, e={e}")
+            raise AICCException("ARK completions ERROR")
+
+    def function_call(self, messages: list, tools: list, output_stream: IO = None) -> str:
+        """
+        :param messages:
+        :param tools:
+        :param output_stream:
+        :return:
+        """
+        if not messages or not tools:
+            raise AICCException("doubao function_call params cannot be empty")
+
+        try:
+            resp = self.client.chat.completions.create(
+                model=self.ep,
+                messages=messages,
+                tools=tools,
+            )
+            return self._chat_completion_to_str(resp)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC ARK service function_call error: {e}")
+            raise AICCException("AICC ARK service function_call error")
+
+    def function_call_stream(self, messages: list, tools: list, output_stream: IO) -> str:
+        return self.function_call(messages, tools, output_stream)
+
+    def asr_recognize(self, file_name: str) -> str:
+        error_info = "Sorry, this version does not support asr at this time."
+        logger.error(error_info)
+        return error_info

src/bytedance/volcengine_aicc_sdk/asr_inference_client.py

@@ -0,0 +1,158 @@
+"""
+DeepSeek机密推理客户端
+"""
+__all__ = ["ASRClient"]
+
+import json
+import uuid
+import requests
+import bytedance.jeddak_secure_channel as jsc
+from typing import IO, Union
+from .aicc_inference import AICCInference
+from .common.constant import Constants
+from .common.error import AICCException
+from .common.log import logger
+from .common.utils import file_to_base64
+
+
+class ASRClient(AICCInference):
+    """
+    ASR客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 aicc_service_name: str = None,
+                 policy_id: str = None,
+                 is_openai: bool = True,
+                 is_secure: bool = True,
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing",
+                 **kwargs):
+        appid = kwargs.get("appid")
+        token = kwargs.get("token")
+        service_id = kwargs.get("service_id")
+
+        if not base_url or not appid or not token or not service_id or not \
+                aicc_service_name or not policy_id or not ak or not sk:
+            raise AICCException("New deepseek client failed, params error")
+
+        super().__init__(base_url, is_openai, is_secure, ak, sk, region)
+
+        self.appid = appid
+        self.token = token
+        self.service_id = service_id
+
+        self.aicc_service_name = aicc_service_name
+        self.policy_id = policy_id
+
+        self.kwargs = kwargs
+
+        self.jsc_client = None
+        self.max_retries = 3
+
+        if is_secure:
+            jsc_config = dict()
+            jsc_config["ra_service_name"] = aicc_service_name
+            jsc_config["ra_policy_id"] = policy_id
+            jsc_config["bytedance_top_info"] = dict()
+            jsc_config["bytedance_top_info"]["ak"] = ak
+            jsc_config["bytedance_top_info"]["sk"] = sk
+            jsc_config["bytedance_top_info"]["region"] = region
+            jsc_config["bytedance_top_info"]["url"] = self.kwargs.get("top_url", "open.volcengineapi.com")
+            jsc_config["bytedance_top_info"]["service"] = self.kwargs.get("top_service", "pcc")
+            jsc_config["log_config"] = self.kwargs.get("log_config", Constants.LOG_CONFIG)
+
+            try:
+                secure_channel_config = jsc.ClientConfig.from_dict(jsc_config)
+                self.jsc_client = jsc.Client(secure_channel_config)
+            except Exception as e:
+                logger.error(f"Error init jsc: {e}")
+                raise AICCException("New deepseek client failed")
+
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param output_stream:
+        :return:
+        """
+        pass
+
+    def function_call(self, messages: list, tools: list, output_stream: IO = None) -> str:
+        """
+        :param messages:
+        :param tools:
+        :param output_stream:
+        :return:
+        """
+        pass
+
+    def asr_recognize(self, file_name: str) -> Union[str, None]:
+        if not file_name:
+            raise AICCException("ASR asr_recognize file_name cannot be empty")
+
+        appid = self.appid
+        token = self.token
+        service_id = self.service_id
+        model_name = "bigmodel"
+
+        if self.base_url.endswith("flash"):
+            recognize_url = self.base_url
+        else:
+            if self.is_secure:
+                recognize_url = f"{self.base_url}/api/v2/auc/{model_name}/recognize/flash"
+            else:
+                recognize_url = f"{self.base_url}/api/v1/auc/{model_name}/recognize/flash"
+
+        try:
+            audio_data = None
+            enc_base64_data = file_to_base64(file_name)  # 转换文件为 Base64
+
+            # 将转换后的base64数据写入文件
+            enc_key = None
+            if self.is_secure:
+                enc_base64_data, enc_key = self.jsc_client.encrypt_with_response(enc_base64_data)
+
+            audio_data = {
+                "data": enc_base64_data,
+            }
+
+            headers = {
+                "X-Api-App-Id": appid,
+                "X-Api-App-Key": token,
+                "X-Api-Aicc-Endpoint": service_id,
+                "X-Api-Request-Id": str(uuid.uuid4()),
+            }
+
+            if not audio_data:
+                raise AICCException("必须提供ASR文件内容")
+
+            request = {
+                "audio": audio_data,
+                "request": {
+                    "model_name": model_name,
+                },
+            }
+
+            response = requests.post(recognize_url, json=request, headers=headers)
+            if response.status_code == 200:  # task finished
+                result = response.json()
+                content = result["result"]
+                if self.is_secure and enc_key:
+                    content = enc_key.decrypt(content)
+                    # 检查 content 类型，如果是 bytes 则解码，否则直接使用
+                    if isinstance(content, bytes):
+                        content = json.loads(content.decode("utf-8"))
+                    else:
+                        content = json.loads(content)
+                result["result"] = content
+                return json.dumps(result)
+            else:
+                raise AICCException("ASR处理失败")
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC ASR service asr_recognize error: {e}")
+            raise AICCException("AICC ASR service asr_recognize error")

src/bytedance/volcengine_aicc_sdk/asr_websocket_client.py

@@ -0,0 +1,647 @@
+"""
+流式语音识别大模型机密推理客户端
+"""
+__all__ = ["AsrWebSocketClient"]
+
+from ..jeddak_secure_channel.crypto import AesKey
+
+"""
+SAUC WebSocket demo client.
+
+Usage:
+1. Minimal run:
+   python3 tests/sauc_websocket_demo.py --file tests/test-sound.mp3 --stop-on-first-result
+
+2. Override any request fields with inline JSON:
+   python3 tests/sauc_websocket_demo.py --file tests/test-sound.mp3 \
+     --request-json '{"request":{"enable_nonstream":true,"show_utterances":true}}'
+
+3. Override any request fields with a JSON file:
+   python3 tests/sauc_websocket_demo.py --file tests/test-sound.mp3 \
+     --request-json-file /path/to/sauc_request.json
+
+Notes:
+- CLI keeps only commonly needed options.
+- Any SAUC full client request fields can still be passed through JSON overrides.
+"""
+
+import io
+import json
+import gzip
+import subprocess
+import struct
+import threading
+import time
+import uuid
+import wave
+from dataclasses import dataclass
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+try:
+    from websocket import WebSocket, WebSocketConnectionClosedException, WebSocketTimeoutException, create_connection
+except ImportError:
+    try:
+        from websocket._core import WebSocket
+        from websocket._exceptions import WebSocketConnectionClosedException, WebSocketTimeoutException
+        from websocket._app import create_connection
+    except ImportError:
+        WebSocket = object
+        WebSocketConnectionClosedException = ConnectionError
+        WebSocketTimeoutException = TimeoutError
+        create_connection = None
+import bytedance.jeddak_secure_channel as jsc
+from .common.log import logger
+from .secure_http_client import AICCConfig
+
+# 常量定义
+DEFAULT_SAMPLE_RATE = 16000
+DEFAULT_MODEL_NAME = "bigmodel"
+INTER_PACKET_DRAIN_TIMEOUT_SECONDS = 0.02
+DEFAULT_FULL_CLIENT_REQUEST = {
+    "audio": {
+        "format": "wav",
+        "codec": "raw",
+        "rate": DEFAULT_SAMPLE_RATE,
+        "bits": 16,
+        "channel": 1,
+    },
+    "request": {
+        "model_name": DEFAULT_MODEL_NAME,
+    },
+}
+
+
+class ProtocolVersion:
+    V1 = 0b0001
+
+
+class MessageType:
+    CLIENT_FULL_REQUEST = 0b0001
+    CLIENT_AUDIO_ONLY_REQUEST = 0b0010
+    SERVER_FULL_RESPONSE = 0b1001
+    SERVER_ERROR_RESPONSE = 0b1111
+
+
+class MessageTypeSpecificFlags:
+    NO_SEQUENCE = 0b0000
+    POS_SEQUENCE = 0b0001
+    NEG_SEQUENCE = 0b0010
+    NEG_WITH_SEQUENCE = 0b0011
+
+
+class SerializationType:
+    NO_SERIALIZATION = 0b0000
+    JSON = 0b0001
+
+
+class CompressionType:
+    GZIP = 0b0001
+
+
+class AsrRequestHeader:
+    def __init__(self):
+        self.message_type = MessageType.CLIENT_FULL_REQUEST
+        self.message_type_specific_flags = MessageTypeSpecificFlags.POS_SEQUENCE
+        self.serialization_type = SerializationType.JSON
+        self.compression_type = CompressionType.GZIP
+        self.reserved_data = bytes([0x00])
+
+    def with_message_type(self, message_type: int) -> "AsrRequestHeader":
+        self.message_type = message_type
+        return self
+
+    def with_message_type_specific_flags(self, flags: int) -> "AsrRequestHeader":
+        self.message_type_specific_flags = flags
+        return self
+
+    def with_serialization_type(self, serialization_type: int) -> "AsrRequestHeader":
+        self.serialization_type = serialization_type
+        return self
+
+    def with_compression_type(self, compression_type: int) -> "AsrRequestHeader":
+        self.compression_type = compression_type
+        return self
+
+    def with_reserved_data(self, reserved_data: bytes) -> "AsrRequestHeader":
+        self.reserved_data = reserved_data
+        return self
+
+    def to_bytes(self) -> bytes:
+        header = bytearray()
+        header.append((ProtocolVersion.V1 << 4) | 1)
+        header.append((self.message_type << 4) | self.message_type_specific_flags)
+        header.append((self.serialization_type << 4) | self.compression_type)
+        header.extend(self.reserved_data)
+        return bytes(header)
+
+    @staticmethod
+    def default_header() -> "AsrRequestHeader":
+        return AsrRequestHeader()
+
+
+@dataclass
+class AsrResponse:
+    code: int = 0
+    event: int = 0
+    is_last_package: bool = False
+    payload_sequence: int = 0
+    payload_size: int = 0
+    payload_msg: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "code": self.code,
+            "event": self.event,
+            "is_last_package": self.is_last_package,
+            "payload_sequence": self.payload_sequence,
+            "payload_size": self.payload_size,
+            "payload_msg": self.payload_msg,
+        }
+
+    def is_terminal(self) -> bool:
+        return self.is_last_package or self.code != 0
+
+
+class _SessionContext:
+    """Holds per-execution state to support concurrent access."""
+
+    def __init__(self, client: "AsrWebSocketClient"):
+        self.client = client
+        self.conn: Optional[WebSocket] = None
+        self.seq: int = 1
+        self.encrypt_key = None
+        self.aes_key = AesKey.generate() if client.jsc_client else None
+
+    def connect(self) -> None:
+        headers = self.client._build_auth_headers(
+            self.client.api_key, self.client.resource_id, self.client._build_connection_headers()
+        )
+        self.conn = create_connection(
+            self.client.base_url,
+            header=[f"{key}: {value}" for key, value in headers.items()],
+            timeout=self.client.connect_timeout,
+        )
+        self.conn.settimeout(self.client.receive_timeout)
+        logger.info("Connected to %s", self.client.base_url)
+
+    def close(self) -> None:
+        if self.conn is not None:
+            self.conn.close()
+            self.conn = None
+
+    def encode_outbound_frame(self, payload: bytes) -> bytes:
+        if self.client.jsc_client:
+            msg, encrypt_key = self.client.jsc_client.encrypt_with_response(payload, self.aes_key)
+            if encrypt_key:
+                self.encrypt_key = encrypt_key
+            if isinstance(msg, str):
+                msg = msg.encode("utf-8")
+            return msg
+        else:
+            return payload
+
+    def decode_inbound_frame(self, payload: bytes) -> bytes:
+        if self.encrypt_key:
+            msg = self.encrypt_key.decrypt(payload)
+            if isinstance(msg, str):
+                msg = msg.encode("utf-8")
+            return msg
+        else:
+            return payload
+
+    def send_binary(self, payload: bytes) -> None:
+        if self.conn is None:
+            raise RuntimeError("WebSocket connection is not established")
+        self.conn.send_binary(self.encode_outbound_frame(payload))
+
+    def recv_one(self, timeout: float) -> Optional[AsrResponse]:
+        if self.conn is None:
+            raise RuntimeError("WebSocket connection is not established")
+
+        self.conn.settimeout(timeout)
+        try:
+            message = self.conn.recv()
+        except WebSocketTimeoutException:
+            return None
+        except WebSocketConnectionClosedException as exc:
+            raise RuntimeError("WebSocket connection closed unexpectedly") from exc
+        finally:
+            self.conn.settimeout(self.client.receive_timeout)
+
+        if isinstance(message, str):
+            logger.info("Received text message: %s", message)
+            return None
+
+        return self.client._parse_response(self.decode_inbound_frame(message))
+
+    def drain_messages(self, timeout: float) -> Iterator[AsrResponse]:
+        while True:
+            response = self.recv_one(timeout)
+            if response is None:
+                return
+            yield response
+            timeout = 0.1
+
+    def send_full_client_request(self) -> Iterator[AsrResponse]:
+        self.send_binary(
+            self.client._build_full_client_request(
+                self.seq,
+                model_name=self.client.model_name,
+                enable_nonstream=self.client.enable_nonstream,
+                enable_itn=self.client.enable_itn,
+                enable_punc=self.client.enable_punc,
+                enable_ddc=self.client.enable_ddc,
+                show_utterances=self.client.show_utterances,
+                result_type=self.client.result_type,
+                enable_accelerate_text=self.client.enable_accelerate_text,
+                vad_segment_duration=self.client.vad_segment_duration,
+                full_client_request_overrides=self.client.full_client_request_overrides,
+            )
+        )
+        logger.info(f"Sent full client request with seq: {self.seq}")
+        self.seq += 1
+        yield from self.drain_messages(self.client.connect_timeout)
+
+    def send_audio_stream(self, content: bytes, segment_size: int) -> Iterator[AsrResponse]:
+        audio_segments = self.client._split_audio(content, segment_size)
+        total_segments = len(audio_segments)
+
+        for index, segment in enumerate(audio_segments):
+            is_last = index == total_segments - 1
+            request = self.client._build_audio_only_request(self.seq, segment, is_last=is_last)
+            self.send_binary(request)
+            logger.debug(f"Sent audio segment with seq: {self.seq} (last: {is_last})")
+            if not is_last:
+                self.seq += 1
+
+            if self.client.segment_duration > 0:
+                time.sleep(self.client.segment_duration / 1000)
+
+            yield from self.drain_messages(INTER_PACKET_DRAIN_TIMEOUT_SECONDS)
+
+    def receive_final_messages(self) -> Iterator[AsrResponse]:
+        deadline = time.monotonic() + self.client.final_wait_seconds
+        while time.monotonic() < deadline:
+            response = self.recv_one(min(self.client.receive_timeout, deadline - time.monotonic()))
+            if response is None:
+                continue
+            yield response
+            if response.is_terminal():
+                return
+
+
+class AsrWebSocketClient:
+    _jsc_client_cache: Dict[str, Any] = {}
+    _jsc_client_lock = threading.Lock()
+
+    def __init__(
+            self,
+            base_url: str = "",
+            api_key: str = "",
+            endpoint: str = "",
+            aicc_config: AICCConfig = None,
+            **kwargs
+    ):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.endpoint = endpoint
+        self.resource_id = self.endpoint
+        self.aicc_config = aicc_config
+        self.kwargs = kwargs
+
+        self.seq = self.kwargs.get("seq", 1)
+        self.segment_duration = self.kwargs.get("segment_duration", 200)
+
+        self.model_name = self.kwargs.get("model_name", None)
+        self.enable_nonstream = self.kwargs.get("enable_nonstream", None)
+        self.enable_itn = self.kwargs.get("enable_itn", None)
+        self.enable_punc = self.kwargs.get("enable_punc", None)
+        self.enable_ddc = self.kwargs.get("enable_ddc", None)
+        self.show_utterances = self.kwargs.get("show_utterances", None)
+        self.result_type = self.kwargs.get("result_type", None)
+        self.enable_accelerate_text = self.kwargs.get("enable_accelerate_text", None)
+        self.vad_segment_duration = self.kwargs.get("vad_segment_duration", None)
+        self.full_client_request_overrides = self.kwargs.get("full_client_request_overrides", {})
+        self.extra_headers = self.kwargs.get("extra_headers", {})
+        self.connect_timeout = self.kwargs.get("connect_timeout", 10.0)
+        self.receive_timeout = self.kwargs.get("receive_timeout", 1.0)
+        self.final_wait_seconds = self.kwargs.get("final_wait_seconds", 15.0)
+        self.is_secure = self.kwargs.get("is_secure", True)
+
+        # 将http地址转成wss地址
+        if self.base_url.startswith("https://"):
+            self.base_url = self.base_url.replace("https", "wss")
+        elif self.base_url.startswith("http://"):
+            self.base_url = self.base_url.replace("http", "wss")
+
+        asr_suffix = self.kwargs.get("asr_suffix", "/sauc/bigmodel_async")
+        if not self.base_url.endswith(asr_suffix):
+            self.base_url = self.base_url + asr_suffix
+
+        self.jsc_client = None
+        if not self.aicc_config:
+            if self.is_secure:
+                self.aicc_config = AICCConfig(ra_url=base_url, auth_token=api_key)
+        if self.aicc_config:
+            self.jsc_client = self._get_or_create_jsc_client(aicc_config=self.aicc_config)
+
+        self.AICC_SDK_VERSION = "0.1.0"
+
+    @classmethod
+    def _get_or_create_jsc_client(cls, config_dict: Dict[str, Any] = None, aicc_config: AICCConfig = None) -> Any:
+        if aicc_config is not None:
+            config_dict = aicc_config.get_jsc_config()
+        cache_key = json.dumps(config_dict, sort_keys=True)
+        if cache_key in cls._jsc_client_cache:
+            return cls._jsc_client_cache[cache_key]
+        with cls._jsc_client_lock:
+            if cache_key in cls._jsc_client_cache:
+                return cls._jsc_client_cache[cache_key]
+            secure_channel_config = jsc.ClientConfig.from_dict(config_dict)
+            client = jsc.Client(secure_channel_config)
+            cls._jsc_client_cache[cache_key] = client
+            return client
+
+    def _build_auth_headers(
+            self,
+            api_key: str,
+            resource_id: str,
+            extra_headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, str]:
+        headers = {
+            "X-Api-Resource-Id": resource_id,
+            "X-Api-Request-Id": str(uuid.uuid4()),
+            "X-Api-Connect-Id": str(uuid.uuid4()),
+            "X-Api-Sequence": "-1",
+            "X-Api-Key": api_key,
+            "Authorization": f"Bearer {api_key}",
+        }
+        if extra_headers:
+            headers.update(extra_headers)
+        return headers
+
+    def _convert_audio_to_wav_bytes(self, audio_path: str, sample_rate: int = DEFAULT_SAMPLE_RATE) -> bytes:
+        cmd = [
+            "ffmpeg",
+            "-v",
+            "quiet",
+            "-y",
+            "-i",
+            audio_path,
+            "-acodec",
+            "pcm_s16le",
+            "-ac",
+            "1",
+            "-ar",
+            str(sample_rate),
+            "-f",
+            "wav",
+            "-",
+        ]
+        try:
+            result = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        except subprocess.CalledProcessError as exc:
+            message = exc.stderr.decode("utf-8", errors="replace")
+            raise RuntimeError(f"Audio conversion failed: {message}") from exc
+
+        try:
+            with open("converted.wav", "wb") as file_handle:
+                file_handle.write(result.stdout)
+        except OSError as exc:
+            logger.warning("Failed to save converted.wav: %s", exc)
+
+        return result.stdout
+
+    def _is_wav_file(self, data: bytes) -> bool:
+        if len(data) < 12:
+            return False
+        return data[:4] == b"RIFF" and data[8:12] == b"WAVE"
+
+    def _read_audio_data(self, file_path: str) -> bytes:
+        with open(file_path, "rb") as file_handle:
+            content = file_handle.read()
+
+        if self._is_wav_file(content):
+            return content
+
+        logger.info("Converting audio to WAV format...")
+        return self._convert_audio_to_wav_bytes(file_path, DEFAULT_SAMPLE_RATE)
+
+    def _read_wav_info(self, data: bytes) -> Tuple[int, int, int, int]:
+        try:
+            with wave.open(io.BytesIO(data), "rb") as wav_file:
+                channel_num = wav_file.getnchannels()
+                samp_width = wav_file.getsampwidth()
+                frame_rate = wav_file.getframerate()
+                frame_count = wav_file.getnframes()
+        except wave.Error as exc:
+            raise ValueError(f"Invalid WAV file: {exc}") from exc
+
+        return channel_num, samp_width, frame_rate, frame_count
+
+    def _get_segment_size(self, content: bytes) -> int:
+        channel_num, samp_width, frame_rate, _ = self._read_wav_info(content)
+        size_per_second = channel_num * samp_width * frame_rate
+        return size_per_second * self.segment_duration // 1000
+
+    def _build_connection_headers(self) -> Dict[str, str]:
+        headers = dict(self.extra_headers)
+        if self.aicc_config:
+            headers["X-AICC-Encryption-Enable"] = "true"
+            headers["X-AICC-Encryption-SDK"] = "aicc"
+            headers["X-AICC-Encryption-Version"] = self.AICC_SDK_VERSION
+        return headers
+
+    def _gzip_compress(self, data: bytes) -> bytes:
+        return gzip.compress(data)
+
+    def _gzip_decompress(self, data: bytes) -> bytes:
+        return gzip.decompress(data)
+
+    def _parse_response(self, message: bytes) -> AsrResponse:
+        if not message:
+            raise ValueError("Cannot parse an empty websocket frame")
+
+        response = AsrResponse()
+        header_size = message[0] & 0x0F
+        message_type = message[1] >> 4
+        message_type_specific_flags = message[1] & 0x0F
+        serialization_method = message[2] >> 4
+        message_compression = message[2] & 0x0F
+        payload = message[header_size * 4:]
+
+        if message_type_specific_flags & 0x01:
+            response.payload_sequence = struct.unpack(">i", payload[:4])[0]
+            payload = payload[4:]
+        if message_type_specific_flags & 0x02:
+            response.is_last_package = True
+        if message_type_specific_flags & 0x04:
+            response.event = struct.unpack(">i", payload[:4])[0]
+            payload = payload[4:]
+
+        if message_type == MessageType.SERVER_FULL_RESPONSE:
+            response.payload_size = struct.unpack(">I", payload[:4])[0]
+            payload = payload[4:]
+        elif message_type == MessageType.SERVER_ERROR_RESPONSE:
+            response.code = struct.unpack(">i", payload[:4])[0]
+            response.payload_size = struct.unpack(">I", payload[4:8])[0]
+            payload = payload[8:]
+
+        if not payload:
+            return response
+
+        if message_compression == CompressionType.GZIP:
+            try:
+                payload = self._gzip_decompress(payload)
+            except Exception as exc:
+                logger.error("Failed to decompress payload: %s", exc)
+                return response
+
+        try:
+            if serialization_method == SerializationType.JSON:
+                response.payload_msg = json.loads(payload.decode("utf-8"))
+        except Exception as exc:
+            logger.error("Failed to parse payload: %s", exc)
+
+        return response
+
+    def _deep_merge_dict(self, base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
+        merged = dict(base)
+        for key, value in override.items():
+            if isinstance(value, dict) and isinstance(merged.get(key), dict):
+                merged[key] = self._deep_merge_dict(merged[key], value)
+            else:
+                merged[key] = value
+        return merged
+
+    def _build_full_request_payload(
+            self,
+            *,
+            model_name: Optional[str] = None,
+            enable_nonstream: Optional[bool] = None,
+            enable_itn: Optional[bool] = None,
+            enable_punc: Optional[bool] = None,
+            enable_ddc: Optional[bool] = None,
+            show_utterances: Optional[bool] = None,
+            result_type: Optional[str] = None,
+            enable_accelerate_text: Optional[bool] = None,
+            vad_segment_duration: Optional[int] = None,
+            full_client_request_overrides: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        payload = self._deep_merge_dict(DEFAULT_FULL_CLIENT_REQUEST, full_client_request_overrides or {})
+        request_payload = dict(payload.get("request", {}))
+
+        if model_name is not None:
+            request_payload["model_name"] = model_name
+        if enable_nonstream is not None:
+            request_payload["enable_nonstream"] = enable_nonstream
+        if enable_itn is not None:
+            request_payload["enable_itn"] = enable_itn
+        if enable_punc is not None:
+            request_payload["enable_punc"] = enable_punc
+        if enable_ddc is not None:
+            request_payload["enable_ddc"] = enable_ddc
+        if show_utterances is not None:
+            request_payload["show_utterances"] = show_utterances
+        if result_type is not None:
+            request_payload["result_type"] = result_type
+        if enable_accelerate_text is not None:
+            request_payload["enable_accelerate_text"] = enable_accelerate_text
+        if vad_segment_duration is not None:
+            request_payload["vad_segment_duration"] = vad_segment_duration
+
+        payload["request"] = request_payload
+        return payload
+
+    def _build_full_client_request(
+            self,
+            seq: int,
+            *,
+            model_name: Optional[str] = None,
+            enable_nonstream: Optional[bool] = None,
+            enable_itn: Optional[bool] = None,
+            enable_punc: Optional[bool] = None,
+            enable_ddc: Optional[bool] = None,
+            show_utterances: Optional[bool] = None,
+            result_type: Optional[str] = None,
+            enable_accelerate_text: Optional[bool] = None,
+            vad_segment_duration: Optional[int] = None,
+            full_client_request_overrides: Optional[Dict[str, Any]] = None,
+    ) -> bytes:
+        header = AsrRequestHeader.default_header().with_message_type_specific_flags(MessageTypeSpecificFlags.POS_SEQUENCE)
+        payload = self._build_full_request_payload(
+            model_name=model_name,
+            enable_nonstream=enable_nonstream,
+            enable_itn=enable_itn,
+            enable_punc=enable_punc,
+            enable_ddc=enable_ddc,
+            show_utterances=show_utterances,
+            result_type=result_type,
+            enable_accelerate_text=enable_accelerate_text,
+            vad_segment_duration=vad_segment_duration,
+            full_client_request_overrides=full_client_request_overrides,
+        )
+        payload_bytes = json.dumps(payload).encode("utf-8")
+        compressed_payload = self._gzip_compress(payload_bytes)
+
+        request = bytearray()
+        request.extend(header.to_bytes())
+        request.extend(struct.pack(">i", seq))
+        request.extend(struct.pack(">I", len(compressed_payload)))
+        request.extend(compressed_payload)
+        return bytes(request)
+
+    def _split_audio(self, data: bytes, segment_size: int) -> List[bytes]:
+        if segment_size <= 0:
+            return []
+        return [data[index:index + segment_size] for index in range(0, len(data), segment_size)]
+
+    def _build_audio_only_request(self, seq: int, segment: bytes, is_last: bool = False) -> bytes:
+        header = AsrRequestHeader.default_header()
+        if is_last:
+            header.with_message_type_specific_flags(MessageTypeSpecificFlags.NEG_WITH_SEQUENCE)
+            seq = -seq
+        else:
+            header.with_message_type_specific_flags(MessageTypeSpecificFlags.POS_SEQUENCE)
+        header.with_message_type(MessageType.CLIENT_AUDIO_ONLY_REQUEST)
+
+        request = bytearray()
+        request.extend(header.to_bytes())
+        request.extend(struct.pack(">i", seq))
+
+        compressed_segment = self._gzip_compress(segment)
+        request.extend(struct.pack(">I", len(compressed_segment)))
+        request.extend(compressed_segment)
+        return bytes(request)
+
+    def execute(self, file_path: str) -> Iterator[AsrResponse]:
+        if not file_path:
+            raise ValueError("File path is empty")
+        if not self.base_url or not (self.base_url.startswith("wss://") or self.base_url.startswith("ws://")):
+            raise ValueError("Base URL is empty or invalid (must start with wss:// or ws://)")
+        if not self.endpoint:
+            raise ValueError("Endpoint is empty")
+        if not self.api_key:
+            raise ValueError("API key is empty")
+
+        session = _SessionContext(self)
+
+        content = self._read_audio_data(file_path)
+        segment_size = self._get_segment_size(content)
+
+        try:
+            session.connect()
+
+            for response in session.send_full_client_request():
+                yield response
+                if response.is_terminal():
+                    return
+
+            for response in session.send_audio_stream(content, segment_size):
+                yield response
+                if response.is_terminal():
+                    return
+
+            yield from session.receive_final_messages()
+        finally:
+            session.close()

src/bytedance/volcengine_aicc_sdk/auc_inference_client.py

@@ -0,0 +1,305 @@
+import uuid
+"""
+seedance机密推理客户端
+"""
+__all__ = ["AUCClient"]
+
+import os
+import base64
+import requests
+import traceback
+from urllib.parse import unquote
+from cryptography.hazmat.primitives import serialization as crypto_serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.backends import default_backend as crypto_default_backend
+from bytedance.jeddak_secure_channel.crypto import PrivateKey
+from .secure_http_client import SecureHttpClient
+from .common.tos_utils import TOSUtils
+from .common.error import AICCException
+from .common.log import logger
+from .secure_http_client import AICCConfig
+from .common.aes import decrypt_file
+
+
+class AUCClient:
+    """
+    AUC机密推理客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = "",
+                 api_key: str = "",
+                 endpoint: str = "",
+                 aicc_config: AICCConfig = None,
+                 **kwargs):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.endpoint = endpoint
+        self.aicc_config = aicc_config
+        self.tos_utils = TOSUtils()
+        self.kwargs = kwargs
+        self.is_secure = self.kwargs.get("is_secure", True)
+
+        if aicc_config:
+            self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+        else:
+            if self.is_secure:
+                aicc_config = AICCConfig(ra_url=base_url, auth_token=api_key)
+                self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+            else:
+                self.secure_http_client = SecureHttpClient()
+
+        # 是否开启AUC文件加密, 默认加密
+        self.is_enable_auc_encrypt = str(self.kwargs.get("is_enable_auc_encrypt", True)).lower() == "true"
+
+        # AUC文件的存储位置, 目前仅支持火山TOS
+        self.auc_file_storage_location = self.kwargs.get("auc_file_storage_location", "TOS").upper()
+
+        self.auc_file_encrypt_pub_key_content = None
+        self.auc_file_encrypt_priv_key_content = None
+
+        self.auc_file_encrypt_priv_key: PrivateKey = None
+
+        self.timeout = self.kwargs.get("timeout", 120.0)
+
+        # API Endpoints
+        self.CREATE_TASK_URL = "/auc/lark/submit"
+        self.QUERY_TASK_URL = "/auc/lark/query"
+
+    def _build_common_headers(self, request_id: str) -> dict[str, str]:
+        if not request_id:
+            request_id = str(uuid.uuid4())
+
+        return {
+            "Content-Type": "application/json",
+            "X-Api-App-Key": self.endpoint,
+            "Authorization": f"Bearer {self.api_key}",
+            "X-Request-Id": request_id,
+        }
+
+    def _get_base_url(self):
+        """
+        Get base URL with trailing slash removed
+        """
+        return self.base_url.rstrip("/")
+
+    def generate_key_pair(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Generate RSA 4096 key pair and save to files
+
+        Args:
+            pub_key_path: Path to save public key
+            priv_key_path: Path to save private key
+
+        Returns:
+            Tuple of (public_key_path, private_key_path)
+        """
+        try:
+            # Generate RSA key pair
+            key = rsa.generate_private_key(
+                backend=crypto_default_backend(),
+                public_exponent=65537,
+                key_size=4096
+            )
+
+            # Save private key
+            private_key = key.private_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PrivateFormat.PKCS8,
+                crypto_serialization.NoEncryption()
+            )
+            with open(priv_key_path, 'wb') as f:
+                f.write(private_key)
+
+            # Save public key
+            public_key = key.public_key().public_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PublicFormat.SubjectPublicKeyInfo
+            )
+            with open(pub_key_path, 'wb') as f:
+                f.write(public_key)
+        except Exception as e:
+            logger.error(f"[Error] Failed to generate key pair: {e}")
+            raise AICCException("Failed to generate RSA key pair") from e
+
+    def set_file_encrypt_key(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Set auc file encrypt key
+        Args:
+            pub_key_path: Public key for encrypt
+            priv_key_path: Private key for decrypt
+        """
+        if os.path.exists(pub_key_path):
+            with open(pub_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.auc_file_encrypt_pub_key_content = content
+
+        if os.path.exists(priv_key_path):
+            with open(priv_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.auc_file_encrypt_priv_key_content = content
+                    self.auc_file_encrypt_priv_key = \
+                        PrivateKey.from_private_pem(self.auc_file_encrypt_priv_key_content)
+
+    def create_task(self, data: dict, custom_header: dict = None, request_id: str = "") -> str:
+        """
+        :param data:
+        :param custom_header:
+        :return:
+        """
+        if not data:
+            raise AICCException("data must be provided")
+
+        if "model" not in data:
+            data["model"] = self.endpoint
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.CREATE_TASK_URL}"
+
+        headers = self._build_common_headers(request_id)
+        if custom_header:
+            headers.update(custom_header)
+
+        # Add encryption headers if provided
+        if self.is_enable_auc_encrypt:
+            if not self.auc_file_encrypt_pub_key_content or not self.auc_file_encrypt_priv_key_content:
+                raise AICCException("pub_key and priv_key must be provided when Is-Enable-Encrypt is 'true'")
+
+            headers["Enable-TOS-Content-Result-Encryption"] = str(self.is_enable_auc_encrypt).lower()
+            headers["X-Encryption-Algorithm"] = "RSA_OAEP_4096_AES_256"
+            if self.auc_file_encrypt_pub_key_content:
+                # Base64 encode the PEM public key for safe HTTP header transmission
+                # This avoids issues with newlines and special characters in headers
+                pk_base64 = base64.b64encode(self.auc_file_encrypt_pub_key_content.encode()).decode()
+                headers["PK"] = pk_base64
+
+        # logger.info(f"[Request] url: {url}, headers: {headers}, data: {data}")
+
+        response = self.secure_http_client.post(url, headers=headers, json=data, timeout=self.timeout)
+        if response.status_code == 200:
+            result = response.json()
+            task_id = result.get("Data", {}).get("TaskID", "")
+            return task_id
+        else:
+            logger.error(f"[Error] url: {url}, response: {response.text}")
+            return ""
+
+    def query_task(self, task_id: str, request_id: str = "") -> dict:
+        """
+        Query auc task
+        API: POST /auc/lark/query
+        """
+
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.QUERY_TASK_URL}"
+
+        headers = self._build_common_headers(request_id)
+        payload = {"TaskID": task_id}
+
+        response = self.secure_http_client.post(url, headers=headers, json=payload, timeout=self.timeout)
+        if response.status_code == 200:
+            result = response.json()
+            return result
+        else:
+            try:
+                result = response.json()
+                return result
+            except Exception as e:
+                logger.error(f"[Error] url: {url}, Exception={e}")
+                raise AICCException("Failed to query auc task")
+
+    def download_file(self, task_id: str, auc_file_path: str) -> bool:
+        """
+        Download and decrypt encrypted auc
+        This function polls the task status until it succeeds, then downloads
+        and decrypts the auc using the provided private key.
+        Args:
+            task_id: AUC generation task ID
+            auc_file_path: Directory to save decrypted auc
+
+        Returns:
+            True if download succeeded, False otherwise
+        """
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        if not auc_file_path:
+            raise AICCException("auc_file_path must be provided")
+
+        task_info = self.query_task(task_id)
+        status = task_info.get("Data", {}).get("Status", "")
+        if not status:
+            logger.error(f"task status is null, task_id={task_id}, query_task return {task_info}")
+            raise AICCException("task status is null")
+        if status != "success":
+            logger.info(f"task status is not success, task_id={task_id}, status={status}")
+            return False
+
+        result = task_info.get("Data", {}).get("Result") or {}
+        for label, key in (
+            ("AudioTranscriptionFile", "AudioTranscriptionFile"),
+            ("SummarizationFile", "SummarizationFile"),
+            ("ChapterFile", "ChapterFile"),
+            ("InformationExtractionFile", "InformationExtractionFile"),
+            ("TranslationFile", "TranslationFile"),
+        ):
+            file_url = result.get(key, "")
+            if file_url:
+                local_file_path = f"./{auc_file_path}_{label}.json"
+                self.download_file_by_url(file_url, local_file_path)
+
+    def download_file_by_url(self, file_url: str, local_file_path: str) -> bool:
+        """
+        Download auc file
+        Args:
+            file_url: File file_url
+            auc_file_path: Directory to save auc file
+
+        Returns:
+            Path to auc file, or None if failed
+        """
+        if not file_url:
+            raise AICCException("file_url must be provided")
+
+        if not local_file_path:
+            raise AICCException("local_file_path must be provided")
+
+        if os.path.exists(local_file_path) and os.path.getsize(local_file_path) > 0:
+            logger.error(f"[Error] File {local_file_path} already exists and is not empty")
+            return False
+
+        try:
+            response = requests.get(file_url, stream=True)
+            with open(local_file_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            if self.is_enable_auc_encrypt:
+                if not self.auc_file_encrypt_priv_key:
+                    logger.error("[Error] No private key set, cannot decrypt encrypted auc")
+                    return False
+
+                headers = response.headers
+                encrypted_key = headers.get("x-tos-meta-enc-dek")
+                if encrypted_key:
+                    try:
+                        encrypted_key_bytes = base64.b64decode(unquote(encrypted_key))
+                        encrypted_key = self.auc_file_encrypt_priv_key.decrypt(encrypted_key_bytes)
+                        tmp_file_path = local_file_path + ".tmp"
+                        decrypt_file(encrypted_key, local_file_path, tmp_file_path, 'b')
+                        os.remove(local_file_path)
+                        os.rename(tmp_file_path, local_file_path)
+                    except Exception as e:
+                        logger.error(f"[Error] Failed to decode encrypted_key: {e}")
+                        traceback.print_exc()
+                        return False
+        except Exception as e:
+            logger.error(f"[Error] Download failed: {e}")
+            traceback.print_exc()
+            return False
+        return True

src/bytedance/volcengine_aicc_sdk/common/aes.py

@@ -0,0 +1,136 @@
+__all__ = [
+    "gcm_encrypt",
+    "gcm_decrypt",
+    "aes_encrypt",
+    "aes_decrypt",
+    "encrypt_file",
+    "decrypt_file"
+]
+
+import base64
+from typing_extensions import IO, Union
+
+import secrets
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+
+SALT = "ai-handler-salt"
+
+
+def gcm_decrypt(key: bytes, ciphertext: bytes) -> bytes:
+    iv, ciphertext = ciphertext[:12], ciphertext[12:]
+    ciphertext, tag = ciphertext[:-16], ciphertext[-16:]
+    decryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv, tag),
+    ).decryptor()
+
+    return decryptor.update(ciphertext) + decryptor.finalize()
+
+
+def gcm_encrypt(key: bytes, plaintext: bytes) -> bytes:
+    iv = secrets.token_bytes(12)
+    encryptor = Cipher(
+        algorithms.AES(key),
+        modes.GCM(iv),
+    ).encryptor()
+
+    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
+    return iv + ciphertext + encryptor.tag
+
+
+def generate_key_from_string(seed_string: str) -> bytes:
+    backend = default_backend()
+    # 将字符串种子转换为字节
+    password = seed_string.encode()
+    salt = SALT.encode()
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(), length=32, salt=salt, iterations=100000, backend=backend
+    )
+    key = kdf.derive(password)
+    return key
+
+
+def aes_encrypt(aes_key: Union[bytes, str], data: Union[bytes, str]) -> bytes:
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+    if isinstance(data, str):
+        data = data.encode()
+    return gcm_encrypt(aes_key, data)
+
+
+def aes_decrypt(aes_key: Union[bytes, str], data: Union[bytes, str]) -> bytes:
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+    if isinstance(data, str):
+        data = data.encode()
+    return gcm_decrypt(aes_key, data)
+
+
+def encrypt_file(
+        aes_key: Union[bytes, str],
+        source_path: str,
+        dest_path: str,
+        mode: str = "b",
+) -> bool:
+    if mode not in ["b", "t"]:
+        raise Exception("mode argument error")
+
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+
+    with open(source_path, f"r{mode}") as source, open(
+            dest_path, f"w{mode}"
+    ) as dest:
+        if mode == "b":
+            source_b: IO[bytes] = source
+            dest_b: IO[bytes] = dest
+            plaintext = source_b.read()
+            ciphertext = aes_encrypt(aes_key, plaintext)
+            dest_b.write(ciphertext)
+        else:
+            source_t: IO[str] = source
+            dest_t: IO[str] = dest
+            for plaintext in source_t:
+                plaintext = plaintext.rstrip("\n")
+                ciphertext = aes_encrypt(aes_key, plaintext)
+                dest_t.write(base64.b64encode(ciphertext).decode() + "\n")
+
+    return True
+
+
+def decrypt_file(
+        aes_key: Union[bytes, str],
+        source_path: str,
+        dest_path: str,
+        mode: str = "b",
+) -> bool:
+    if mode not in ["b", "t"]:
+        raise Exception("mode argument error")
+
+    if isinstance(aes_key, str):
+        aes_key = aes_key.encode()
+
+    with open(source_path, f"r{mode}") as source, open(
+            dest_path, f"w{mode}"
+    ) as dest:
+        if mode == "b":
+            source_b: IO[bytes] = source
+            dest_b: IO[bytes] = dest
+            ciphertext = source_b.read()
+            plaintext = aes_decrypt(aes_key, ciphertext)
+            dest_b.write(plaintext)
+        else:
+            source_t: IO[str] = source
+            dest_t: IO[str] = dest
+            for ciphertext in source_t:
+                ciphertext = ciphertext.rstrip("\n")
+                ciphertext = base64.b64decode(ciphertext.encode())
+                plaintext = aes_decrypt(aes_key, ciphertext)
+                dest_t.write(plaintext.decode() + "\n")
+                # dest_t.write(base64.b64encode(plaintext).decode() + "\n")
+    return True

src/bytedance/volcengine_aicc_sdk/common/constant.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+
+class Constants:
+    LOG_CONFIG = "{\"dir\": \"aicc_log\"}"

src/bytedance/volcengine_aicc_sdk/common/error.py

@@ -0,0 +1,21 @@
+"""
+安全通信错误类型.
+"""
+
+__all__ = [
+    "AICCException"
+]
+
+
+class AICCException(Exception):
+    """AICC 异常"""
+
+    def __init__(self, msg):
+        Exception.__init__(self)
+        self._msg = msg
+
+    def get_msg(self):
+        return self._msg
+
+    def __str__(self):
+        return self.get_msg()

src/bytedance/volcengine_aicc_sdk/common/log/__init__.py

@@ -0,0 +1,24 @@
+__all__ = [
+    "DEBUG",
+    "INFO",
+    "WARNING",
+    "ERROR",
+    "CRITICAL",
+    "FORMAT",
+    "LogConfig",
+    "debug",
+    "info",
+    "warning",
+    "error",
+    "critical",
+    "init_log_config"
+]
+
+from .config import DEBUG, INFO, WARNING, ERROR, CRITICAL, FORMAT, LogConfig
+from .logger import init_log_config, info, debug, warning, error, critical
+
+_initialized = False
+if not _initialized:
+    config = LogConfig()
+    init_log_config(config)
+    _initialized = True

src/bytedance/volcengine_aicc_sdk/common/log/config.py

@@ -0,0 +1,31 @@
+__all__ = [
+    "DEBUG",
+    "INFO",
+    "WARNING",
+    "ERROR",
+    "CRITICAL",
+    "FORMAT",
+    "LogConfig"
+]
+
+# Log level constants
+DEBUG = "DEBUG"
+INFO = "INFO"
+WARNING = "WARNING"
+ERROR = "ERROR"
+CRITICAL = "CRITICAL"
+# Log output format
+FORMAT = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {file.path}:{line} | {message}"
+
+
+class LogConfig:
+    """Log configuration class."""
+
+    def __init__(self):
+        self.rotation = 100  # Size limit for log files in MB
+        self.retention = 7  # Days to retain log files
+        self.compression = "zip"
+        self.dir = "aicc_log"
+        self.filename = "aicc.log"
+        self.format = FORMAT
+        self.level = INFO

src/bytedance/volcengine_aicc_sdk/common/log/logger.py

@@ -0,0 +1,88 @@
+__all__ = [
+    "debug",
+    "info",
+    "warning",
+    "error",
+    "critical",
+    "init_log_config"
+]
+
+import os.path
+import sys
+from typing import Any
+
+from loguru import logger
+
+from . import LogConfig, FORMAT
+
+
+def init_log_config(*configs: LogConfig):
+    """Initialize logging configuration with given configs."""
+    logger.remove()
+
+    log_level = "INFO"
+    if os.environ.get("log_level"):
+        log_level = os.environ.get("log_level")
+    if os.environ.get("LOG_LEVEL"):
+        log_level = os.environ.get("LOG_LEVEL")
+
+    logger.add(sink=sys.stdout, format=FORMAT, level=log_level)
+
+    for config in configs:
+        if not config.dir:
+            config.dir = os.environ.get("LOG_DIR", "aicc_log")
+
+        if not config.filename:
+            config.filename = "aicc.log"
+
+        os.makedirs(config.dir, exist_ok=True)
+        fullpath = os.path.join(config.dir, config.filename)
+
+        if config.rotation <= 0:
+            config.rotation = 100
+        rotation = f"{config.rotation} MB"
+
+        if config.retention <= 0:
+            config.retention = 7
+        retention = f"{config.retention} days"
+
+        logger.add(
+            sink=fullpath,
+            rotation=rotation,
+            retention=retention,
+            compression=config.compression,
+            format=config.format,
+            level=config.level
+        )
+
+
+def msg_proc(message: str) -> str:
+    if message:
+        return f"[aicc]:{message}"
+    else:
+        return message
+
+
+def debug(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log debug message."""
+    logger.opt(depth=1).debug(msg_proc(message), *args, **kwargs)
+
+
+def info(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log info message."""
+    logger.opt(depth=1).info(msg_proc(message), *args, **kwargs)
+
+
+def warning(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log warning message."""
+    logger.opt(depth=1).warning(msg_proc(message), *args, **kwargs)
+
+
+def error(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log error message."""
+    logger.opt(depth=1).error(msg_proc(message), *args, **kwargs)
+
+
+def critical(message: str, *args: Any, **kwargs: Any) -> None:
+    """Log critical message."""
+    logger.opt(depth=1).critical(msg_proc(message), *args, **kwargs)

src/bytedance/volcengine_aicc_sdk/common/tos_utils.py

@@ -0,0 +1,341 @@
+__all__ = ["TOSUtils"]
+
+from urllib import request
+from .aes import decrypt_file
+import os
+import wget
+import time
+import httpx
+from .error import AICCException
+from .log import logger
+
+
+class TOSUtils:
+    def __init__(self, workdir: str = "./tmp"):
+        self.workdir = workdir
+        if not os.path.exists(self.workdir):
+            os.makedirs(self.workdir, exist_ok=True)
+
+        self.default_large_file_threshold = 100 * 1024 * 1024
+        self.default_download_timeout = 300.0
+        self.default_chunk_size = 1024 * 1024
+
+        self.max_retries = 3
+        self.retry_delay = 2  # seconds
+
+    def _download_small_file_wget(self, tos_url: str, local_file_path: str) -> bool:
+        """
+        Download small files using Python wget module
+        
+        Args:
+            tos_url: TOS文件路径
+            local_file_path: 本地文件路径
+            
+        Returns:
+            bool: 是否下载成功
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                logger.info(f"Download attempt {attempt}/{self.max_retries}")
+                
+                # Download file using wget
+                wget.download(tos_url, local_file_path)
+                
+                # Verify file was downloaded
+                if not os.path.exists(local_file_path):
+                    raise AICCException("File download failed, try again...")
+                
+                file_size = os.path.getsize(local_file_path)
+                if file_size == 0:
+                    raise AICCException("Downloaded file is empty, try again...")
+                
+                logger.info(f"Download successful, file size: {file_size} bytes")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Download attempt {attempt} failed: {str(e)}, try again...")
+                if os.path.exists(local_file_path):
+                    try:
+                        os.remove(local_file_path)
+                    except:
+                        pass
+                
+                if attempt == self.max_retries:
+                    raise AICCException(
+                        f"Failed to download encrypted file after {self.max_retries} attempts\n"
+                        f"Last error: {str(e)}\n"
+                        "Please check:\n"
+                        "  1. TOS URL is accessible\n"
+                        "  2. Network connection is stable\n"
+                        "  3. File exists on TOS"
+                    )
+                
+                logger.info(f"Waiting {self.retry_delay} seconds before retry...")
+                time.sleep(self.retry_delay)
+                self.retry_delay *= 2
+        
+        raise AICCException("Download failed unexpectedly")
+
+    def _download_small_file_httpx(self, tos_url: str, local_file_path: str) -> bool:
+        """
+        Download small files using httpx (fallback when wget is not available)
+        
+        Args:
+            tos_url: TOS文件路径
+            local_file_path: Path to save the downloaded file            
+        Returns:
+            bool: 是否下载成功
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                logger.info(f"Download attempt {attempt}/{self.max_retries}")
+                
+                # Download file using httpx
+                with httpx.stream('GET', tos_url, timeout=self.default_download_timeout, follow_redirects=True) as response:
+                    response.raise_for_status()
+                    
+                    with open(local_file_path, 'wb') as f:
+                        for chunk in response.iter_bytes(chunk_size=self.default_chunk_size):
+                            f.write(chunk)
+                
+                # Verify file was downloaded
+                if not os.path.exists(local_file_path):
+                    raise AICCException("File download failed, try again...")
+                
+                file_size = os.path.getsize(local_file_path)
+                if file_size == 0:
+                    raise AICCException("Downloaded file is empty")
+                
+                logger.info(f"Download successful, file size: {file_size} bytes")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Download attempt {attempt} failed: {str(e)}")
+                
+                # Clean up failed download
+                if os.path.exists(local_file_path):
+                    try:
+                        os.remove(local_file_path)
+                    except:
+                        pass
+                
+                # If this was the last attempt, raise error
+                if attempt == self.max_retries:
+                    raise AICCException(
+                        f"Failed to download encrypted file after {self.max_retries} attempts\n"
+                        f"Last error: {str(e)}\n"
+                        "Please check:\n"
+                        "  1. TOS URL is accessible\n"
+                        "  2. Network connection is stable\n"
+                        "  3. File exists on TOS"
+                    )
+                
+                # Wait before retry
+                logger.info(f"Waiting {self.retry_delay} seconds before retry...")
+                time.sleep(self.retry_delay)
+                self.retry_delay *= 2
+
+        # Should never reach here, but satisfy type checker
+        raise AICCException("Download failed unexpectedly")
+
+    def _download_large_file_streaming(self, tos_url: str, local_file_path: str, expected_size: int) -> bool:
+        """
+        Download large files using httpx streaming
+        
+        Args:
+            tos_url: TOS文件路径
+            local_file_path: Path to save the downloaded file            
+            expected_size: Expected file size in bytes
+            
+        Returns:
+            Path to the downloaded file
+        """
+        if not tos_url or not local_file_path or expected_size <= 0:
+            raise AICCException("TOS URL, local file path, and expected file size are required and must be greater than 0")
+
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                logger.info(f"Streaming download attempt {attempt}/{self.max_retries}")
+                
+                # Streaming download with timeout
+                with httpx.stream('GET', tos_url, timeout=self.default_download_timeout, follow_redirects=True) as response:
+                    response.raise_for_status()
+                    
+                    downloaded_bytes = 0
+                    
+                    with open(local_file_path, 'wb') as f:
+                        for chunk in response.iter_bytes(chunk_size=self.default_chunk_size):
+                            f.write(chunk)
+                            downloaded_bytes += len(chunk)
+                            
+                            # Log progress every 10MB
+                            if downloaded_bytes % (10 * 1024 * 1024) < self.default_chunk_size:
+                                progress = (downloaded_bytes / expected_size) * 100 if expected_size > 0 else 0
+                                logger.info(f"Download progress: {progress:.1f}% ({downloaded_bytes / (1024*1024):.1f} MB / {expected_size / (1024*1024):.1f} MB)")
+                
+                # Verify file was downloaded
+                if not os.path.exists(local_file_path):
+                    raise AICCException("File download failed, try again...")
+                
+                file_size = os.path.getsize(local_file_path)
+                if file_size == 0:
+                    raise AICCException("Downloaded file is empty")
+                
+                logger.info(f"Download successful, file size: {file_size} bytes")
+                return True
+                
+            except Exception as e:
+                logger.warning(f"Streaming download attempt {attempt} failed: {str(e)}")
+                
+                # Clean up failed download
+                if os.path.exists(local_file_path):
+                    try:
+                        os.remove(local_file_path)
+                    except:
+                        pass
+                
+                # If this was the last attempt, raise error
+                if attempt == self.max_retries:
+                    raise AICCException(
+                        f"Failed to download large file after {self.max_retries} attempts\n"
+                        f"Last error: {str(e)}\n"
+                        "Please check:\n"
+                        "  1. TOS URL is accessible\n"
+                        "  2. Network connection is stable\n"
+                        "  3. File exists on TOS\n"
+                        "  4. Timeout is sufficient for file size"
+                    )
+                
+                # Wait before retry
+                logger.info(f"Waiting {self.retry_delay} seconds before retry...")
+                time.sleep(self.retry_delay)
+                self.retry_delay *= 2
+        
+        # Should never reach here, but satisfy type checker
+        raise AICCException("Download failed unexpectedly")
+
+    def _extract_filename_from_url(self, url: str) -> str:
+        """
+        Extract filename from TOS URL
+        
+        Args:
+            url: TOS URL
+            
+        Returns:
+            Filename
+        """
+        # Remove query parameters
+        path = url.split('?')[0]
+        filename = os.path.basename(path)
+        return filename
+    
+    def download_file(self, tos_url: str, local_file_path: str, encrypted_key: bytes) -> bool:
+        """
+        下载文件到本地路径, encrypted_key有值时需要先解密文件, 再保存到本地路径
+        :return: 是否成功
+        :param tos_url: TOS文件路径
+        :param local_file_path: 本地文件路径
+        :param encrypted_key: 加密密钥, 为空时表示文件未加密
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+        
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        
+        if not tos_url.startswith('http'):
+            raise AICCException(f"Invalid TOS URL: {tos_url}\n URL must start with http:// or https://")
+        
+        try:
+            if not self._download_small_file_wget(tos_url, local_file_path):
+                raise AICCException("download file from TOS failed failed, please check the URL")
+
+            if encrypted_key:
+                tmp_file_path = local_file_path + ".tmp"
+                decrypt_file(encrypted_key, local_file_path, tmp_file_path, mode="b")
+                os.remove(local_file_path)
+                os.rename(tmp_file_path, local_file_path)
+
+            return True
+
+        except AICCException as e:
+            logger.error(f"Error extracting filename from TOS URL: {e}")
+            raise e
+        except Exception as e:
+            logger.error(f"Error extracting filename from TOS URL: {str(e)}")
+            raise AICCException("download filename from TOS file failed, please check the URL")
+
+    def download_files(self, tos_url: str, local_file_path: str, encrypted_key: bytes) -> bool:
+        """
+        下载多个文件并合并到,local_file_path, encrypted_key有值时需要先解密文件, 再保存
+        :return: 是否成功
+        :param tos_url: TOS文件路径
+        :param local_file_path: 本地文件路径
+        :param encrypted_key: AES加密密钥, 为空时表示文件未加密
+        """
+        if not tos_url or not local_file_path:
+            raise AICCException("TOS URL and local file path are required")
+        
+        if os.path.exists(local_file_path):
+            raise AICCException("Local file already exists, please delete it first")
+        
+        if not tos_url.startswith('http'):
+            raise AICCException(f"Invalid TOS URL: {tos_url}\n URL must start with http:// or https://")
+        
+        try:
+            meta_file_npath = tos_url + "/meta.json" 
+            resp = request.urlopen(meta_file_npath)
+            resp.raise_for_status()
+            meta_json = resp.json()
+            logger.info(f"Meta file content: {meta_json}")
+
+            file_list = meta_json.get("files", [])
+            if not file_list:
+                raise AICCException("No files found in the TOS directory")
+            
+            # sort file_list by filename
+            file_list.sort()
+
+            video_files = []
+            for file_name in file_list:
+                tos_file_path = tos_url + "/" + file_name
+                tmp_file_path = local_file_path + "." + file_name
+                if not self._download_small_file_wget(tos_file_path, tmp_file_path):
+                    raise AICCException(f"download file from TOS failed failed, please check the URL{tos_file_path}")
+
+                if encrypted_key:
+                    decrypt_file_path = tmp_file_path + ".dec"
+                    decrypt_file(encrypted_key, tmp_file_path, decrypt_file_path, mode="b")
+                    os.remove(tmp_file_path)
+                    os.rename(decrypt_file_path, tmp_file_path)
+
+                video_files.append(tmp_file_path)
+            
+            # 将ret_video_file中的文件合并输出到ret_video_file_path
+            with open(local_file_path, 'wb') as f:
+                for video_file in video_files:
+                    with open(video_file, 'rb') as video:
+                        f.write(video.read())
+
+            return True
+
+        except AICCException as e:
+            logger.error(f"Error extracting filename from TOS URL: {e}")
+            raise e
+        except Exception as e:
+            logger.error(f"Error extracting filename from TOS URL: {str(e)}")
+            raise AICCException("download filename from TOS file failed, please check the URL")

src/bytedance/volcengine_aicc_sdk/common/utils.py

@@ -0,0 +1,45 @@
+__all__ = [
+    "format_query",
+    "file_to_base64",
+    "base64_to_file"
+]
+
+import base64
+import json
+import urllib.parse
+from typing import Dict
+
+
+def file_to_base64(file_path):
+    with open(file_path, 'rb') as file:
+        file_data = file.read()  # 读取文件内容
+        base64_data = base64.b64encode(file_data).decode('utf-8')  # Base64 编码
+    return base64_data
+
+
+# 辅助函数：将Base64数据写入文件
+def base64_to_file(base64_data, output_file_path):
+    """将Base64编码的数据写入文件"""
+    if not base64_data:
+        print("没有Base64数据可写入文件。")
+        return False
+    try:
+        # 将Base64字符串解码为二进制数据
+        binary_data = base64.b64decode(base64_data)
+        # 写入文件
+        with open(output_file_path, 'wb') as f:
+            f.write(binary_data)
+        print(f"成功将Base64数据写入文件：{output_file_path}")
+        return True
+    except Exception as e:
+        print(f"写入文件失败：{e}")
+        return False
+
+
+def format_query(params: Dict[str, str]) -> str:
+    return "&".join(
+        urllib.parse.quote(key, safe="-_.~")
+        + "="
+        + urllib.parse.quote(value, safe="-_.~")
+        for key, value in params.items()
+    )

src/bytedance/volcengine_aicc_sdk/cu_inference_client.py

@@ -0,0 +1,231 @@
+"""
+CU开源模型机密推理客户端
+"""
+__all__ = ["CUClient"]
+
+import base64
+import json
+import time
+import requests
+import bytedance.jeddak_secure_channel as jsc
+from typing import IO, Union
+from .aicc_inference import AICCInference
+from .common.constant import Constants
+from .common.error import AICCException
+from .common.log import logger
+
+
+class CUClient(AICCInference):
+    """
+    AICC机密推理访问客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 aicc_service_name: str = None,
+                 policy_id: str = None,
+                 is_openai: bool = True,
+                 is_secure: bool = True,
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing",
+                 **kwargs):
+        if not base_url or not aicc_service_name or not policy_id or not ak or not sk:
+            raise AICCException(f"New CU client failed, params error")
+
+        super().__init__(base_url, is_openai, is_secure, ak, sk, region)
+
+        self.aicc_service_name = aicc_service_name
+        self.policy_id = policy_id
+        self.kwargs = kwargs
+
+        self.jsc_client = None
+        self.max_retries = 3
+
+        if is_secure:
+            jsc_config = dict()
+            jsc_config["ra_service_name"] = aicc_service_name
+            jsc_config["ra_policy_id"] = policy_id
+            jsc_config["bytedance_top_info"] = dict()
+            jsc_config["bytedance_top_info"]["ak"] = ak
+            jsc_config["bytedance_top_info"]["sk"] = sk
+            jsc_config["bytedance_top_info"]["region"] = region
+            jsc_config["bytedance_top_info"]["url"] = self.kwargs.get("top_url", "open.volcengineapi.com")
+            jsc_config["bytedance_top_info"]["service"] = self.kwargs.get("top_service", "pcc")
+            jsc_config["log_config"] = self.kwargs.get("log_config", Constants.LOG_CONFIG)
+
+            try:
+                secure_channel_config = jsc.ClientConfig.from_dict(jsc_config)
+                self.jsc_client = jsc.Client(secure_channel_config)
+            except Exception as e:
+                logger.error(f"Error init jsc: {e}")
+                raise AICCException("create cu client failed")
+
+    def completions(self, prompt: str, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param prompt:
+        :param output_stream:
+        :return:
+        """
+        if not prompt:
+            raise AICCException("vllm completions Prompt cannot be empty")
+
+        try:
+            if self.is_secure:
+                url = f"{self.base_url}/v2/completions"
+            else:
+                url = f"{self.base_url}/v1/completions"
+
+            headers = {"Content-Type": "application/json"}
+            payload = {
+                "prompt": prompt,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "stream": True if output_stream else False
+            }
+
+            encrypt_key = None
+            if self.is_secure:
+                # 加密数据并发送
+                try:
+                    encrypted_prompt, encrypt_key = self.jsc_client.encrypt_with_response(json.dumps(payload["prompt"]))
+                    payload["prompt"] = encrypted_prompt  # 敏感字段
+                except Exception as e:
+                    logger.error(f"Error during encrypt: {e}")
+                    raise AICCException("error during encrypt")
+
+            if output_stream:
+                try:
+                    # 启用流式响应
+                    with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                        response.raise_for_status()
+                        for chunk in response.iter_lines():
+                            if chunk:
+                                if self.is_secure and encrypt_key:
+                                    data = encrypt_key.decrypt(base64.b64decode(chunk))
+                                    output_stream.write(f"{data}")
+                                else:
+                                    output_stream.write(f"{chunk}")
+                except Exception as e:
+                    logger.error(f"call AICC VLLM service completions stream error:{e}")
+                    raise AICCException("vllm completions_stream error")
+            else:
+                for attempt in range(self.max_retries):
+                    try:
+                        response = requests.post(url, json=payload, headers=headers)
+                        response.raise_for_status()
+                        data = response.json()
+
+                        if self.is_secure and encrypt_key:
+                            for elm in data['choices']:
+                                elm['text'] = encrypt_key.decrypt(elm['text']).decode()
+
+                        return json.dumps(data)
+                    except Exception as e:
+                        if attempt == self.max_retries - 1:
+                            logger.error(f"call AICC VLLM service completions error:{e}")
+                            raise AICCException("call vllm completions error")
+                        else:
+                            logger.error(f"call AICC VLLM service completions error:{e}, try again....")
+                            time.sleep(1)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"completions error,Exception={e}")
+            raise AICCException("completions error")
+
+    def completions_stream(self, prompt: str, output_stream: IO) -> Union[str, None]:
+        return self.completions(prompt, output_stream)
+
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param output_stream:
+        :return:
+        """
+        try:
+            if not messages:
+                raise AICCException("vllm chat_completions Messages cannot be empty")
+
+            if self.is_secure:
+                url = f"{self.base_url}/v2/chat/completions"
+            else:
+                url = f"{self.base_url}/v1/chat/completions"
+
+            headers = {"Content-Type": "application/json"}
+            payload = {
+                "messages": messages,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "stream": True if output_stream else False  # 启用流式响应
+            }
+
+            encrypt_key = None
+            if self.is_secure:
+                # 加密数据并发送
+                try:
+                    encrypted_messages, encrypt_key = \
+                        self.jsc_client.encrypt_with_response(json.dumps(payload["messages"]))
+                    payload["messages"] = encrypted_messages  # 敏感字段
+                except Exception as e:
+                    logger.error(f"Error during encrypt: {e}")
+                    raise AICCException("Error during encrypt")
+
+            if output_stream:
+                with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                    response.raise_for_status()
+                    for chunk in response.iter_lines():
+                        if chunk:
+                            if self.is_secure and encrypt_key:
+                                data = encrypt_key.decrypt(base64.b64decode(chunk))
+                                output_stream.write(f"{data}")
+                            else:
+                                output_stream.write(f"{chunk}")
+            else:
+                for attempt in range(self.max_retries):
+                    try:
+                        response = requests.post(url, json=payload, headers=headers)
+                        response.raise_for_status()
+                        data = response.json()
+
+                        if self.is_secure and encrypt_key:
+                            for elm in data['choices']:
+                                elm['message'] = json.loads(encrypt_key.decrypt(elm['message']))
+
+                        return json.dumps(data)
+                    except Exception as e:
+                        if attempt == self.max_retries - 1:
+                            logger.error(f"call AICC VLLM service chat_completions error:{e}")
+                            raise AICCException("vllm chat_completions error")
+                        else:
+                            logger.error(f"call AICC VLLM service chat_completions error:{e}, try again....")
+                            time.sleep(1)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC VLLM service chat_completions error: {e}")
+            raise AICCException("AICC VLLM service chat_completions error")
+
+    def chat_completions_stream(self, messages: list, output_stream: IO) -> Union[str, None]:
+        return self.chat_completions(messages, output_stream)
+
+    def function_call(self, messages: list, tools: list, output_stream: IO = None):
+        """
+        :param messages:
+        :param tools:
+        :param output_stream:
+        :return:
+        """
+        error_info = "Sorry, this version does not support function_call at this time."
+        logger.error(error_info)
+        raise AICCException(error_info)
+
+    def function_call_stream(self, messages: list, tools: list, output_stream: IO):
+        return self.function_call(messages, tools, output_stream)
+
+    def asr_recognize(self, file_name: str) -> Union[str, None]:
+        error_info = "Sorry, this version does not support asr at this time."
+        logger.error(error_info)
+        return error_info

src/bytedance/volcengine_aicc_sdk/deepseek_inference_client.py

@@ -0,0 +1,501 @@
+"""
+DeepSeek机密推理客户端
+"""
+__all__ = ["DeepSeekClient"]
+
+import base64
+import copy
+import json
+import time
+import requests
+import pickle
+import bytedance.jeddak_secure_channel as jsc
+from typing import IO, Union
+
+from openai.types.chat import ChatCompletionMessage, ChatCompletionChunk
+
+from .aicc_inference import AICCInference
+from .common.constant import Constants
+from .common.error import AICCException
+from .common.log import logger
+from openai import OpenAI
+
+
+class DeepSeekClient(AICCInference):
+    """
+    DeepSeek机密推理客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = None,
+                 api_key: str = None,
+                 ep: str = None,
+                 aicc_service_name: str = None,
+                 policy_id: str = None,
+                 is_openai: bool = True,
+                 is_secure: bool = True,
+                 ak: str = None,
+                 sk: str = None,
+                 region: str = "cn-beijing",
+                 **kwargs):
+        if not base_url or not api_key or not aicc_service_name or not policy_id or not ak or not sk:
+            raise AICCException("New deepseek client failed, params error")
+
+        super().__init__(base_url, is_openai, is_secure, ak, sk, region)
+
+        self.api_key = api_key
+        self.ep = ep
+        self.aicc_service_name = aicc_service_name
+        self.policy_id = policy_id
+
+        self.kwargs = kwargs
+
+        self.dp_version = self.kwargs.get("dp_version", "3.1")
+        self.jsc_client = None
+        self.max_retries = 3
+
+        if is_secure:
+            jsc_config = dict()
+            jsc_config["ra_service_name"] = aicc_service_name
+            jsc_config["ra_policy_id"] = policy_id
+            jsc_config["bytedance_top_info"] = dict()
+            jsc_config["bytedance_top_info"]["ak"] = ak
+            jsc_config["bytedance_top_info"]["sk"] = sk
+            jsc_config["bytedance_top_info"]["region"] = region
+            jsc_config["bytedance_top_info"]["url"] = self.kwargs.get("top_url", "open.volcengineapi.com")
+            jsc_config["bytedance_top_info"]["service"] = self.kwargs.get("top_service", "pcc")
+            jsc_config["log_config"] = self.kwargs.get("log_config", Constants.LOG_CONFIG)
+
+            try:
+                secure_channel_config = jsc.ClientConfig.from_dict(jsc_config)
+                self.jsc_client = jsc.Client(secure_channel_config)
+            except Exception as e:
+                logger.error(f"Error init jsc: {e}")
+                raise AICCException("New deepseek client failed")
+
+    def _encrypt_prompt_with_jsc(self, messages: list) -> object:
+        """
+        使用端云互信SDK对prompt进行加密
+        :param messages:
+        :return: 返回对称解密密钥（多密钥时返回 role user 的content加密 密钥）
+        """
+        if not messages or not self.jsc_client:
+            logger.error(f"messages is NULL")
+            return None
+
+        ret_encrypt_key = None
+        for i in range(len(messages)):
+            messages[i]["content"], encrypt_key = self.jsc_client.encrypt_with_response(
+                json.dumps(messages[i].get("content")))
+            if messages[i].get("role") == "user":
+                ret_encrypt_key = encrypt_key
+
+        return ret_encrypt_key
+
+    def _chat_completion_message_to_str(self, chat_message: ChatCompletionMessage) -> str:
+        ret = dict()
+        ret["role"] = chat_message.role
+        ret["content"] = chat_message.content
+        ret["reasoning_content"] = chat_message.reasoning_content
+        ret["tool_calls"] = chat_message.tool_calls
+        ret["refusal"] = chat_message.refusal
+        ret["annotations"] = chat_message.annotations
+        ret["audio"] = chat_message.audio
+        ret["function_call"] = chat_message.function_call
+        return json.dumps(ret)
+
+    def _chat_completion_chunk_to_str(self, chat_completion_chunk: ChatCompletionChunk) -> str:
+        ret = dict()
+        ret["id"] = chat_completion_chunk.id
+        ret["object"] = chat_completion_chunk.object
+        ret["created"] = chat_completion_chunk.created
+        ret["model"] = chat_completion_chunk.model
+
+        if chat_completion_chunk.usage:
+            ret["usage"] = dict()
+            ret["usage"]["completion_tokens"] = chat_completion_chunk.usage.completion_tokens
+            ret["usage"]["prompt_tokens"] = chat_completion_chunk.usage.prompt_tokens
+            ret["usage"]["total_tokens"] = chat_completion_chunk.usage.total_tokens
+            ret["usage"]["completion_tokens_details"] = chat_completion_chunk.usage.completion_tokens_details
+            ret["usage"]["prompt_tokens_details"] = chat_completion_chunk.usage.prompt_tokens_details
+            ret["usage"]["reasoning_tokens"] = chat_completion_chunk.usage.reasoning_tokens
+        else:
+            ret["usage"] = None
+
+        ret["choices"] = list()
+
+        for choice in chat_completion_chunk.choices:
+            choice_info = dict()
+            choice_info["index"] = choice.index
+            choice_info["logprobs"] = choice.logprobs
+            choice_info["finish_reason"] = choice.finish_reason
+            choice_info["matched_stop"] = choice.matched_stop
+            choice_info["delta"] = dict()
+            choice_info["delta"]["content"] = choice.delta.content
+            choice_info["delta"]["function_call"] = choice.delta.function_call
+            choice_info["delta"]["refusal"] = choice.delta.refusal
+            choice_info["delta"]["role"] = choice.delta.role
+            choice_info["delta"]["tool_calls"] = choice.delta.tool_calls
+            choice_info["delta"]["reasoning_content"] = choice.delta.reasoning_content
+
+            ret["choices"].append(choice_info)
+
+        return json.dumps(ret)
+
+    def chat_completions(self, messages: list, output_stream: IO = None) -> Union[str, None]:
+        """
+        生成聊天补全
+        :param messages:
+        :param output_stream:
+        :return:
+        """
+        try:
+            if not messages:
+                raise AICCException("DeepSeek chat_completions Messages cannot be empty")
+
+            # DeepSeek目前仅3.1支持openai的方式
+            if self.dp_version < "3.1" and self.is_openai:
+                error_info = f"Sorry, DeepSeek{self.dp_version} not support is_openai={self.is_openai}"
+                logger.error(error_info)
+                raise AICCException(error_info)
+
+            if self.is_secure:
+                # V3.1用V4， V3.0用V2
+                if self.dp_version < "3.1":
+                    url = f"{self.base_url}/v2/chat/completions"
+                else:
+                    url = f"{self.base_url}/v4/chat/completions"
+            else:
+                url = f"{self.base_url}/v1/chat/completions"
+
+            headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
+            payload = {
+                "messages": messages,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "chat_template_kwargs": {"thinking": self.kwargs.get("thinking", True)},
+                "separate_reasoning": True
+            }
+            if output_stream:
+                payload["stream"] = True
+                payload["stop"] = "<|im_end|>"
+                payload["eos_token"] = "<|im_end|>"
+
+            if self.is_openai:
+                openai_messages = copy.deepcopy(messages)
+                encrypt_key = None
+                if self.is_secure:
+                    encrypt_key = self._encrypt_prompt_with_jsc(openai_messages)
+                    if not encrypt_key:
+                        raise AICCException("Error during openai encrypt")
+
+                url = url.rstrip("/chat/completions")
+                openai_client = OpenAI(base_url=url, api_key=self.api_key)
+
+                resp = openai_client.chat.completions.create(
+                    model=self.ep if self.ep else "",
+                    messages=openai_messages,
+                    temperature=payload.get("temperature"),
+                    max_tokens=payload.get("max_tokens"),
+                    stream=True if output_stream else False,
+                )
+
+                if output_stream:
+                    for chunk in resp:
+                        if self.is_secure and encrypt_key:
+                            if chunk.content:
+                                data = encrypt_key.decrypt(base64.b64decode(chunk.content)).decode("utf-8")
+                                if data.startswith("data:"):
+                                    data = data.strip("data:").strip()
+                                output_stream.write(f"{data}\n")
+                        else:
+                            cont = self._chat_completion_chunk_to_str(chunk)
+                            output_stream.write(f"{cont}\n")
+                    return
+                else:
+                    resp_msg = resp.choices[0].message
+                    if self.is_secure and encrypt_key:
+                        resp_msg = encrypt_key.decrypt(resp_msg).decode("utf-8")
+                        return resp_msg
+                    else:
+                        return self._chat_completion_message_to_str(resp_msg)
+
+            encrypt_key = None
+            if self.is_secure:
+                # 加密数据并发送
+                try:
+                    if self.dp_version < "3.1":
+                        encrypted_messages, encrypt_key = \
+                            self.jsc_client.encrypt_with_response(pickle.dumps(payload["messages"]))
+                    else:
+                        encrypted_messages, encrypt_key = \
+                            self.jsc_client.encrypt_with_response(json.dumps(payload["messages"]))
+
+                    payload["messages"] = encrypted_messages  # 敏感字段
+                except Exception as e:
+                    logger.error(f"Error during encrypt: {e}")
+                    raise AICCException("Error during encrypt")
+
+            if output_stream:
+                try:
+                    with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                        response.raise_for_status()
+                        for chunk in response.iter_lines():
+                            if chunk:
+                                if self.is_secure and encrypt_key:
+                                    if self.dp_version < "3.1":
+                                        data = encrypt_key.decrypt(base64.b64decode(chunk)).decode("utf-8")
+                                        output_stream.write(f"{data}\n")
+                                    else:
+                                        chunk_str = chunk.decode("utf-8")
+                                        if chunk_str.startswith("data:"):
+                                            chunk_str = chunk_str.strip("data:").strip()
+                                        if chunk_str == "[DONE]":
+                                            output_stream.write(chunk_str + "\n")
+                                        else:
+                                            content = json.loads(chunk_str).get("content")
+                                            data = encrypt_key.decrypt(base64.b64decode(content)).decode("utf-8")
+                                            output_stream.write(f"{data}\n")
+                                else:
+                                    chunk_str = chunk.decode("utf-8")
+                                    if chunk_str.startswith("data:"):
+                                        chunk_str = chunk_str.strip("data:").strip()
+                                    output_stream.write(f"{chunk_str}\n")
+                except Exception as e:
+                    logger.error(f"call AICC DeepSeek service chat_completions error:{e}")
+                    raise AICCException("DeepSeek chat_completions_stream error")
+            else:
+                for attempt in range(self.max_retries):
+                    try:
+                        response = requests.post(url, json=payload, headers=headers)
+                        response.raise_for_status()
+                        data = response.json()
+
+                        if self.is_secure and encrypt_key:
+                            for elm in data['choices']:
+                                if self.dp_version < "3.1":
+                                    elm['message'] = pickle.loads(encrypt_key.decrypt(elm['message']))
+                                else:
+                                    elm['message'] = json.loads(encrypt_key.decrypt(elm['message']))
+
+                        return json.dumps(data)
+                    except Exception as e:
+                        if attempt == self.max_retries - 1:
+                            logger.error(f"call AICC DeepSeek service chat_completions error:{e}")
+                            raise AICCException("DeepSeek chat_completions error")
+                        else:
+                            logger.error(f"call AICC DeepSeek service chat_completions error:{e}, try again....")
+                            time.sleep(1)
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC deepseek service chat_completions error: {e}")
+            raise AICCException("AICC deepseek service chat_completions error")
+
+    def chat_completions_stream(self, messages: list, output_stream: IO) -> Union[str, None]:
+        return self.chat_completions(messages, output_stream)
+
+    def chat_completions_stream_yield(self, messages: list):
+        """
+        生成聊天补全
+        :param messages:
+        :return:
+        """
+        try:
+            if not messages:
+                raise AICCException("DeepSeek chat_completions Messages cannot be empty")
+
+            # DeepSeek目前仅3.1支持openai的方式
+            if self.dp_version < "3.1" and self.is_openai:
+                error_info = f"Sorry, DeepSeek{self.dp_version} not support is_openai={self.is_openai}"
+                logger.error(error_info)
+                raise AICCException(error_info)
+
+            if self.is_secure:
+                # V3.1用V4， V3.0用V2
+                if self.dp_version < "3.1":
+                    url = f"{self.base_url}/v2/chat/completions"
+                else:
+                    url = f"{self.base_url}/v4/chat/completions"
+            else:
+                url = f"{self.base_url}/v1/chat/completions"
+
+            headers = {"Content-Type": "application/json", "Authorization": f"Bearer {self.api_key}"}
+            payload = {
+                "messages": messages,
+                "max_tokens": int(self.kwargs.get("max_tokens", 100)),
+                "temperature": float(self.kwargs.get("temperature", 1.0)),
+                "chat_template_kwargs": {"thinking": self.kwargs.get("thinking", True)},
+                "separate_reasoning": True,
+                "stream": True,
+                "stop": "<|im_end|>",
+                "eos_token": "<|im_end|>"
+            }
+
+            if self.is_openai:
+                openai_messages = copy.deepcopy(messages)
+                encrypt_key = None
+                if self.is_secure:
+                    encrypt_key = self._encrypt_prompt_with_jsc(openai_messages)
+                    if not encrypt_key:
+                        raise AICCException("Error during openai encrypt")
+
+                url = url.rstrip("/chat/completions")
+                openai_client = OpenAI(base_url=url, api_key=self.api_key)
+
+                resp = openai_client.chat.completions.create(
+                    model=self.ep if self.ep else "",
+                    messages=openai_messages,
+                    temperature=payload.get("temperature"),
+                    max_tokens=payload.get("max_tokens"),
+                    stream=True,
+                )
+
+                for chunk in resp:
+                    if self.is_secure and encrypt_key:
+                        if chunk.content:
+                            data = encrypt_key.decrypt(base64.b64decode(chunk.content)).decode("utf-8")
+                            if data.startswith("data:"):
+                                data = data.strip("data:").strip()
+                            yield data
+                    else:
+                        cont = self._chat_completion_chunk_to_str(chunk)
+                        yield cont
+            else:
+                encrypt_key = None
+                if self.is_secure:
+                    # 加密数据并发送
+                    try:
+                        if self.dp_version < "3.1":
+                            encrypted_messages, encrypt_key = \
+                                self.jsc_client.encrypt_with_response(pickle.dumps(payload["messages"]))
+                        else:
+                            encrypted_messages, encrypt_key = \
+                                self.jsc_client.encrypt_with_response(json.dumps(payload["messages"]))
+
+                        payload["messages"] = encrypted_messages  # 敏感字段
+                    except Exception as e:
+                        logger.error(f"Error during encrypt: {e}")
+                        raise AICCException("Error during encrypt")
+
+                try:
+                    with requests.post(url, json=payload, headers=headers, stream=True) as response:
+                        response.raise_for_status()
+                        for chunk in response.iter_lines():
+                            if chunk:
+                                if self.is_secure and encrypt_key:
+                                    chunk_str = chunk.decode("utf-8")
+                                    if chunk_str.startswith("data:"):
+                                        chunk_str = chunk_str.strip("data:").strip()
+                                    if chunk_str == "[DONE]":
+                                        yield chunk_str
+                                    else:
+                                        if self.dp_version < "3.1":
+                                            data = encrypt_key.decrypt(base64.b64decode(chunk_str)).decode("utf-8")
+                                        else:
+                                            content = json.loads(chunk_str).get("content")
+                                            data = encrypt_key.decrypt(base64.b64decode(content)).decode("utf-8")
+                                        yield data
+                                else:
+                                    chunk_str = chunk.decode("utf-8")
+                                    if chunk_str.startswith("data:"):
+                                        chunk_str = chunk_str.strip("data:").strip()
+                                    yield chunk_str
+                except Exception as e:
+                    logger.error(f"call AICC DeepSeek service chat_completions_stream_yield error:{e}")
+                    raise AICCException("DeepSeek chat_completions_stream_yield error")
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC deepseek service chat_completions_stream_yield error: {e}")
+            raise AICCException("AICC deepseek service chat_completions_stream_yield error")
+
+    def function_call(self, messages: list, tools: list, output_stream: IO = None) -> str:
+        """
+        :param messages:
+        :param tools:
+        :param output_stream:
+        :return:
+        """
+        if not messages or not tools:
+            raise AICCException("DeepSeek function_call params cannot be empty")
+
+        try:
+            # DeepSeek目前仅支持通过openai的方式使用function_call
+            if not self.is_openai:
+                error_info = f"Sorry, DeepSeek is_openai={self.is_openai} not support function_call."
+                logger.error(error_info)
+                raise AICCException(error_info)
+
+            # DeepSeek function_call目前仅支持密文访问
+            if not self.is_secure:
+                error_info = f"Sorry, DeepSeek secure={self.is_secure} not support function_call."
+                logger.error(error_info)
+                raise AICCException(error_info)
+
+            if self.is_secure:
+                url = f"{self.base_url}/v4"
+            else:
+                url = f"{self.base_url}/v1"
+
+            openai_client = OpenAI(base_url=url, api_key=self.api_key)
+
+            openai_messages = copy.deepcopy(messages)
+            encrypt_key = None
+            if self.is_secure:
+                encrypt_key = self._encrypt_prompt_with_jsc(openai_messages)
+                if not encrypt_key:
+                    raise AICCException("Error during openai encrypt")
+
+            if self.is_secure:
+                if self.dp_version < "3.1":
+                    encrypted_tools, _ = self.jsc_client.encrypt_with_response(pickle.dumps(tools))
+                else:
+                    encrypted_tools, _ = self.jsc_client.encrypt_with_response(json.dumps(tools))
+            else:
+                encrypted_tools = tools
+
+            extra_fields = {
+                "aicc_enc_tools": encrypted_tools
+            }
+
+            completion = openai_client.chat.completions.create(
+                model=self.ep,
+                messages=openai_messages,
+                tool_choice="required",
+                extra_body=extra_fields,
+                stream=True if output_stream else False,
+            )
+
+            if output_stream:
+                for chunk in completion:
+                    if not hasattr(chunk, 'content') or chunk.content is None:
+                        continue
+                    if chunk.content == "[DONE]":
+                        output_stream.write("[DONE]\n")
+                    if self.is_secure and encrypt_key:
+                        decrypted_data = encrypt_key.decrypt(base64.b64decode(chunk.content)).decode("utf-8")
+                        if decrypted_data.startswith("data:"):
+                            decrypted_data = decrypted_data.strip("data:").strip()
+                        output_stream.write(f"{decrypted_data}\n")
+                    else:
+                        output_stream.write(f"{chunk.content}\n")
+            else:
+                resp_msg = completion.choices[0].message
+                if self.is_secure and encrypt_key:
+                    resp_msg = encrypt_key.decrypt(resp_msg).decode("utf-8")
+
+                return resp_msg
+        except AICCException as e:
+            raise e
+        except Exception as e:
+            logger.error(f"call AICC deepseek service chat_completions error: {e}")
+            raise AICCException("AICC deepseek service chat_completions error")
+
+    def function_call_stream(self, messages: list, tools: list, output_stream: IO) -> str:
+        return self.function_call(messages, tools, output_stream)
+
+    def asr_recognize(self, file_name: str) -> str:
+        error_info = "Sorry, this version does not support asr at this time."
+        logger.error(error_info)
+        return error_info

src/bytedance/volcengine_aicc_sdk/secure_http_client.py

@@ -0,0 +1,297 @@
+import httpx
+from typing import Optional, Tuple
+from urllib.parse import urlparse
+from bytedance.jeddak_secure_channel import ResponseKey
+import bytedance.jeddak_secure_channel as jsc
+from .common.error import AICCException
+from .common.log import logger
+
+
+class AICCConfig:
+    """AICC 安全通信配置类"""
+
+    def __init__(self,
+                 ak: str = "",
+                 sk: str = "",
+                 aicc_llm_service_name: str = "AICC.ConfidentialInference",
+                 policy_id: str = "router_policy",
+                 region: str = "cn-beijing",
+                 top_url: str = "pcc.volcengineapi.com",
+                 top_service: str = "pcc",
+                 ra_url: str = "",
+                 auth_token: str = "",
+                 **kwargs):
+        """
+        :param ak: 火山引擎用户的ak
+        :param sk: 火山引擎用户的sk
+        :param aicc_llm_service_name: aicc推理服务的服务名称
+        :param policy_id: aicc推理服务做远程证明锁对应的policy
+        :param region: aicc推理服务所在的区域
+        :param service: aicc推理服务TOP网关的service
+        :param ra_url: 非TCA接口时的RA接口地址
+        :param auth_token: 非TCA接口时的鉴权凭据
+        :param kwargs:
+        """
+        self.ak = ak
+        self.sk = sk
+        self.aicc_llm_service_name = aicc_llm_service_name
+        self.policy_id = policy_id
+        self.region = region
+        self.top_url = top_url
+        self.top_service = top_service
+        self.ra_url = ra_url
+        self.auth_token = auth_token
+
+        self.kwargs = kwargs
+
+        self.HTTP_RA_PATH = "v1/security/token"
+        if self.ra_url:
+            if self.ra_url.startswith("wss://"):
+                self.ra_url = self.ra_url.replace("wss", "https")
+            parsed = urlparse(self.ra_url)
+            base_domain = f"{parsed.scheme}://{parsed.netloc}"
+            self.ra_url = f"{base_domain}/{self.HTTP_RA_PATH}"
+
+    def to_dict(self) -> dict:
+        """将配置对象转换为字典格式
+
+        :return: 包含所有配置项的字典
+        """
+        config_dict = {
+            "ak": self.ak,
+            "sk": self.sk,
+            "aicc_llm_service_name": self.aicc_llm_service_name,
+            "policy_id": self.policy_id,
+            "region": self.region,
+            "top_url": self.top_url,
+            "top_service": self.top_service,
+            "ra_url": self.ra_url,
+            "auth_token": self.auth_token
+        }
+
+        # 添加额外的 kwargs
+        if self.kwargs:
+            config_dict.update(self.kwargs)
+
+        return config_dict
+
+    def get_jsc_config(self) -> dict:
+        """生成aicc端云互信配置
+        """
+        jsc_config = dict()
+        jsc_config["ra_url"] = self.ra_url
+        jsc_config["auth_token"] = self.auth_token
+        jsc_config["ra_service_name"] = self.aicc_llm_service_name
+        jsc_config["ra_policy_id"] = self.policy_id
+
+        if self.ak and self.sk:
+            jsc_config["bytedance_top_info"] = dict()
+            jsc_config["bytedance_top_info"]["ak"] = self.ak
+            jsc_config["bytedance_top_info"]["sk"] = self.sk
+            jsc_config["bytedance_top_info"]["region"] = self.region
+            jsc_config["bytedance_top_info"]["url"] = self.top_url
+            jsc_config["bytedance_top_info"]["service"] = self.top_service
+
+        return jsc_config
+
+
+AICC_SDK_VERSION = "0.1.0"
+jsc_clients = dict()
+
+
+class DemoSyncByteStream(httpx.SyncByteStream):
+    def __init__(self, raw_stream, encrypt_key: ResponseKey = None) -> None:
+        self._raw_stream = raw_stream
+        self.encrypt_key = encrypt_key
+        self.buffer = b""
+
+    def __iter__(self):
+        for part in self._raw_stream:
+            if not isinstance(part, bytes):
+                part = part.encode("utf-8")
+
+            for line in part.splitlines(keepends=True):
+                self.buffer += line
+                if self.buffer.endswith((b"\r", b"\n")):
+                    cipher_events = self.encrypt_key.decrypt(self.buffer.strip())
+                    cipher_events += b"\n" if self.buffer[-1] == 10 else b""  # 10 == "\n"
+                    self.buffer = b""
+                    yield cipher_events
+
+        if self.buffer:
+            yield self.buffer
+
+    def close(self) -> None:
+        if hasattr(self._raw_stream, "close"):
+            self._raw_stream.close()
+
+
+class DemoAsyncByteStream(httpx.AsyncByteStream):
+    def __init__(self, raw_stream, encrypt_key: ResponseKey = None) -> None:
+        self._raw_stream = raw_stream
+        self.encrypt_key = encrypt_key
+        self.buffer = b""
+
+    async def __aiter__(self):
+        async for part in self._raw_stream:
+            if not isinstance(part, bytes):
+                part = part.encode("utf-8")
+
+            for line in part.splitlines(keepends=True):
+                self.buffer += line
+                if self.buffer.endswith((b"\r", b"\n")):
+                    cipher_events = self.encrypt_key.decrypt(self.buffer.strip())
+                    cipher_events += b"\n" if self.buffer[-1] == 10 else b""  # 10 == "\n"
+                    self.buffer = b""
+                    yield cipher_events
+
+        if self.buffer:
+            yield self.buffer
+
+    async def aclose(self) -> None:
+        if hasattr(self._raw_stream, "aclose"):
+            await self._raw_stream.aclose()
+        elif hasattr(self._raw_stream, "close"):
+            self._raw_stream.close()
+
+
+class SecureHttpClient(httpx.Client):
+    def __init__(self, *args, **kwargs):
+        self.aicc_config: Optional[AICCConfig] = None
+
+        if kwargs.get("aicc_config"):
+            try:
+                self.aicc_config = kwargs.pop("aicc_config")
+                secure_channel_config = jsc.ClientConfig.from_dict(self.aicc_config.get_jsc_config())
+                jsc_clients[self.aicc_config.aicc_llm_service_name] = jsc.Client(secure_channel_config)
+            except Exception as e:
+                logger.error(f"Error init jsc, aicc_llm_service_name={self.aicc_config.aicc_llm_service_name}, {e}")
+                self.aicc_config = None
+                raise AICCException("New deepseek client failed")
+
+        super().__init__(*args, **kwargs)
+
+    def jsc_encryption(self, msg: bytes) -> Tuple[str, ResponseKey]:
+        encrypt_key: Optional[ResponseKey] = None
+
+        if self.aicc_config:
+            jsc_client = jsc_clients.get(self.aicc_config.aicc_llm_service_name)
+            if jsc_client:
+                msg, encrypt_key = jsc_client.encrypt_with_response(msg)
+
+        return msg, encrypt_key
+
+    def send(self, request: httpx.Request, *args, **kwargs):
+        encrypt_key: Optional[ResponseKey] = None
+
+        if self.aicc_config:
+            # 对所有请求进行加密（包括multipart请求）
+            request.headers["X-AICC-Encryption-Enable"] = "true"
+            request.headers["X-AICC-Encryption-SDK"] = "aicc"
+            request.headers["X-AICC-Encryption-Version"] = AICC_SDK_VERSION
+
+            # 先读取内容（处理流式请求）
+            try:
+                content = request.content
+            except httpx.RequestNotRead:
+                content = request.read()
+            
+            encryption_content, encrypt_key = self.jsc_encryption(content)
+            if isinstance(encryption_content, str):
+                encryption_content = encryption_content.encode("utf-8")
+
+            request.stream = httpx.ByteStream(encryption_content)
+            delattr(request, "_content")
+            # 不修改Content-Type，保持原始的Content-Type（如multipart/form-data）
+            request.headers["Content-Length"] = str(len(encryption_content))
+            request.read()
+
+        resp = super().send(request, *args, **kwargs)
+
+        if 200 <= resp.status_code < 300:
+            if encrypt_key:
+                if not kwargs.get("stream", False):
+                    content = resp.content
+                    decrypted_content = encrypt_key.decrypt(content)
+                    resp._content = decrypted_content
+                else:
+                    resp.stream = DemoSyncByteStream(resp.stream, encrypt_key)
+        else:
+            if encrypt_key:
+                try:
+                    if not kwargs.get("stream", False):
+                        content = resp.content
+                        decrypted_content = encrypt_key.decrypt(content)
+                        resp._content = decrypted_content
+                    else:
+                        resp.stream = DemoSyncByteStream(resp.stream, encrypt_key)
+                except Exception as e:
+                    pass
+        return resp
+
+
+class AsyncSecureHttpClient(httpx.AsyncClient):
+    def __init__(self, *args, **kwargs):
+        self.aicc_config: Optional[AICCConfig] = None
+
+        if kwargs.get("aicc_config"):
+            try:
+                self.aicc_config = kwargs.pop("aicc_config")
+                secure_channel_config = jsc.ClientConfig.from_dict(self.aicc_config.get_jsc_config())
+                jsc_clients[self.aicc_config.aicc_llm_service_name] = jsc.Client(secure_channel_config)
+            except Exception as e:
+                logger.error(f"Error init jsc, aicc_llm_service_name={self.aicc_config.aicc_llm_service_name}, {e}")
+                self.aicc_config = None
+                raise AICCException("New async secure client failed")
+
+        super().__init__(*args, **kwargs)
+
+    def jsc_encryption(self, msg: bytes) -> Tuple[str, ResponseKey]:
+        encrypt_key: Optional[ResponseKey] = None
+
+        if self.aicc_config:
+            jsc_client = jsc_clients.get(self.aicc_config.aicc_llm_service_name)
+            if jsc_client:
+                msg, encrypt_key = jsc_client.encrypt_with_response(msg)
+
+        return msg, encrypt_key
+
+    async def send(self, request: httpx.Request, *args, **kwargs):
+        encrypt_key: Optional[ResponseKey] = None
+
+        if self.aicc_config:
+            request.headers["X-AICC-Encryption-Enable"] = "true"
+            request.headers["X-AICC-Encryption-SDK"] = "aicc"
+            request.headers["X-AICC-Encryption-Version"] = AICC_SDK_VERSION
+
+            encryption_content, encrypt_key = self.jsc_encryption(request.content)
+            if isinstance(encryption_content, str):
+                encryption_content = encryption_content.encode("utf-8")
+
+            request.stream = httpx.ByteStream(encryption_content)
+            delattr(request, "_content")
+            request.headers["Content-Type"] = "application/json"
+            request.headers["Content-Length"] = str(len(encryption_content))
+
+        resp = await super().send(request, *args, **kwargs)
+
+        if 200 <= resp.status_code < 300:
+            if encrypt_key:
+                if not kwargs.get("stream", False):
+                    content = await resp.aread()
+                    decrypted_content = encrypt_key.decrypt(content)
+                    resp._content = decrypted_content
+                else:
+                    resp.stream = DemoAsyncByteStream(resp.stream, encrypt_key)
+        else:
+            if encrypt_key:
+                try:
+                    if not kwargs.get("stream", False):
+                        content = await resp.aread()
+                        decrypted_content = encrypt_key.decrypt(content)
+                        resp._content = decrypted_content
+                    else:
+                        resp.stream = DemoAsyncByteStream(resp.stream, encrypt_key)
+                except Exception as e:
+                    pass
+        return resp

src/bytedance/volcengine_aicc_sdk/seedance_inference_client.py

@@ -0,0 +1,500 @@
+"""
+seedance机密推理客户端
+"""
+__all__ = ["SeedanceClient"]
+
+import os
+import base64
+import requests
+import traceback
+from urllib.parse import unquote
+from cryptography.hazmat.primitives import serialization as crypto_serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.backends import default_backend as crypto_default_backend
+from bytedance.jeddak_secure_channel.crypto import PrivateKey
+from .secure_http_client import SecureHttpClient
+from .common.tos_utils import TOSUtils
+from .common.error import AICCException
+from .common.log import logger
+from .secure_http_client import AICCConfig
+from .common.aes import decrypt_file
+
+
+class SeedanceClient:
+    """
+    DeepSeek机密推理客户端.
+    """
+
+    def __init__(self,
+                 base_url: str = "",
+                 api_key: str = "",
+                 endpoint: str = "",
+                 aicc_config: AICCConfig = None,
+                 **kwargs):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.endpoint = endpoint
+        self.aicc_config = aicc_config
+        self.tos_utils = TOSUtils()
+        self.kwargs = kwargs
+        self.is_secure = self.kwargs.get("is_secure", True)
+
+        if aicc_config:
+            self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+        else:
+            if self.is_secure:
+                aicc_config = AICCConfig(ra_url=base_url, auth_token=api_key)
+                self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+            else:
+                self.secure_http_client = SecureHttpClient()
+
+        # 是否开启视频文件加密, 默认加密
+        self.is_enable_video_encrypt = str(self.kwargs.get("is_enable_video_encrypt", True)).lower() == "true"
+
+        # 视频文件的存储位置, 目前仅支持火山TOS
+        self.video_file_storage_location = self.kwargs.get("video_file_storage_location", "TOS").upper()
+
+        self.video_file_encrypt_pub_key_content = None
+        self.video_file_encrypt_priv_key_content = None
+
+        self.video_file_encrypt_priv_key: PrivateKey = None
+
+        self.timeout = self.kwargs.get("timeout", 120.0)
+
+        # API Endpoints
+        self.CREATE_TASK_URL = "/contents/generations/tasks"
+        self.QUERY_TASK_URL = "/contents/generations/tasks"
+        self.QUERY_TASK_LIST_URL = "/contents/generations/tasks"
+        self.DELETE_TASK_URL = "/contents/generations/tasks"
+
+    def _get_base_url(self):
+        """
+        Get base URL with trailing slash removed
+        """
+        return self.base_url.rstrip("/")
+
+    def set_video_file_encrypt_key(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Set video file encrypt key
+        Args:
+            pub_key_path: Public key for encrypt
+            priv_key_path: Private key for decrypt
+        """
+        if os.path.exists(pub_key_path):
+            with open(pub_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.video_file_encrypt_pub_key_content = content
+
+        if os.path.exists(priv_key_path):
+            with open(priv_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.video_file_encrypt_priv_key_content = content
+                    self.video_file_encrypt_priv_key = \
+                        PrivateKey.from_private_pem(self.video_file_encrypt_priv_key_content)
+
+    def _generate_key_pair(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Generate RSA 4096 key pair and save to files
+
+        Args:
+            pub_key_path: Path to save public key
+            priv_key_path: Path to save private key
+
+        Returns:
+            Tuple of (public_key_path, private_key_path)
+        """
+        try:
+            # Generate RSA key pair
+            key = rsa.generate_private_key(
+                backend=crypto_default_backend(),
+                public_exponent=65537,
+                key_size=4096
+            )
+
+            # Save private key
+            private_key = key.private_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PrivateFormat.PKCS8,
+                crypto_serialization.NoEncryption()
+            )
+            with open(priv_key_path, 'wb') as f:
+                f.write(private_key)
+
+            # Save public key
+            public_key = key.public_key().public_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PublicFormat.SubjectPublicKeyInfo
+            )
+            with open(pub_key_path, 'wb') as f:
+                f.write(public_key)
+        except Exception as e:
+            logger.error(f"[Error] Failed to generate key pair: {e}")
+            raise AICCException("Failed to generate RSA key pair") from e
+
+    def generate_video_file_encrypt_key(self, pub_key_path: str = "", priv_key_path: str = ""):
+        """
+        Generate video file encrypt key
+        Args:
+            pub_key_path: Public key for encrypt
+            priv_key_path: Private key for decrypt
+        """
+        # 如果没有提供路径报错
+        if not pub_key_path or not priv_key_path:
+            raise ValueError("pub_key_path and priv_key_path must be provided")
+        
+        self._generate_key_pair(pub_key_path, priv_key_path)
+        # self.set_video_file_encrypt_key(pub_key_path, priv_key_path)
+
+    def create_video_generation_task(self, data: dict, custom_header: dict = None) -> str:
+        """
+        :param data:
+        :param custom_header:
+        :return:
+        """
+        if not data:
+            raise AICCException("data must be provided")
+
+        if "model" not in data:
+            data["model"] = self.endpoint
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.CREATE_TASK_URL}"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+        if custom_header:
+            headers.update(custom_header)
+
+        # Add encryption headers if provided
+        if self.is_enable_video_encrypt:
+            if not self.video_file_encrypt_pub_key_content or not self.video_file_encrypt_priv_key_content:
+                raise AICCException("pub_key and priv_key must be provided when Is-Enable-Encrypt is 'true'")
+
+            headers["Enable-TOS-Content-Result-Encryption"] = str(self.is_enable_video_encrypt).lower()
+            headers["X-Encryption-Algorithm"] = "RSA_OAEP_4096_AES_256"
+            if self.video_file_encrypt_pub_key_content:
+                # Base64 encode the PEM public key for safe HTTP header transmission
+                # This avoids issues with newlines and special characters in headers
+                pk_base64 = base64.b64encode(self.video_file_encrypt_pub_key_content.encode()).decode()
+                headers["PK"] = pk_base64
+
+        # logger.info(f"[Request] Data: {json.dumps(data, ensure_ascii=False, indent=2)}")
+
+        response = self.secure_http_client.post(url, headers=headers, json=data, timeout=self.timeout)
+        if response.status_code == 200:
+            result = response.json()
+            task_id = result.get("id")
+            return task_id
+        else:
+            logger.error(f"[Error] url: {url}, response: {response.text}")
+            return ""
+
+    def create_video_generation_task_by_message(self, message: list, custom_header: dict = None,
+                                                generate_audio: bool = True,
+                                                ratio: str = "16:9", duration: int = 11,
+                                                watermark: bool = False) -> str:
+        """
+           message: List of messages to generate video
+        """
+        if not message:
+            raise AICCException("message must be provided")
+
+        data = {
+            "model": self.endpoint,
+            "content": message,
+            "generate_audio": generate_audio,
+            "ratio": ratio,
+            "duration": duration,
+            "watermark": watermark
+        }
+        return self.create_video_generation_task(data, custom_header)
+
+    def query_video_generation_task(self, task_id: str) -> dict:
+        """
+        Query video generation task
+        API: GET /api/v3/contents/generations/tasks/{task_id}
+        """
+
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.QUERY_TASK_URL}/{task_id}"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+
+        # logger.info(f"[Request] URL: {url}")
+
+        # response = self.secure_http_client.get(url, headers=headers, timeout=self.timeout)
+        response = self.secure_http_client.request(
+            "GET", url, headers=headers, content=b"{}", timeout=self.timeout
+        )
+
+        if response.status_code == 200:
+            result = response.json()
+            # logger.info(f"[Response] Body: {json.dumps(result, ensure_ascii=False, indent=2)}")
+
+            # status = result.get("status")
+            # if status == "succeeded":
+            #     video_url = result.get("content", {}).get("video_url")
+            #     print(f"[Video URL] {video_url}")
+
+            return result
+        else:
+            try:
+                result = response.json()
+                return result
+            except Exception as e:
+                logger.error(f"[Error] url: {url}, Exception={e}")
+                raise AICCException(f"Failed to query video generation task")
+
+    def query_video_generation_task_list(
+            self,
+            page_num: int = 0,
+            page_size: int = 0,
+            status: str = "",
+            task_ids: list = None,
+            model: str = "",
+            service_tier: str = ""
+    ) -> dict:
+        """
+        Query video generation task list
+        API: GET /api/v3/contents/generations/tasks
+
+        Args:
+            page_num: Page number (only added if > 0)
+            page_size: Page size (only added if > 0)
+            status: Filter by status (only added if not empty)
+            task_ids: Filter by task IDs (only added if not empty list)
+            model: Filter by model (only added if not empty)
+            service_tier: Filter by service tier (only added if not empty)
+
+        Returns:
+            Query result dict
+        """
+        from urllib.parse import urlencode
+
+        base_url = self._get_base_url()
+
+        # Build query parameters (only add non-empty values)
+        params = []
+
+        if page_num > 0:
+            params.append(f"page_num={page_num}")
+
+        if page_size > 0:
+            params.append(f"page_size={page_size}")
+
+        if status:
+            params.append(f"filter.status={status}")
+
+        if task_ids and len(task_ids) > 0:
+            for task_id in task_ids:
+                params.append(f"filter.task_ids={task_id}")
+
+        if model:
+            params.append(f"filter.model={model}")
+
+        if service_tier:
+            params.append(f"filter.service_tier={service_tier}")
+
+        # Build URL with query parameters
+        query_string = "&".join(params)
+        if query_string:
+            url = f"{base_url}{self.QUERY_TASK_LIST_URL}?{query_string}"
+        else:
+            url = f"{base_url}{self.QUERY_TASK_LIST_URL}"
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+
+        response = self.secure_http_client.get(url, headers=headers, timeout=self.timeout)
+        try:
+            result = response.json()
+            return result
+        except Exception as e:
+            logger.error(f"query_video_generation_task_list error: {e}")
+            raise AICCException(f"Failed to query video generation task list")
+
+    def delete_video_generation_task(self, task_id: str):
+        """
+        Delete video generation task
+        API: DELETE /api/v3/contents/generations/tasks/{task_id}
+        """
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.DELETE_TASK_URL}/{task_id}"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+        response = self.secure_http_client.delete(url, headers=headers, timeout=self.timeout)
+        try:
+            result = response.json()
+            return result
+        except Exception as e:
+            logger.error(f"delete_video_generation_task error: {e}")
+            raise AICCException(f"Failed to delete_video_generation_task")
+
+    def video_file_download(self, task_id: str, video_file_path: str) -> bool:
+        """
+        Download and decrypt encrypted video
+        This function polls the task status until it succeeds, then downloads
+        and decrypts the video using the provided private key.
+        Args:
+            task_id: Video generation task ID
+            video_file_path: Directory to save decrypted video
+
+        Returns:
+            Path to decrypted video file, or None if failed
+        """
+        if not task_id:
+            raise AICCException("task_id must be provided")
+
+        if not video_file_path:
+            raise AICCException("video_file_path must be provided")
+
+        result = self.query_video_generation_task(task_id)
+        if not result:
+            logger.error(f"[Error] Failed to query task,task_id={task_id}")
+            return False
+
+        status = result.get("status")
+        if status == "succeeded":
+            video_url = result.get("content", {}).get("video_url")
+            if not video_url:
+                logger.error("[Error] No video_url in response")
+                return False
+
+            return self.download_file_normal(video_url, video_file_path)
+            # encrypted_key = b""
+            # if self.is_enable_video_encrypt:
+            #     encrypted_key = result.get("encrypted_key")
+            #     if not encrypted_key:
+            #         logger.error("[Error] No encrypted_key in response")
+            #         return False
+            #     if not self.video_file_encrypt_priv_key:
+            #         logger.error("[Error] No private key set, cannot decrypt encrypted video")
+            #         return False
+            #     else:
+            #         # Base64 decode the encrypted key before RSA decryption
+            #         encrypted_key_bytes = base64.b64decode(encrypted_key)
+            #         encrypted_key = self.video_file_encrypt_priv_key.decrypt(encrypted_key_bytes)
+            #
+            # logger.info(f"\n[Downloading] Starting download and decryption...")
+            #
+            # try:
+            #     if self.video_file_storage_location == "TOS":
+            #         self.tos_utils.download_file(tos_url=video_url, local_file_path=video_file_path,
+            #                                       encrypted_key=encrypted_key)
+            #     else:
+            #         logger.error(f"[Error] Unsupported video storage location: {self.video_file_storage_location}")
+            #         return False
+            #
+            #     logger.info(f"\n[Success] Decrypted video saved to: {video_file_path}")
+            #     return True
+            #
+            # except Exception as e:
+            #     logger.error(f"[Error] Download and decryption failed: {str(e)}")
+            #     return False
+
+        elif status == "failed":
+            error_info = result.get("error", "Unknown error")
+            logger.error(f"[Error] Task failed: {error_info}")
+            return False
+
+        elif status in ["pending", "running", "queued"]:
+            logger.info(f"[Waiting] Task is {status}, please wait for it to finish...")
+            return False
+        else:
+            logger.warning(f"[Warning] Unknown status: {status}")
+            return False
+
+    def video_file_download_by_url(self, video_url: str, video_file_path: str) -> bool:
+        """
+        Download and decrypt encrypted video
+        This function polls the task status until it succeeds, then downloads
+        and decrypts the video using the provided private key.
+        Args:
+            video_url: Video video_url
+            video_file_path: Directory to save decrypted video
+
+        Returns:
+            Path to decrypted video file, or None if failed
+        """
+        if not video_url:
+            raise AICCException("video_url must be provided")
+
+        if not video_file_path:
+            raise AICCException("video_file_path must be provided")
+
+        logger.info(f"\n[Downloading] Starting download and decryption...")
+
+        try:
+            return self.download_file_normal(video_url, video_file_path)
+        except Exception as e:
+            logger.error(f"[Error] Download and decryption failed: {str(e)}")
+            return False
+
+    def download_file_normal(self, video_url: str, video_file_path: str) -> bool:
+        """
+        Download video file
+        Args:
+            video_url: Video video_url
+            video_file_path: Directory to save video file
+
+        Returns:
+            Path to video file, or None if failed
+        """
+        if not video_url:
+            raise AICCException("video_url must be provided")
+
+        if not video_file_path:
+            raise AICCException("video_file_path must be provided")
+
+        if os.path.exists(video_file_path) and os.path.getsize(video_file_path) > 0:
+            logger.error(f"[Error] File {video_file_path} already exists and is not empty")
+            return False
+
+        try:
+            response = requests.get(video_url, stream=True)
+            with open(video_file_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            if self.is_enable_video_encrypt:
+                if not self.video_file_encrypt_priv_key:
+                    logger.error("[Error] No private key set, cannot decrypt encrypted video")
+                    return False
+
+                headers = response.headers
+                encrypted_key = headers.get("x-tos-meta-enc-dek")
+                # nonce = headers.get("x-tos-meta-x-tos-meta-base-nonce")
+                if not encrypted_key:
+                    logger.error("[Error] No encrypted_key in response")
+                    return False
+
+                try:
+                    encrypted_key_bytes = base64.b64decode(unquote(encrypted_key))
+                    encrypted_key = self.video_file_encrypt_priv_key.decrypt(encrypted_key_bytes)
+                    tmp_file_path = video_file_path + ".tmp"
+                    decrypt_file(encrypted_key, video_file_path, tmp_file_path, 'b')
+                    os.remove(video_file_path)
+                    os.rename(tmp_file_path, video_file_path)
+                except Exception as e:
+                    logger.error(f"[Error] Failed to decode encrypted_key: {e}")
+                    return False
+        except Exception as e:
+            logger.error(f"[Error] Download failed: {e}")
+            traceback.print_exc()
+            return False
+        return True

src/bytedance/volcengine_aicc_sdk/seedream_inference_client.py

@@ -0,0 +1,213 @@
+"""
+seedream机密推理客户端
+"""
+__all__ = ["SeedreamClient"]
+
+import os
+import uuid
+import base64
+import traceback
+from urllib.parse import unquote
+from cryptography.hazmat.primitives import serialization as crypto_serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.backends import default_backend as crypto_default_backend
+from bytedance.jeddak_secure_channel.crypto import PrivateKey
+from .secure_http_client import SecureHttpClient
+from .common.error import AICCException
+from .common.log import logger
+from .secure_http_client import AICCConfig
+from .common.aes import decrypt_file
+
+
+class SeedreamClient:
+
+    def __init__(self,
+                 base_url: str = "",
+                 api_key: str = "",
+                 endpoint: str = "",
+                 aicc_config: AICCConfig = None,
+                 **kwargs):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.endpoint = endpoint
+        self.aicc_config = aicc_config
+        self.kwargs = kwargs
+        self.is_secure = self.kwargs.get("is_secure", True)
+
+        if aicc_config:
+            self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+        else:
+            if self.is_secure:
+                aicc_config = AICCConfig(ra_url=base_url, auth_token=api_key)
+                self.secure_http_client = SecureHttpClient(aicc_config=aicc_config)
+            else:
+                self.secure_http_client = SecureHttpClient()
+
+        self.is_enable_image_encrypt = str(self.kwargs.get("is_enable_image_encrypt", True)).lower() == "true"
+
+        self.image_file_encrypt_pub_key_content = None
+        self.image_file_encrypt_priv_key_content = None
+
+        self.image_file_encrypt_priv_key: PrivateKey = None
+
+        self.timeout = self.kwargs.get("timeout", 120.0)
+
+        self.CREATE_IMAGE_URL = "/images/generations"
+
+    def _get_base_url(self):
+        return self.base_url.rstrip("/")
+
+    def _build_common_headers(self, request_id: str = "") -> dict[str, str]:
+        if not request_id:
+            request_id = str(uuid.uuid4())
+
+        return {
+            "Content-Type": "application/json",
+            "X-Api-App-Key": self.endpoint,
+            "Authorization": f"Bearer {self.api_key}",
+            "X-Request-Id": request_id,
+        }
+
+    def generate_key_pair(self, pub_key_path: str = "", priv_key_path: str = ""):
+        try:
+            key = rsa.generate_private_key(
+                backend=crypto_default_backend(),
+                public_exponent=65537,
+                key_size=4096
+            )
+
+            private_key = key.private_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PrivateFormat.PKCS8,
+                crypto_serialization.NoEncryption()
+            )
+            with open(priv_key_path, 'wb') as f:
+                f.write(private_key)
+
+            public_key = key.public_key().public_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PublicFormat.SubjectPublicKeyInfo
+            )
+            with open(pub_key_path, 'wb') as f:
+                f.write(public_key)
+        except Exception as e:
+            logger.error(f"[Error] Failed to generate key pair: {e}")
+            raise AICCException("Failed to generate RSA key pair") from e
+
+    def set_image_encrypt_key(self, pub_key_path: str = "", priv_key_path: str = ""):
+        if os.path.exists(pub_key_path):
+            with open(pub_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.image_file_encrypt_pub_key_content = content
+
+        if os.path.exists(priv_key_path):
+            with open(priv_key_path, 'r') as f:
+                content = f.read()
+                if '-----BEGIN' in content and '-----END' in content:
+                    self.image_file_encrypt_priv_key_content = content
+                    self.image_file_encrypt_priv_key = \
+                        PrivateKey.from_private_pem(self.image_file_encrypt_priv_key_content)
+
+    def create_image(self, data: dict, custom_header: dict = None, request_id: str = ""):
+        if not data:
+            raise AICCException("data must be provided")
+
+        if "model" not in data:
+            data["model"] = self.endpoint
+
+        base_url = self._get_base_url()
+        url = f"{base_url}{self.CREATE_IMAGE_URL}"
+
+        headers = self._build_common_headers(request_id)
+        if custom_header:
+            headers.update(custom_header)
+
+        if self.is_enable_image_encrypt:
+            if not self.image_file_encrypt_pub_key_content or not self.image_file_encrypt_priv_key_content:
+                raise AICCException("pub_key and priv_key must be provided when Is-Enable-Encrypt is 'true'")
+
+            headers["Enable-TOS-Content-Result-Encryption"] = str(self.is_enable_image_encrypt).lower()
+            headers["X-Encryption-Algorithm"] = "RSA_OAEP_4096_AES_256"
+            if self.image_file_encrypt_pub_key_content:
+                pk_base64 = base64.b64encode(self.image_file_encrypt_pub_key_content.encode()).decode()
+                headers["PK"] = pk_base64
+
+        stream = data.get("stream", False)
+        if stream:
+            return self.secure_http_client.send(
+                self.secure_http_client.build_request("POST", url, headers=headers, json=data, timeout=self.timeout),
+                stream=True,
+            )
+
+        response = self.secure_http_client.post(url, headers=headers, json=data, timeout=self.timeout)
+        if response.status_code == 200:
+            return response.json()
+        else:
+            logger.error(f"[Error] url: {url}, response: {response.text}")
+            try:
+                return response.json()
+            except Exception:
+                raise AICCException(f"Failed to create image, status_code={response.status_code}")
+
+    def create_image_by_prompt(self, prompt: str, custom_header: dict = None, request_id: str = "",
+                               size: str = "1024x1024", n: int = 1, response_format: str = "url",
+                               stream: bool = False, extra_body: dict = None):
+        if not prompt:
+            raise AICCException("prompt must be provided")
+
+        data = {
+            "model": self.endpoint,
+            "prompt": prompt,
+            "size": size,
+            "n": n,
+            "response_format": response_format,
+        }
+        if stream:
+            data["stream"] = True
+        if extra_body:
+            data.update(extra_body)
+        return self.create_image(data, custom_header, request_id)
+
+    def download_image(self, image_url: str, local_file_path: str) -> bool:
+        if not image_url:
+            raise AICCException("image_url must be provided")
+
+        if not local_file_path:
+            raise AICCException("local_file_path must be provided")
+
+        if os.path.exists(local_file_path) and os.path.getsize(local_file_path) > 0:
+            logger.error(f"[Error] File {local_file_path} already exists and is not empty")
+            return False
+
+        try:
+            import requests as req
+            response = req.get(image_url, stream=True)
+            with open(local_file_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            if self.is_enable_image_encrypt:
+                if not self.image_file_encrypt_priv_key:
+                    logger.error("[Error] No private key set, cannot decrypt encrypted image")
+                    return False
+
+                headers = response.headers
+                encrypted_key = headers.get("x-tos-meta-enc-dek")
+                if encrypted_key:
+                    try:
+                        encrypted_key_bytes = base64.b64decode(unquote(encrypted_key))
+                        encrypted_key = self.image_file_encrypt_priv_key.decrypt(encrypted_key_bytes)
+                        tmp_file_path = local_file_path + ".tmp"
+                        decrypt_file(encrypted_key, local_file_path, tmp_file_path, 'b')
+                        os.remove(local_file_path)
+                        os.rename(tmp_file_path, local_file_path)
+                    except Exception as e:
+                        logger.error(f"[Error] Failed to decode encrypted_key: {e}")
+                        traceback.print_exc()
+                        return False
+        except Exception as e:
+            logger.error(f"[Error] Download failed: {e}")
+            traceback.print_exc()
+            return False
+        return True