From a86c598f72ea2c84ba11945caa5d3dd6f2262bb9 Mon Sep 17 00:00:00 2001 From: crayon <873217631@qq.com> Date: Thu, 20 Jul 2023 22:27:16 +0800 Subject: [PATCH] fix: Fix the issue of "Too many arguments for this mode" in AES CTR mode (fixed #16) --- bkcrypto/symmetric/ciphers/aes.py | 32 +++++++++++++++++++------------ pyproject.toml | 2 +- release.md | 7 +++++++ 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/bkcrypto/symmetric/ciphers/aes.py b/bkcrypto/symmetric/ciphers/aes.py index 3663f3e..814cb49 100644 --- a/bkcrypto/symmetric/ciphers/aes.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index a4404f7..a2bf9d1 100644 --- a/pyproject.toml +++ b/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 "] readme = "readme.md" diff --git a/release.md b/release.md index 38d1d5d..806cc7d 100644 --- a/release.md +++ b/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))