fix: Fix the issue of "Too many arguments for this mode" in AES CTR mode (fixed #16)

Author: ZhuoZhuoCrayon

bkcrypto/symmetric/ciphers/aes.py

@@ -13,6 +13,7 @@
 from dataclasses import dataclass
 
 from Cryptodome.Cipher import AES
+from Cryptodome.Util import Counter
 
 from bkcrypto import constants, types
 
@@ -54,16 +55,29 @@ class AESSymmetricCipher(base.BaseSymmetricCipher):
     def get_block_size(self) -> int:
         return self.config.key_size
 
-    def _encrypt(self, plaintext_bytes: bytes, encryption_metadata: base.EncryptionMetadata) -> bytes:
-
+    def init_ctx(self, encryption_metadata: base.EncryptionMetadata):
         mode_init_args: typing.List[bytes] = []
-        if self.config.enable_iv:
-            mode_init_args.append(encryption_metadata.iv)
+        mode_init_kwargs: typing.Dict[str : typing.Any] = {}
 
-        cipher_ctx = AES.new(self.config.key, self.config.mode_class, *mode_init_args)
+        if self.config.enable_iv:
+            if self.config.mode == constants.SymmetricMode.CTR:
+                # Size of the counter block must match block size
+                mode_init_kwargs["counter"] = Counter.new(
+                    self.get_block_size() * 8, initial_value=int.from_bytes(encryption_metadata.iv, byteorder="big")
+                )
+            else:
+                mode_init_args.append(encryption_metadata.iv)
+
+        cipher_ctx = AES.new(self.config.key, self.config.mode_class, *mode_init_args, **mode_init_kwargs)
         if self.config.enable_aad:
             cipher_ctx.update(encryption_metadata.aad)
 
+        return cipher_ctx
+
+    def _encrypt(self, plaintext_bytes: bytes, encryption_metadata: base.EncryptionMetadata) -> bytes:
+
+        cipher_ctx = self.init_ctx(encryption_metadata)
+
         if self.config.mode == constants.SymmetricMode.GCM:
             ciphertext_bytes, tag = cipher_ctx.encrypt_and_digest(plaintext_bytes)
             encryption_metadata.tag = tag
@@ -73,13 +87,7 @@ def _encrypt(self, plaintext_bytes: bytes, encryption_metadata: base.EncryptionM
 
     def _decrypt(self, ciphertext_bytes: bytes, encryption_metadata: base.EncryptionMetadata) -> bytes:
 
-        mode_init_args: typing.List[bytes] = []
-        if self.config.enable_iv:
-            mode_init_args.append(encryption_metadata.iv)
-
-        cipher_ctx = AES.new(self.config.key, self.config.mode_class, *mode_init_args)
-        if self.config.enable_aad:
-            cipher_ctx.update(encryption_metadata.aad)
+        cipher_ctx = self.init_ctx(encryption_metadata)
 
         if self.config.mode == constants.SymmetricMode.GCM:
             plaintext_bytes: bytes = cipher_ctx.decrypt_and_verify(ciphertext_bytes, encryption_metadata.tag)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bk-crypto-python-sdk"
-version = "1.0.3"
+version = "1.0.4"
 description = "bk-crypto-python-sdk is a lightweight cryptography toolkit for Python applications based on Cryptodome / tongsuopy and other encryption libraries."
 authors = ["TencentBlueKing <[email protected]>"]
 readme = "readme.md"

release.md

@@ -31,3 +31,10 @@
 ### Feature
 
 * [ Feature ] Support configuring AsymmetricCipherManager through Django settings ([#14](https://github.com/TencentBlueKing/crypto-python-sdk/issues/14))
+
+
+## 1.0.4 - 2023-07-20
+
+### Fixed
+
+* [ Fixed ] Fix the issue of "Too many arguments for this mode" in AES CTR mode ([#16](https://github.com/TencentBlueKing/crypto-python-sdk/issues/16))