volcengine-python-sdk-aicc 0.0.41__py3-none-any.whl

bytedance/__init__.py

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

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}")

bytedance/jeddak_secure_channel/client.py

@@ -0,0 +1,305 @@
+"""
+数据发送方接口.
+"""
+
+__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)

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)

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  # 超过一定阈值被认为是不活跃的客户端，单位秒

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

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