feat: Add support for prefix ciphertext decryption (closed #20)

Author: ZhuoZhuoCrayon

bkcrypto/contrib/django/ciphers.py

@@ -9,6 +9,7 @@
 an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 specific language governing permissions and limitations under the License.
 """
+import abc
 import typing
 
 from bkcrypto.asymmetric.ciphers import BaseAsymmetricCipher
@@ -17,7 +18,7 @@
 from bkcrypto.symmetric.ciphers import BaseSymmetricCipher
 from bkcrypto.symmetric.options import SymmetricOptions
 
-from .init_configs import AsymmetricCipherInitConfig, SymmetricCipherInitConfig
+from .init_configs import CipherInitConfig
 from .settings import crypto_settings
 
 
@@ -47,54 +48,86 @@ def get_symmetric_cipher(
     )
 
 
-class SymmetricCipherManager:
-    _cache: typing.Optional[typing.Dict[str, BaseSymmetricCipher]] = None
+class BaseCipherManager(abc.ABC):
 
-    def __init__(self):
-        self._cache: [str, BaseSymmetricCipher] = {}
+    _cache: typing.Dict[str, typing.Any] = None
 
-    def cipher(
-        self, using: typing.Optional[str] = None, cipher_type: typing.Optional[str] = None
-    ) -> BaseSymmetricCipher:
+    def __init__(self):
+        self._cache = {}
 
+    def _get_init_config(self, using: typing.Optional[str] = None) -> CipherInitConfig:
         using: str = using or "default"
-        if using not in crypto_settings.SYMMETRIC_CIPHERS:
+        init_configs: typing.Dict[str, CipherInitConfig] = self._get_init_configs_from_settings()
+        if using not in init_configs:
             raise RuntimeError(f"Invalid using {using}")
+        return init_configs[using]
+
+    @abc.abstractmethod
+    def _get_init_configs_from_settings(self) -> typing.Dict[str, CipherInitConfig]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _get_cipher_type_from_settings(self) -> str:
+        raise NotImplementedError
 
-        cipher_type: str = cipher_type or crypto_settings.SYMMETRIC_CIPHER_TYPE
+    @abc.abstractmethod
+    def _get_cipher(self, cipher_type: str, init_config: CipherInitConfig):
+        raise NotImplementedError
+
+    def _cipher(self, using: typing.Optional[str] = None, cipher_type: typing.Optional[str] = None):
+
+        # try to get cipher from cache
+        cipher_type: str = cipher_type or self._get_cipher_type_from_settings()
         cache_key: str = f"{using}-{cipher_type}"
         if cache_key in self._cache:
             return self._cache[cache_key]
 
-        init_config: SymmetricCipherInitConfig = crypto_settings.SYMMETRIC_CIPHERS[using]
-        cipher: BaseSymmetricCipher = get_symmetric_cipher(**init_config.as_get_cipher_params(cipher_type))
-        self._cache[cache_key] = cipher
-        return cipher
+        # create & cache instance
+        init_config: CipherInitConfig = self._get_init_config(using=using)
+        self._cache[cache_key] = self._get_cipher(cipher_type, init_config)
+        return self._cache[cache_key]
 
+    @abc.abstractmethod
+    def cipher(self, using: typing.Optional[str] = None, cipher_type: typing.Optional[str] = None):
+        raise NotImplementedError
 
-class AsymmetricCipherManager:
-    _cache: typing.Optional[typing.Dict[str, BaseAsymmetricCipher]] = None
 
-    def __init__(self):
-        self._cache: [str, BaseAsymmetricCipher] = {}
+class SymmetricCipherManager(BaseCipherManager):
+
+    _cache: typing.Optional[typing.Dict[str, BaseSymmetricCipher]] = None
+
+    def _get_init_configs_from_settings(self) -> typing.Dict[str, CipherInitConfig]:
+        return crypto_settings.SYMMETRIC_CIPHERS
+
+    def _get_cipher_type_from_settings(self) -> str:
+        return crypto_settings.SYMMETRIC_CIPHER_TYPE
+
+    def _get_cipher(self, cipher_type: str, init_config: CipherInitConfig) -> BaseSymmetricCipher:
+        return get_symmetric_cipher(**init_config.as_get_cipher_params(cipher_type))
 
     def cipher(
         self, using: typing.Optional[str] = None, cipher_type: typing.Optional[str] = None
-    ) -> BaseAsymmetricCipher:
+    ) -> BaseSymmetricCipher:
+        return self._cipher(using, cipher_type)
 
-        using: str = using or "default"
-        if using not in crypto_settings.ASYMMETRIC_CIPHERS:
-            raise RuntimeError(f"Invalid using {using}")
 
-        cipher_type: str = cipher_type or crypto_settings.ASYMMETRIC_CIPHER_TYPE
-        cache_key: str = f"{using}-{cipher_type}"
-        if cache_key in self._cache:
-            return self._cache[cache_key]
+class AsymmetricCipherManager(BaseCipherManager):
 
-        init_config: AsymmetricCipherInitConfig = crypto_settings.ASYMMETRIC_CIPHERS[using]
-        cipher: BaseAsymmetricCipher = get_asymmetric_cipher(**init_config.as_get_cipher_params(cipher_type))
-        self._cache[cache_key] = cipher
-        return cipher
+    _cache: typing.Optional[typing.Dict[str, BaseAsymmetricCipher]] = None
+
+    def _get_init_configs_from_settings(self) -> typing.Dict[str, CipherInitConfig]:
+        return crypto_settings.ASYMMETRIC_CIPHERS
+
+    def _get_cipher_type_from_settings(self) -> str:
+        return crypto_settings.ASYMMETRIC_CIPHER_TYPE
+
+    def _get_cipher(self, cipher_type: str, init_config: CipherInitConfig) -> BaseAsymmetricCipher:
+        return get_asymmetric_cipher(**init_config.as_get_cipher_params(cipher_type))
+
+    def cipher(
+        self, using: typing.Optional[str] = None, cipher_type: typing.Optional[str] = None
+    ) -> BaseAsymmetricCipher:
+        return self._cipher(using, cipher_type)
 
 
 symmetric_cipher_manager = SymmetricCipherManager()

bkcrypto/contrib/django/fields.py

@@ -13,39 +13,10 @@
 
 from django.db import models
 
-from bkcrypto.contrib.django.ciphers import symmetric_cipher_manager
-from bkcrypto.contrib.django.init_configs import SymmetricCipherInitConfig
-from bkcrypto.contrib.django.settings import crypto_settings
-from bkcrypto.symmetric.ciphers.base import BaseSymmetricCipher
+from bkcrypto.contrib.django.selectors import SymmetricCipherSelectorMixin
 
 
-class SymmetricFieldMixin:
-
-    cipher: BaseSymmetricCipher = None
-
-    # 是否指定固定前缀，如果不为 None，密文将统一使用 prefix 作为前缀
-    prefix: str = None
-    # 指定对称加密实例，默认使用 `default`
-    using: str = None
-
-    def prefix_selector(self, value: str) -> typing.Tuple[bool, str, typing.Optional[str]]:
-        """
-        密文前缀匹配，用于提取可能存在的加密类型
-        :param value:
-        :return:
-        """
-        if self.prefix is not None:
-            if value.startswith(self.prefix):
-                return True, value[len(self.prefix) :], None
-            else:
-                return False, value, None
-        else:
-            init_config: SymmetricCipherInitConfig = crypto_settings.SYMMETRIC_CIPHERS[self.using]
-            for prefix, cipher_type in init_config.prefix_cipher_type_map.items():
-                if value.startswith(prefix):
-                    return True, value[len(prefix) :], cipher_type
-            return False, value, None
-
+class SymmetricFieldMixin(SymmetricCipherSelectorMixin):
     def deconstruct(self):
         name, path, args, kwargs = super().deconstruct()
         if self.prefix is not None:
@@ -63,15 +34,7 @@ def get_decrypted_value(self, value):
         if value is None:
             return value
 
-        is_match, trusted_value, cipher_type = self.prefix_selector(value)
-        if is_match:
-            try:
-                cipher: BaseSymmetricCipher = symmetric_cipher_manager.cipher(using=self.using, cipher_type=cipher_type)
-                value = cipher.decrypt(trusted_value)
-            except Exception:
-                pass
-
-        return value
+        return self.decrypt(value)
 
     def from_db_value(self, value, expression, connection, context=None):
         """出库后解密数据"""
@@ -102,16 +65,7 @@ def get_prep_value(self, value):
         if hasattr(sp, "get_prep_value"):
             value = sp.get_prep_value(value)
 
-        if self.prefix is not None:
-            prefix: str = self.prefix
-        else:
-            init_config: SymmetricCipherInitConfig = crypto_settings.SYMMETRIC_CIPHERS[self.using]
-            prefix: str = init_config.db_prefix_map[crypto_settings.SYMMETRIC_CIPHER_TYPE]
-
-        cipher: BaseSymmetricCipher = symmetric_cipher_manager.cipher(using=self.using)
-        value = prefix + cipher.encrypt(value)
-
-        return value
+        return self.encrypt(value)
 
 
 class SymmetricTextField(SymmetricFieldMixin, models.TextField):

bkcrypto/contrib/django/selectors.py

@@ -0,0 +1,123 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云 - crypto-python-sdk
+(BlueKing - crypto-python-sdk) available.
+Copyright (C) 2017-2023 THL A29 Limited, a Tencent company. All rights reserved.
+Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
+You may obtain a copy of the License at https://opensource.org/licenses/MIT
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+"""
+import abc
+import typing
+
+from bkcrypto.asymmetric.ciphers import BaseAsymmetricCipher
+from bkcrypto.contrib.django.ciphers import asymmetric_cipher_manager, symmetric_cipher_manager
+from bkcrypto.contrib.django.init_configs import CipherInitConfig
+from bkcrypto.contrib.django.settings import crypto_settings
+from bkcrypto.symmetric.ciphers.base import BaseSymmetricCipher
+
+
+class CipherSelectorMixin:
+    init_config: CipherInitConfig = None
+
+    # 是否指定固定前缀，如果不为 None，密文将统一使用 prefix 作为前缀
+    prefix: str = None
+    # 指定对称加密实例，默认使用 `default`
+    using: str = None
+
+    @abc.abstractmethod
+    def _get_cipher_type_from_settings(self) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_cipher(self, cipher_type: typing.Optional[str] = None):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_init_config(self) -> CipherInitConfig:
+        raise NotImplementedError
+
+    def prefix_selector(self, ciphertext_with_prefix: str) -> typing.Tuple[bool, str, typing.Optional[str]]:
+        """
+        密文前缀匹配，用于提取可能存在的加密类型
+        :param ciphertext_with_prefix:
+        :return:
+        """
+        if self.prefix is not None:
+            if ciphertext_with_prefix.startswith(self.prefix):
+                return True, ciphertext_with_prefix[len(self.prefix) :], None
+            else:
+                return False, ciphertext_with_prefix, None
+        else:
+            for prefix, cipher_type in self.get_init_config().prefix_cipher_type_map.items():
+                if ciphertext_with_prefix.startswith(prefix):
+                    return True, ciphertext_with_prefix[len(prefix) :], cipher_type
+            return False, ciphertext_with_prefix, None
+
+    def encrypt(self, plaintext: str) -> str:
+        if self.prefix is not None:
+            prefix: str = self.prefix
+        else:
+            prefix: str = self.get_init_config().db_prefix_map[self._get_cipher_type_from_settings()]
+
+        cipher = self.get_cipher()
+        ciphertext_with_prefix: str = prefix + cipher.encrypt(plaintext)
+        return ciphertext_with_prefix
+
+    def decrypt(self, ciphertext_with_prefix: str) -> str:
+
+        is_match, trusted_value, cipher_type = self.prefix_selector(ciphertext_with_prefix)
+        if is_match:
+            try:
+                # 解密时使用前缀匹配到的算法
+                cipher = self.get_cipher(cipher_type=cipher_type)
+                plaintext: str = cipher.decrypt(trusted_value)
+            except Exception:
+                return ciphertext_with_prefix
+        else:
+            return ciphertext_with_prefix
+
+        return plaintext
+
+
+class SymmetricCipherSelectorMixin(CipherSelectorMixin):
+    def _get_cipher_type_from_settings(self) -> str:
+        return crypto_settings.SYMMETRIC_CIPHER_TYPE
+
+    def get_cipher(self, cipher_type: typing.Optional[str] = None) -> BaseSymmetricCipher:
+        return symmetric_cipher_manager.cipher(using=self.using, cipher_type=cipher_type)
+
+    def get_init_config(self) -> CipherInitConfig:
+        return crypto_settings.SYMMETRIC_CIPHERS[self.using]
+
+
+class AsymmetricCipherSelectorMixin(CipherSelectorMixin):
+    def _get_cipher_type_from_settings(self) -> str:
+        return crypto_settings.ASYMMETRIC_CIPHER_TYPE
+
+    def get_cipher(self, cipher_type: typing.Optional[str] = None) -> BaseAsymmetricCipher:
+        return asymmetric_cipher_manager.cipher(using=self.using, cipher_type=cipher_type)
+
+    def get_init_config(self) -> CipherInitConfig:
+        return crypto_settings.ASYMMETRIC_CIPHERS[self.using]
+
+
+class CipherSelector:
+    def __init__(self, using: typing.Optional[str] = None, prefix: typing.Optional[str] = None):
+        """
+        对称加密
+        :param using: 指定对称加密实例，默认使用 `default`
+        :param prefix: 是否指定固定前缀，如果不为 None，密文将统一使用 prefix 作为前缀
+        """
+        self.prefix = prefix
+        self.using = using or "default"
+
+
+class SymmetricCipherSelector(SymmetricCipherSelectorMixin, CipherSelector):
+    pass
+
+
+class AsymmetricCipherSelector(AsymmetricCipherSelectorMixin, CipherSelector):
+    pass

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bk-crypto-python-sdk"
-version = "1.0.4"
+version = "1.1.0"
 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

@@ -1,4 +1,4 @@
-# 版本日志
+**# 版本日志
 
 ## 1.0.0 - 2023-07-03
 
@@ -38,3 +38,11 @@
 ### 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))
+
+
+## 1.1.0 - 2023-08-07
+
+### Feature
+
+* [ Feature ] Add support for non-Django projects ([#19](https://github.com/TencentBlueKing/crypto-python-sdk/issues/19))
+* [ Feature ] Add support for prefix ciphertext decryption ([#20](https://github.com/TencentBlueKing/crypto-python-sdk/issues/20))