From 6aa44199be04bedfe578711dc1c771522c1efcbc Mon Sep 17 00:00:00 2001 From: chakcy <947105045@qq.com> Date: Fri, 31 Jul 2026 16:55:35 +0800 Subject: [PATCH] first commit --- .gitignore | 72 ++ README.md | 0 pyproject.toml | 21 + src/bytedance/__init__.py | 4 + .../jeddak_secure_channel/__init__.py | 28 + src/bytedance/jeddak_secure_channel/client.py | 300 +++++++ src/bytedance/jeddak_secure_channel/config.py | 171 ++++ src/bytedance/jeddak_secure_channel/consts.py | 5 + .../jeddak_secure_channel/crypto/__init__.py | 14 + .../jeddak_secure_channel/crypto/aes.py | 174 ++++ .../jeddak_secure_channel/crypto/envelope.py | 395 +++++++++ .../jeddak_secure_channel/crypto/rsa.py | 164 ++++ src/bytedance/jeddak_secure_channel/error.py | 131 +++ .../jeddak_secure_channel/log/__init__.py | 24 + .../jeddak_secure_channel/log/config.py | 31 + .../jeddak_secure_channel/log/logger.py | 88 ++ .../jeddak_secure_channel/logger_factory.py | 134 +++ src/bytedance/jeddak_secure_channel/ra.py | 768 ++++++++++++++++++ src/bytedance/jeddak_secure_channel/server.py | 734 +++++++++++++++++ .../jeddak_secure_channel/service/__init__.py | 16 + .../jeddak_secure_channel/service/eps.py | 133 +++ .../service/eps_models.py | 47 ++ .../jeddak_secure_channel/service/ras.py | 88 ++ .../jeddak_secure_channel/service/tks.py | 51 ++ .../service/tks_models.py | 334 ++++++++ src/bytedance/jeddak_secure_channel/utils.py | 469 +++++++++++ src/bytedance/tks/attester.py | 41 + src/bytedance/tks/client.py | 628 ++++++++++++++ src/bytedance/tks/consts.py | 122 +++ src/bytedance/tks/crypto/aes.py | 47 ++ src/bytedance/tks/crypto/ec.py | 39 + src/bytedance/tks/exceptions.py | 22 + src/bytedance/volcengine_aicc_sdk/__init__.py | 14 + .../volcengine_aicc_sdk/aicc_client.py | 225 +++++ .../volcengine_aicc_sdk/aicc_inference.py | 64 ++ .../ark_inference_client.py | 343 ++++++++ .../asr_inference_client.py | 158 ++++ .../asr_websocket_client.py | 647 +++++++++++++++ .../auc_inference_client.py | 305 +++++++ .../volcengine_aicc_sdk/common/__init__.py | 0 .../volcengine_aicc_sdk/common/aes.py | 136 ++++ .../volcengine_aicc_sdk/common/constant.py | 4 + .../volcengine_aicc_sdk/common/error.py | 21 + .../common/log/__init__.py | 24 + .../volcengine_aicc_sdk/common/log/config.py | 31 + .../volcengine_aicc_sdk/common/log/logger.py | 88 ++ .../volcengine_aicc_sdk/common/tos_utils.py | 341 ++++++++ .../volcengine_aicc_sdk/common/utils.py | 45 + .../cu_inference_client.py | 231 ++++++ .../deepseek_inference_client.py | 501 ++++++++++++ .../volcengine_aicc_sdk/secure_http_client.py | 297 +++++++ .../seedance_inference_client.py | 500 ++++++++++++ .../seedream_inference_client.py | 213 +++++ 53 files changed, 9483 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 src/bytedance/__init__.py create mode 100644 src/bytedance/jeddak_secure_channel/__init__.py create mode 100644 src/bytedance/jeddak_secure_channel/client.py create mode 100644 src/bytedance/jeddak_secure_channel/config.py create mode 100644 src/bytedance/jeddak_secure_channel/consts.py create mode 100644 src/bytedance/jeddak_secure_channel/crypto/__init__.py create mode 100644 src/bytedance/jeddak_secure_channel/crypto/aes.py create mode 100644 src/bytedance/jeddak_secure_channel/crypto/envelope.py create mode 100644 src/bytedance/jeddak_secure_channel/crypto/rsa.py create mode 100644 src/bytedance/jeddak_secure_channel/error.py create mode 100644 src/bytedance/jeddak_secure_channel/log/__init__.py create mode 100644 src/bytedance/jeddak_secure_channel/log/config.py create mode 100644 src/bytedance/jeddak_secure_channel/log/logger.py create mode 100644 src/bytedance/jeddak_secure_channel/logger_factory.py create mode 100644 src/bytedance/jeddak_secure_channel/ra.py create mode 100644 src/bytedance/jeddak_secure_channel/server.py create mode 100644 src/bytedance/jeddak_secure_channel/service/__init__.py create mode 100644 src/bytedance/jeddak_secure_channel/service/eps.py create mode 100644 src/bytedance/jeddak_secure_channel/service/eps_models.py create mode 100644 src/bytedance/jeddak_secure_channel/service/ras.py create mode 100644 src/bytedance/jeddak_secure_channel/service/tks.py create mode 100644 src/bytedance/jeddak_secure_channel/service/tks_models.py create mode 100644 src/bytedance/jeddak_secure_channel/utils.py create mode 100644 src/bytedance/tks/attester.py create mode 100644 src/bytedance/tks/client.py create mode 100644 src/bytedance/tks/consts.py create mode 100644 src/bytedance/tks/crypto/aes.py create mode 100644 src/bytedance/tks/crypto/ec.py create mode 100644 src/bytedance/tks/exceptions.py create mode 100644 src/bytedance/volcengine_aicc_sdk/__init__.py create mode 100644 src/bytedance/volcengine_aicc_sdk/aicc_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/aicc_inference.py create mode 100644 src/bytedance/volcengine_aicc_sdk/ark_inference_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/asr_inference_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/asr_websocket_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/auc_inference_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/__init__.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/aes.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/constant.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/error.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/log/__init__.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/log/config.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/log/logger.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/tos_utils.py create mode 100644 src/bytedance/volcengine_aicc_sdk/common/utils.py create mode 100644 src/bytedance/volcengine_aicc_sdk/cu_inference_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/deepseek_inference_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/secure_http_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/seedance_inference_client.py create mode 100644 src/bytedance/volcengine_aicc_sdk/seedream_inference_client.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c8f0442 --- /dev/null +++ b/.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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..e69de29 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..c666818 --- /dev/null +++ b/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"] diff --git a/src/bytedance/__init__.py b/src/bytedance/__init__.py new file mode 100644 index 0000000..ef0f915 --- /dev/null +++ b/src/bytedance/__init__.py @@ -0,0 +1,4 @@ +from pkgutil import extend_path + +# Make "bytedance" a namespace package +__path__ = extend_path(__path__, __name__) diff --git a/src/bytedance/jeddak_secure_channel/__init__.py b/src/bytedance/jeddak_secure_channel/__init__.py new file mode 100644 index 0000000..e01f789 --- /dev/null +++ b/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}") diff --git a/src/bytedance/jeddak_secure_channel/client.py b/src/bytedance/jeddak_secure_channel/client.py new file mode 100644 index 0000000..8e1bda7 --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/config.py b/src/bytedance/jeddak_secure_channel/config.py new file mode 100644 index 0000000..20e1863 --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/consts.py b/src/bytedance/jeddak_secure_channel/consts.py new file mode 100644 index 0000000..4b3e84f --- /dev/null +++ b/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 # 超过一定阈值被认为是不活跃的客户端，单位秒 diff --git a/src/bytedance/jeddak_secure_channel/crypto/__init__.py b/src/bytedance/jeddak_secure_channel/crypto/__init__.py new file mode 100644 index 0000000..9bdb314 --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/crypto/aes.py b/src/bytedance/jeddak_secure_channel/crypto/aes.py new file mode 100644 index 0000000..dcf97cb --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/crypto/envelope.py b/src/bytedance/jeddak_secure_channel/crypto/envelope.py new file mode 100644 index 0000000..137e10c --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/crypto/rsa.py b/src/bytedance/jeddak_secure_channel/crypto/rsa.py new file mode 100644 index 0000000..16e4b2d --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/error.py b/src/bytedance/jeddak_secure_channel/error.py new file mode 100644 index 0000000..cd8d2ee --- /dev/null +++ b/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" diff --git a/src/bytedance/jeddak_secure_channel/log/__init__.py b/src/bytedance/jeddak_secure_channel/log/__init__.py new file mode 100644 index 0000000..e8b1587 --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/log/config.py b/src/bytedance/jeddak_secure_channel/log/config.py new file mode 100644 index 0000000..86c10bb --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/log/logger.py b/src/bytedance/jeddak_secure_channel/log/logger.py new file mode 100644 index 0000000..64b9273 --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/logger_factory.py b/src/bytedance/jeddak_secure_channel/logger_factory.py new file mode 100644 index 0000000..6994f0f --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/ra.py b/src/bytedance/jeddak_secure_channel/ra.py new file mode 100644 index 0000000..1cd818e --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/server.py b/src/bytedance/jeddak_secure_channel/server.py new file mode 100644 index 0000000..1b90d12 --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/service/__init__.py b/src/bytedance/jeddak_secure_channel/service/__init__.py new file mode 100644 index 0000000..689f738 --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/service/eps.py b/src/bytedance/jeddak_secure_channel/service/eps.py new file mode 100644 index 0000000..965ec68 --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/service/eps_models.py b/src/bytedance/jeddak_secure_channel/service/eps_models.py new file mode 100644 index 0000000..a671bb1 --- /dev/null +++ b/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 + """明文的运行时度量事件日志.""" diff --git a/src/bytedance/jeddak_secure_channel/service/ras.py b/src/bytedance/jeddak_secure_channel/service/ras.py new file mode 100644 index 0000000..86d0c0c --- /dev/null +++ b/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 diff --git a/src/bytedance/jeddak_secure_channel/service/tks.py b/src/bytedance/jeddak_secure_channel/service/tks.py new file mode 100644 index 0000000..4ca0a65 --- /dev/null +++ b/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) diff --git a/src/bytedance/jeddak_secure_channel/service/tks_models.py b/src/bytedance/jeddak_secure_channel/service/tks_models.py new file mode 100644 index 0000000..939c083 --- /dev/null +++ b/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 时使用.""" diff --git a/src/bytedance/jeddak_secure_channel/utils.py b/src/bytedance/jeddak_secure_channel/utils.py new file mode 100644 index 0000000..f497cf9 --- /dev/null +++ b/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 diff --git a/src/bytedance/tks/attester.py b/src/bytedance/tks/attester.py new file mode 100644 index 0000000..f208e98 --- /dev/null +++ b/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 diff --git a/src/bytedance/tks/client.py b/src/bytedance/tks/client.py new file mode 100644 index 0000000..9689f8f --- /dev/null +++ b/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" + # } + # ] + # ]) + # )) diff --git a/src/bytedance/tks/consts.py b/src/bytedance/tks/consts.py new file mode 100644 index 0000000..b734927 --- /dev/null +++ b/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() diff --git a/src/bytedance/tks/crypto/aes.py b/src/bytedance/tks/crypto/aes.py new file mode 100644 index 0000000..072e4fe --- /dev/null +++ b/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 diff --git a/src/bytedance/tks/crypto/ec.py b/src/bytedance/tks/crypto/ec.py new file mode 100644 index 0000000..9585d6c --- /dev/null +++ b/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 diff --git a/src/bytedance/tks/exceptions.py b/src/bytedance/tks/exceptions.py new file mode 100644 index 0000000..c1293b9 --- /dev/null +++ b/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) diff --git a/src/bytedance/volcengine_aicc_sdk/__init__.py b/src/bytedance/volcengine_aicc_sdk/__init__.py new file mode 100644 index 0000000..1a228ff --- /dev/null +++ b/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}") diff --git a/src/bytedance/volcengine_aicc_sdk/aicc_client.py b/src/bytedance/volcengine_aicc_sdk/aicc_client.py new file mode 100644 index 0000000..c441568 --- /dev/null +++ b/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) diff --git a/src/bytedance/volcengine_aicc_sdk/aicc_inference.py b/src/bytedance/volcengine_aicc_sdk/aicc_inference.py new file mode 100644 index 0000000..c8af808 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/ark_inference_client.py b/src/bytedance/volcengine_aicc_sdk/ark_inference_client.py new file mode 100644 index 0000000..3feed2e --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/asr_inference_client.py b/src/bytedance/volcengine_aicc_sdk/asr_inference_client.py new file mode 100644 index 0000000..6d8e8cf --- /dev/null +++ b/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") diff --git a/src/bytedance/volcengine_aicc_sdk/asr_websocket_client.py b/src/bytedance/volcengine_aicc_sdk/asr_websocket_client.py new file mode 100644 index 0000000..42de538 --- /dev/null +++ b/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() diff --git a/src/bytedance/volcengine_aicc_sdk/auc_inference_client.py b/src/bytedance/volcengine_aicc_sdk/auc_inference_client.py new file mode 100644 index 0000000..dd764f7 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/common/__init__.py b/src/bytedance/volcengine_aicc_sdk/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/bytedance/volcengine_aicc_sdk/common/aes.py b/src/bytedance/volcengine_aicc_sdk/common/aes.py new file mode 100644 index 0000000..c4593fa --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/common/constant.py b/src/bytedance/volcengine_aicc_sdk/common/constant.py new file mode 100644 index 0000000..cca6e47 --- /dev/null +++ b/src/bytedance/volcengine_aicc_sdk/common/constant.py @@ -0,0 +1,4 @@ +# -*- coding: utf-8 -*- + +class Constants: + LOG_CONFIG = "{\"dir\": \"aicc_log\"}" diff --git a/src/bytedance/volcengine_aicc_sdk/common/error.py b/src/bytedance/volcengine_aicc_sdk/common/error.py new file mode 100644 index 0000000..8e13cff --- /dev/null +++ b/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() diff --git a/src/bytedance/volcengine_aicc_sdk/common/log/__init__.py b/src/bytedance/volcengine_aicc_sdk/common/log/__init__.py new file mode 100644 index 0000000..e8b1587 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/common/log/config.py b/src/bytedance/volcengine_aicc_sdk/common/log/config.py new file mode 100644 index 0000000..58224d2 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/common/log/logger.py b/src/bytedance/volcengine_aicc_sdk/common/log/logger.py new file mode 100644 index 0000000..474efea --- /dev/null +++ b/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) diff --git a/src/bytedance/volcengine_aicc_sdk/common/tos_utils.py b/src/bytedance/volcengine_aicc_sdk/common/tos_utils.py new file mode 100644 index 0000000..6c8fc58 --- /dev/null +++ b/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") diff --git a/src/bytedance/volcengine_aicc_sdk/common/utils.py b/src/bytedance/volcengine_aicc_sdk/common/utils.py new file mode 100644 index 0000000..48aa958 --- /dev/null +++ b/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() + ) diff --git a/src/bytedance/volcengine_aicc_sdk/cu_inference_client.py b/src/bytedance/volcengine_aicc_sdk/cu_inference_client.py new file mode 100644 index 0000000..bc34ff5 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/deepseek_inference_client.py b/src/bytedance/volcengine_aicc_sdk/deepseek_inference_client.py new file mode 100644 index 0000000..f25a547 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/secure_http_client.py b/src/bytedance/volcengine_aicc_sdk/secure_http_client.py new file mode 100644 index 0000000..8a9ba42 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/seedance_inference_client.py b/src/bytedance/volcengine_aicc_sdk/seedance_inference_client.py new file mode 100644 index 0000000..40c8a46 --- /dev/null +++ b/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 diff --git a/src/bytedance/volcengine_aicc_sdk/seedream_inference_client.py b/src/bytedance/volcengine_aicc_sdk/seedream_inference_client.py new file mode 100644 index 0000000..a915550 --- /dev/null +++ b/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