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

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)

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

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"

bytedance/jeddak_secure_channel/log/__init__.py

@@ -0,0 +1,23 @@
+__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