From 75be45ce43838dc77c317b4b9fc09a9276330d52 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Thu, 11 Jun 2020 12:10:00 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=BD=BF=E7=94=A8=E6=96=B0=E7=9A=84?=
 =?UTF-8?q?=E5=AF=B9=E7=A7=B0=E5=8A=A0=E5=AF=86=E6=96=B9=E5=BC=8F:=20aes?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/fields/model.py   | 24 ++++++++++++----
 apps/common/utils/__init__.py |  1 +
 apps/common/utils/crypto.py   | 54 +++++++++++++++++++++++++++++++++++
 apps/common/utils/encode.py   | 50 --------------------------------
 4 files changed, 74 insertions(+), 55 deletions(-)
 create mode 100644 apps/common/utils/crypto.py

diff --git a/apps/common/fields/model.py b/apps/common/fields/model.py
index ddb61f7c0..bc5c5538d 100644
--- a/apps/common/fields/model.py
+++ b/apps/common/fields/model.py
@@ -3,8 +3,9 @@
 import json
 from django.db import models
 from django.utils.translation import ugettext_lazy as _
+from django.utils.encoding import force_text
 
-from ..utils import signer
+from ..utils import signer, aes_crypto
 
 
 __all__ = [
@@ -114,11 +115,22 @@ class EncryptMixin:
     def from_db_value(self, value, expression, connection, context):
         if value is None:
             return value
-        value = signer.unsign(value)
+        value = force_text(value)
+
+        plain_value = ''
+        # 优先采用 aes 解密
+        try:
+            plain_value = aes_crypto.decrypt(value)
+        except (TypeError, ValueError):
+            pass
+
+        # 如果没有解开，使用原来的signer解密
+        if not plain_value:
+            plain_value = signer.unsign(value) or ''
         sp = super()
         if hasattr(sp, 'from_db_value'):
-            return sp.from_db_value(value, expression, connection, context)
-        return value
+            plain_value = sp.from_db_value(plain_value, expression, connection, context)
+        return plain_value
 
     def get_prep_value(self, value):
         if value is None:
@@ -126,7 +138,9 @@ class EncryptMixin:
         sp = super()
         if hasattr(sp, 'get_prep_value'):
             value = sp.get_prep_value(value)
-        return signer.sign(value)
+        value = force_text(value)
+        # 替换新的加密方式
+        return aes_crypto.encrypt(value)
 
 
 class EncryptTextField(EncryptMixin, models.TextField):
diff --git a/apps/common/utils/__init__.py b/apps/common/utils/__init__.py
index 802943072..01850f0cf 100644
--- a/apps/common/utils/__init__.py
+++ b/apps/common/utils/__init__.py
@@ -6,3 +6,4 @@ from .django import *
 from .encode import *
 from .http import *
 from .ipip import *
+from .crypto import *
diff --git a/apps/common/utils/crypto.py b/apps/common/utils/crypto.py
new file mode 100644
index 000000000..ea3590d6c
--- /dev/null
+++ b/apps/common/utils/crypto.py
@@ -0,0 +1,54 @@
+import base64
+from Crypto.Cipher import AES
+
+from django.conf import settings
+
+
+class AESCrypto:
+    """
+    AES
+    除了MODE_SIV模式key长度为：32, 48, or 64,
+    其余key长度为16, 24 or 32
+    详细见AES内部文档
+    CBC模式传入iv参数
+    本例使用常用的ECB模式
+    """
+
+    def __init__(self, key):
+        if len(key) > 32:
+            key = key[:32]
+        self.key = self.to_16(key)
+
+    @staticmethod
+    def to_16(key):
+        """
+        转为16倍数的bytes数据
+        :param key:
+        :return:
+        """
+        key = bytes(key, encoding="utf8")
+        while len(key) % 16 != 0:
+            key += b'\0'
+        return key  # 返回bytes
+
+    def aes(self):
+        return AES.new(self.key, AES.MODE_ECB)  # 初始化加密器
+
+    def encrypt(self, text):
+        aes = self.aes()
+        return str(base64.encodebytes(aes.encrypt(self.to_16(text))),
+                   encoding='utf8').replace('\n', '')  # 加密
+
+    def decrypt(self, text):
+        aes = self.aes()
+        return str(aes.decrypt(base64.decodebytes(bytes(text, encoding='utf8'))).rstrip(b'\0').decode("utf8"))  # 解密
+
+
+def get_aes_crypto(key=None):
+    if key is None:
+        key = settings.SECRET_KEY
+    a = AESCrypto(key)
+    return a
+
+
+aes_crypto = get_aes_crypto()
diff --git a/apps/common/utils/encode.py b/apps/common/utils/encode.py
index 72e011917..cd130e7fe 100644
--- a/apps/common/utils/encode.py
+++ b/apps/common/utils/encode.py
@@ -9,7 +9,6 @@ import time
 import hashlib
 from io import StringIO
 from itertools import chain
-from Crypto.Cipher import AES
 
 import paramiko
 import sshpubkeys
@@ -227,52 +226,3 @@ def model_to_json(instance, sort_keys=True, indent=2, cls=None):
         cls = DjangoJSONEncoder
     return json.dumps(data, sort_keys=sort_keys, indent=indent, cls=cls)
 
-
-class AESCrypto:
-    """
-    AES
-    除了MODE_SIV模式key长度为：32, 48, or 64,
-    其余key长度为16, 24 or 32
-    详细见AES内部文档
-    CBC模式传入iv参数
-    本例使用常用的ECB模式
-    """
-
-    def __init__(self, key):
-        if len(key) > 32:
-            key = key[:32]
-        self.key = self.to_16(key)
-
-    @staticmethod
-    def to_16(key):
-        """
-        转为16倍数的bytes数据
-        :param key:
-        :return:
-        """
-        key = bytes(key, encoding="utf8")
-        while len(key) % 16 != 0:
-            key += b'\0'
-        return key  # 返回bytes
-
-    def aes(self):
-        return AES.new(self.key, AES.MODE_ECB)  # 初始化加密器
-
-    def encrypt(self, text):
-        aes = self.aes()
-        return str(base64.encodebytes(aes.encrypt(self.to_16(text))),
-                   encoding='utf8').replace('\n', '')  # 加密
-
-    def decrypt(self, text):
-        aes = self.aes()
-        return str(aes.decrypt(base64.decodebytes(bytes(text, encoding='utf8'))).rstrip(b'\0').decode("utf8"))  # 解密
-
-
-def get_aes_crypto(key=None):
-    if key is None:
-        key = settings.SECRET_KEY
-    a = AESCrypto(key)
-    return a
-
-
-aes = get_aes_crypto()