fix: EncryptField unable to process values that starts with special header (#237)

Author: piglei

.github/workflows/blue-krill.yml

@@ -19,19 +19,19 @@ jobs:
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: 3.11
-      - name: Check with ruff
-        working-directory: sdks/blue-krill
-        run: |
-          pip install ruff==0.9.6
-          ruff check --config ../../pyproject.toml
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v5
+      - name: Install Poetry system-wide
+        # Use uv install to install poetry directly instead github action for simplicity
+        run: uv pip install --system poetry==2.1.3
+      - name: Ruff check
+        run: poetry install && poetry run ruff check ./sdks/blue-krill --config ./pyproject.toml
       - name: Lint with mypy
         working-directory: sdks/blue-krill
-        run: |
-          pip install mypy==0.910 types-requests==2.31.0.2 types-setuptools==57.4.18 types-dataclasses==0.1.7 types-redis==3.5.18 types-PyMySQL==1.1.0.1 types-six==0.1.9 types-toml==0.1.5
-          mypy . --config-file=pyproject.toml
+        run: poetry install && poetry run mypy . --config-file=./pyproject.toml
   test:
     strategy:
       fail-fast: false

pyproject.toml

@@ -1,10 +1,13 @@
 [tool.poetry]
-name = "蓝鲸 PaaS 平台 Python 工具集"
-version = "0.0.1"
-description = "BlueKing PaaS Python SDK"
-authors = ["blueking <blueking@tencent.com>"]
 package-mode = false
 
+[project]
+name = "bkpaas-python-sdk"
+description = "BlueKing PaaS Python SDK"
+version = "0.0.1"
+authors = [{ name = "blueking", email = "blueking@tencent.com" }]
+requires-python = '>=3.8.1,<3.12'
+
 [tool.poetry.dependencies]
 python = ">=3.8.1,<3.12"
 

sdks/blue-krill/blue_krill/encrypt/handler.py

@@ -56,11 +56,6 @@ def encrypt_cipher_type(self):
 
     def encrypt(self, text: str) -> str:
         """根据指定加密算法，加密字段"""
-        # 已加密则不处理
-        for cls in self.cipher_classes.values():
-            if cls.header.contain_header(text):
-                return text
-
         # 根据加密类型配置选择不同的加密算法
         try:
             cipher_class = self.cipher_classes[self.encrypt_cipher_type]

sdks/blue-krill/blue_krill/models/fields.py

@@ -22,6 +22,11 @@
 from blue_krill.encrypt.handler import EncryptHandler
 
 
+class _EncryptedString(str):
+    """A string that is encrypted, this type is used to distinguish between normal
+    strings and encrypted ones."""
+
+
 class EncryptField(models.TextField):
     """a field which will be encrypted via cryptography/fernet"""
 
@@ -34,7 +39,10 @@ def __init__(self, encrypt_cipher_type: Optional[str] = None, secret_key: Option
     def get_prep_value(self, value):
         if value is None:
             return value
-        return self.handler.encrypt(value)
+        # 如果值已经是 _EncryptedString 实例，则直接返回，避免重复触发加密逻辑
+        if isinstance(value, _EncryptedString):
+            return value
+        return _EncryptedString(self.handler.encrypt(value))
 
     def get_db_prep_value(self, value, connection, prepared=False):
         return self.get_prep_value(value)

sdks/blue-krill/pyproject.toml

@@ -72,3 +72,4 @@ build-backend = "poetry.core.masonry.api"
 ignore_missing_imports = true
 show_error_codes = true
 check_untyped_defs = true
+exclude = ['^tests/models/apps/.*/migrations/.*$']

sdks/blue-krill/tests/conftest.py

@@ -17,6 +17,7 @@
 
 
 def pytest_configure():
+    from cryptography.fernet import Fernet
     from django.conf import settings
 
     settings.configure(
@@ -28,12 +29,21 @@ def pytest_configure():
         TIME_ZONE="Asia/Shanghai",
         MIDDLEWARE=[],
         MIDDLEWARE_CLASSES=[],
-        INSTALLED_APPS=(),
+        # If the models of these apps have been updated, remember to remove all migration files so
+        # that the migrations could stay minimal because the test suite will create them from scratch.
+        INSTALLED_APPS=("tests.models.apps.test_encrypt",),
+        # For testing encrypt module
+        BKKRILL_ENCRYPT_SECRET_KEY=Fernet.generate_key(),
     )
 
     try:
         import django
+        from django.core.management import execute_from_command_line
 
         django.setup()
+
+        # Create migrations for test apps
+        execute_from_command_line(["manage.py", "makemigrations"])
+
     except AttributeError:
         pass

sdks/blue-krill/tests/models/__init__.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class TestEncryptConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "tests.models.apps.test_encrypt"