feat(decrypt): multi b64 string

feat(decrypt): multi b64 string

Author: liuhuiqi.7

volcenginesdkarkruntime/_utils/_key_agreement.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+import re
 import base64
 from typing import Tuple
 
@@ -72,6 +73,28 @@ def aes_gcm_decrypt_base64_string(key: bytes, nonce: bytes, ciphertext: str) ->
     return aes_gcm_decrypt_bytes(key, nonce, cipher_bytes).decode()
 
 
+base64_pattern = r'(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{4})'
+
+
+def aes_gcm_decrypt_base64_list(key: bytes, nonce: bytes, ciphertext: str) -> str:
+    # Decrypt
+    base64_array = re.findall(base64_pattern, ciphertext)
+    result = []
+    for b64 in base64_array:
+        try:
+            result.append(aes_gcm_decrypt_base64_string(key, nonce, b64))
+        except Exception:
+            for i in range(20, len(b64), 4):
+                try:
+                    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:])
+                    result.append(decrypted)
+                except Exception:
+                    result.append('')
+    return ''.join(result)
+
+
 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

@@ -29,7 +29,7 @@
 
 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
+from ..._utils._key_agreement import aes_gcm_decrypt_base64_string, aes_gcm_decrypt_base64_list
 from ..._base_client import make_request_options
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._compat import cached_property
@@ -142,9 +142,14 @@ def _decrypt(
                         choice.message is not None and choice.finish_reason != 'content_filter'
                         and choice.message.content is not None
                     ):
-                        choice.message.content = aes_gcm_decrypt_base64_string(
+                        content = aes_gcm_decrypt_base64_string(
                             key, nonce, choice.message.content
                         )
+                        if content == '':
+                            content = aes_gcm_decrypt_base64_list(
+                                key, nonce, choice.message.content
+                            )
+                        choice.message.content = content
                     resp.choices[index] = choice
             return resp
         else:
@@ -291,9 +296,14 @@ async def _decrypt(
                         choice.message is not None and choice.finish_reason != 'content_filter'
                         and choice.message.content is not None
                     ):
-                        choice.message.content = aes_gcm_decrypt_base64_string(
+                        content = aes_gcm_decrypt_base64_string(
                             key, nonce, choice.message.content
                         )
+                        if content == '':
+                            content = aes_gcm_decrypt_base64_list(
+                                key, nonce, choice.message.content
+                            )
+                        choice.message.content = content
                     resp.choices[index] = choice
             return resp
         else: