Merge branch 'feat/aicc_case' into 'integration_2025-10-16_1070862462978'

feat: [development task] ark runtime (1736563)

See merge request iaasng/volcengine-python-sdk!854

Author: BitsAdmin

volcenginesdkark/models/get_endpoint_certificate_request.py

@@ -33,23 +33,28 @@ class GetEndpointCertificateRequest(object):
                             and the value is json key in definition.
     """
     swagger_types = {
-        'id': 'str'
+        'id': 'str',
+        'type': 'str'
     }
 
     attribute_map = {
-        'id': 'Id'
+        'id': 'Id',
+        'type': 'Type'
     }
 
-    def __init__(self, id=None, _configuration=None):  # noqa: E501
+    def __init__(self, id=None, type=None, _configuration=None):  # noqa: E501
         """GetEndpointCertificateRequest - a model defined in Swagger"""  # noqa: E501
         if _configuration is None:
             _configuration = Configuration()
         self._configuration = _configuration
 
         self._id = None
+        self._type = None
         self.discriminator = None
 
         self.id = id
+        if type is not None:
+            self.type = type
 
     @property
     def id(self):
@@ -74,6 +79,27 @@ def id(self, id):
 
         self._id = id
 
+    @property
+    def type(self):
+        """Gets the type of this GetEndpointCertificateRequest.  # noqa: E501
+
+
+        :return: The type of this GetEndpointCertificateRequest.  # noqa: E501
+        :rtype: str
+        """
+        return self._type
+
+    @type.setter
+    def type(self, type):
+        """Sets the type of this GetEndpointCertificateRequest.
+
+
+        :param type: The type of this GetEndpointCertificateRequest.  # noqa: E501
+        :type: str
+        """
+
+        self._type = type
+
     def to_dict(self):
         """Returns the model properties as a dict"""
         result = {}

volcenginesdkarkruntime/_client.py

@@ -42,7 +42,7 @@
 )
 from ._streaming import Stream
 
-from ._utils._key_agreement import key_agreement_client
+from ._utils._key_agreement import key_agreement_client, get_cert_info
 from ._utils._model_breaker import ModelBreaker
 
 __all__ = ["Ark", "AsyncArk"]
@@ -143,7 +143,7 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
             self._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(endpoint_id)
 
-    def _get_endpoint_certificate(self, endpoint_id: str) -> key_agreement_client:
+    def _get_endpoint_certificate(self, endpoint_id: str) -> tuple[key_agreement_client, str, str]:
         if self._certificate_manager is None:
             cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
             if (
@@ -279,7 +279,7 @@ def _get_bot_sts_token(self, bot_id: str):
             self._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(bot_id, resource_type="bot")
 
-    def _get_endpoint_certificate(self, endpoint_id: str) -> key_agreement_client:
+    def _get_endpoint_certificate(self, endpoint_id: str) -> tuple[key_agreement_client, str, str]:
         if self._certificate_manager is None:
             cert_path = os.environ.get("E2E_CERTIFICATE_PATH")
             if (
@@ -429,7 +429,7 @@ def __init__(
         base_url: str | URL = BASE_URL,
         api_key: str | None = None,
     ):
-        self._certificate_manager: Dict[str, key_agreement_client] = {}
+        self._certificate_manager: Dict[str, Tuple[key_agreement_client, str, str]] = {}
 
         # local cache prepare
         self._init_local_cert_cache()
@@ -460,6 +460,10 @@ def __init__(
         self._e2e_uri = "/e2e/get/certificate"
         self._x_session_token = {"X-Session-Token": self._e2e_uri}
 
+        self._aicc_enabled = False
+        if os.environ.get("VOLC_ARK_ENCRYPTION") == "AICC":
+            self._aicc_enabled = True
+
     def _load_cert_by_cert_path(self) -> str:
         with open(self.cert_path, "r") as f:
             cert_pem = f.read()
@@ -469,6 +473,10 @@ def _load_cert_by_ak_sk(self, ep: str) -> str:
         get_endpoint_certificate_request = (
             volcenginesdkark.GetEndpointCertificateRequest(id=ep)
         )
+        if self._aicc_enabled:
+            get_endpoint_certificate_request = (
+                volcenginesdkark.GetEndpointCertificateRequest(id=ep, type="AICCv0.1")
+            )
         try:
             resp: volcenginesdkark.GetEndpointCertificateResponse = (
                 self.api_instance.get_endpoint_certificate(
@@ -484,10 +492,13 @@ def _load_cert_by_ak_sk(self, ep: str) -> str:
 
     def _sync_load_cert_by_auth(self, ep: str) -> str:
         try:  # try to make request with session header (used for header statistic)
+            req_body = {"model": ep}
+            if self._aicc_enabled:
+                req_body["type"] = "AICCv0.1"
             resp = self.client.post(
                 self._e2e_uri,
                 options={"headers": self._x_session_token},
-                body={"model": ep},
+                body=req_body,
                 cast_to=CertificateResponse,
             )
         except Exception as e:
@@ -508,10 +519,16 @@ def _load_cert_locally(self, ep: str) -> str | None:
             current_time = time.time()
             time_difference = current_time - last_modified_time
             if time_difference <= self._cert_expiration_seconds:
+                cert_pem = None
                 with open(cert_file_path, "r") as f:
-                    return f.read()
-            else:
-                os.remove(cert_file_path)
+                    cert_pem = f.read()
+                ring, key = get_cert_info(cert_pem)
+                # check cert is complement with AICC/PCA
+                if (ring == "" or key == "") and not self._aicc_enabled:
+                    return cert_pem
+                if ring != "" and key != "" and self._aicc_enabled:
+                    return cert_pem
+            os.remove(cert_file_path)
         return None
 
     def _init_local_cert_cache(self):
@@ -528,7 +545,7 @@ def _init_local_cert_cache(self):
                     "failed to create certificate directory %s: %s\n" % (self._cert_storage_path, e)
                 )
 
-    def get(self, ep: str) -> key_agreement_client:
+    def get(self, ep: str) -> tuple[key_agreement_client, str, str]:
         if ep not in self._certificate_manager:
             cert_pem = self._load_cert_locally(ep)
             if cert_pem is None:
@@ -539,7 +556,12 @@ def get(self, ep: str) -> key_agreement_client:
                 else:
                     cert_pem = self._load_cert_by_ak_sk(ep)
                 self._save_cert_to_file(ep, cert_pem)
-            self._certificate_manager[ep] = key_agreement_client(
-                certificate_pem_string=cert_pem
+            ring, key = get_cert_info(cert_pem)
+            self._certificate_manager[ep] = (
+                key_agreement_client(
+                    certificate_pem_string=cert_pem
+                ),
+                ring,
+                key,
             )
         return self._certificate_manager[ep]

volcenginesdkarkruntime/_utils/_key_agreement.py

@@ -17,6 +17,22 @@
 from typing import Tuple
 
 
+def get_cert_info(cert_pem: str) -> Tuple[str, str]:
+    import re
+    from cryptography import x509
+    from cryptography.hazmat.backends import default_backend
+
+    cert = x509.load_pem_x509_certificate(cert_pem.encode(), default_backend())
+    try:
+        dns = cert.extensions.get_extension_for_class(
+            x509.SubjectAlternativeName).value.get_values_for_type(x509.DNSName)
+        if dns and len(dns) > 1 and re.match(r"^ring\..*$", dns[0]) and re.match(r"^key\..*$", dns[1]):
+            return dns[0][5:], dns[1][4:]
+    except Exception:
+        pass
+    return "", ""
+
+
 def aes_gcm_encrypt_bytes(
     key: bytes, iv: bytes, plain_bytes: bytes, associated_data: bytes = b""
 ) -> bytes:
@@ -86,15 +102,24 @@ def aes_gcm_decrypt_base64_list(key: bytes, nonce: bytes, ciphertext: str) -> st
         except Exception:
             for i in range(20, len(b64), 4):
                 try:
-                    decrypted = aes_gcm_decrypt_base64_string(key, nonce, b64[:i+4])
+                    decrypted = aes_gcm_decrypt_base64_string(
+                        key, nonce, b64[:i+4])
                     result.append(decrypted)
-                    decrypted = aes_gcm_decrypt_base64_string(key, nonce, b64[i+4:])
+                    decrypted = aes_gcm_decrypt_base64_string(
+                        key, nonce, b64[i+4:])
                     result.append(decrypted)
+                    break
                 except Exception:
-                    result.append('')
+                    pass
     return ''.join(result)
 
 
+def decrypt_validate(ciphertext: str) -> bool:
+    cipher_bytes = ciphertext.encode()
+    cipher_b64_bytes = base64.decodebytes(cipher_bytes)
+    return len(cipher_bytes)/4 >= len(cipher_b64_bytes)/3 >= len(cipher_bytes)/4 - 1
+
+
 def marshal_cryptography_pub_key(key) -> bytes:
     # python version of crypto/elliptic/elliptic.go Marshal
     # without point on curve check

volcenginesdkarkruntime/resources/chat/completions.py

@@ -23,13 +23,15 @@
     AsyncIterator,
 )
 
+import os
+import json
 import httpx
 import warnings
 from typing_extensions import Literal
 
 from ..._types import Body, Query, Headers
 from ..._utils._utils import deepcopy_minimal, with_sts_token, async_with_sts_token
-from ..._utils._key_agreement import aes_gcm_decrypt_base64_string, aes_gcm_decrypt_base64_list
+from ..._utils._key_agreement import aes_gcm_decrypt_base64_string, aes_gcm_decrypt_base64_list, decrypt_validate
 from ..._base_client import make_request_options
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._compat import cached_property
@@ -69,7 +71,8 @@ def _process_messages(
                         part["text"] = f(part["text"])
                     elif part.get("type", None) == "image_url":
                         if part["image_url"]["url"].startswith("data:"):
-                            part["image_url"]["url"] = f(part["image_url"]["url"])
+                            part["image_url"]["url"] = f(
+                                part["image_url"]["url"])
                         else:
                             warnings.warn(
                                 "encryption is not supported for image url, "
@@ -103,15 +106,16 @@ def _encrypt(
         model: str,
         messages: Iterable[ChatCompletionMessageParam],
         extra_headers: Headers,
-    ) -> tuple[bytes, bytes]:
-        client = self._client._get_endpoint_certificate(model)
+    ) -> tuple[bytes, bytes, str, str]:
+        client, ring_id, key_id = self._client._get_endpoint_certificate(model)
         _crypto_key, _crypto_nonce, session_token = client.generate_ecies_key_pair()
         extra_headers["X-Session-Token"] = session_token
         _process_messages(
             messages,
-            lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x),
+            lambda x: client.encrypt_string_with_key(
+                _crypto_key, _crypto_nonce, x),
         )
-        return _crypto_key, _crypto_nonce
+        return _crypto_key, _crypto_nonce, ring_id, key_id
 
     def _decrypt_chunk(
         self, key: bytes, nonce: bytes, resp: Stream[ChatCompletionChunk]
@@ -142,10 +146,13 @@ def _decrypt(
                         choice.message is not None and choice.finish_reason != 'content_filter'
                         and choice.message.content is not None
                     ):
-                        content = aes_gcm_decrypt_base64_string(
-                            key, nonce, choice.message.content
-                        )
-                        if content == '':
+                        try:
+                            content = aes_gcm_decrypt_base64_string(
+                                key, nonce, choice.message.content
+                            )
+                        except Exception:
+                            content = ''
+                        if content == '' or not decrypt_validate(choice.message.content):
                             content = aes_gcm_decrypt_base64_list(
                                 key, nonce, choice.message.content
                             )
@@ -197,7 +204,15 @@ def create(
         ):
             is_encrypt = True
             messages = deepcopy_minimal(messages)
-            e2e_key, e2e_nonce = self._encrypt(model, messages, extra_headers)
+            e2e_key, e2e_nonce, ring_id, key_id = self._encrypt(
+                model, messages, extra_headers)
+            if os.environ.get("VOLC_ARK_ENCRYPTION") == "AICC":
+                info = {
+                    'Version': 'AICCv0.1',
+                    'RingID': ring_id,
+                    'KeyID': key_id,
+                }
+                extra_headers["X-Encrypt-Info"] = json.dumps(info)
 
         resp = self._post(
             "/chat/completions",
@@ -257,15 +272,16 @@ def _encrypt(
         model: str,
         messages: Iterable[ChatCompletionMessageParam],
         extra_headers: Headers,
-    ) -> tuple[bytes, bytes]:
-        client = self._client._get_endpoint_certificate(model)
+    ) -> tuple[bytes, bytes, str, str]:
+        client, ring_id, key_id = self._client._get_endpoint_certificate(model)
         _crypto_key, _crypto_nonce, session_token = client.generate_ecies_key_pair()
         extra_headers["X-Session-Token"] = session_token
         _process_messages(
             messages,
-            lambda x: client.encrypt_string_with_key(_crypto_key, _crypto_nonce, x),
+            lambda x: client.encrypt_string_with_key(
+                _crypto_key, _crypto_nonce, x),
         )
-        return _crypto_key, _crypto_nonce
+        return _crypto_key, _crypto_nonce, ring_id, key_id
 
     async def _decrypt_chunk(
         self, key: bytes, nonce: bytes, resp: AsyncStream[ChatCompletionChunk]
@@ -296,10 +312,13 @@ async def _decrypt(
                         choice.message is not None and choice.finish_reason != 'content_filter'
                         and choice.message.content is not None
                     ):
-                        content = aes_gcm_decrypt_base64_string(
-                            key, nonce, choice.message.content
-                        )
-                        if content == '':
+                        try:
+                            content = aes_gcm_decrypt_base64_string(
+                                key, nonce, choice.message.content
+                            )
+                        except Exception:
+                            content = ''
+                        if content == '' or not decrypt_validate(choice.message.content):
                             content = aes_gcm_decrypt_base64_list(
                                 key, nonce, choice.message.content
                             )
@@ -351,7 +370,15 @@ async def create(
         ):
             is_encrypt = True
             messages = deepcopy_minimal(messages)
-            e2e_key, e2e_nonce = self._encrypt(model, messages, extra_headers)
+            e2e_key, e2e_nonce, ring_id, key_id = self._encrypt(
+                model, messages, extra_headers)
+            if os.environ.get("VOLC_ARK_ENCRYPTION") == "AICC":
+                info = {
+                    'Version': 'AICCv0.1',
+                    'RingID': ring_id,
+                    'KeyID': key_id,
+                }
+                extra_headers["X-Encrypt-Info"] = json.dumps(info)
 
         resp = await self._post(
             "/chat/completions",