From 1b2a97890ac4904f403e3a1e524e9c606a2d8cfe Mon Sep 17 00:00:00 2001
From: MimoAtHome <mingruirui@xd.com>
Date: Sat, 21 Oct 2023 02:17:20 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AF=86=E7=A0=81=E5=AD=98=E5=82=A8=20?=
 =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=86=85=E7=BD=AE=20AES=20=E5=8A=A0=E5=AF=86?=
 =?UTF-8?q?=EF=BC=9B=E5=A2=9E=E5=8A=A0=20Vault=20Transit=20=E5=8A=A0?=
 =?UTF-8?q?=E5=AF=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmdb-api/Pipfile             |  1 +
 cmdb-api/api/lib/utils.py    | 33 +++++++++++++++++++++++++++++++++
 cmdb-api/api/models/acl.py   | 16 ++++++++++++----
 cmdb-api/requirements.txt    |  1 +
 cmdb-api/settings.example.py |  4 ++++
 5 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/cmdb-api/Pipfile b/cmdb-api/Pipfile
index 60206f8..02ecea2 100644
--- a/cmdb-api/Pipfile
+++ b/cmdb-api/Pipfile
@@ -26,6 +26,7 @@ Flask-Bcrypt = "==1.0.1"
 Flask-Cors = ">=3.0.8"
 ldap3 = "==2.9.1"
 pycryptodome = "==3.12.0"
+hvac = "==1.2.1"
 # Caching
 Flask-Caching = ">=1.0.0"
 # Environment variable parsing
diff --git a/cmdb-api/api/lib/utils.py b/cmdb-api/api/lib/utils.py
index 6482856..555ec3b 100644
--- a/cmdb-api/api/lib/utils.py
+++ b/cmdb-api/api/lib/utils.py
@@ -6,6 +6,7 @@ import time
 from typing import Set
 
 import elasticsearch
+import hvac
 import redis
 import six
 from Crypto.Cipher import AES
@@ -286,3 +287,35 @@ class AESCrypto(object):
         text_decrypted = cipher.decrypt(encode_bytes)
 
         return cls.unpad(text_decrypted).decode('utf8')
+
+
+class VaultTransitCrypto:
+    TRANSIT_KEY_NAME = 'cmdb-hvac-key'
+    client = None
+
+    @classmethod
+    def init_client(cls):
+        if not cls.client or not cls.client.is_authenticated():
+            cls.client = hvac.Client(
+                url=current_app.config.get('VAULT_URL'),
+                token=current_app.config.get('VAULT_TOKEN'),
+            )
+
+    @classmethod
+    def encrypt(cls, text):
+        cls.init_client()
+        cls.client.secrets.transit.create_key(name=cls.TRANSIT_KEY_NAME)
+        encrypt_data_response = cls.client.secrets.transit.encrypt_data(
+            name=cls.TRANSIT_KEY_NAME,
+            plaintext=base64.b64encode(text.encode()).decode(),
+        )
+        return encrypt_data_response['data']['ciphertext']
+
+    @classmethod
+    def decrypt(cls, ciphertext):
+        cls.init_client()
+        decrypt_data_response = cls.client.secrets.transit.decrypt_data(
+            name=cls.TRANSIT_KEY_NAME,
+            ciphertext=ciphertext,
+        )
+        return base64.b64decode(decrypt_data_response['data']['plaintext']).decode()
diff --git a/cmdb-api/api/models/acl.py b/cmdb-api/api/models/acl.py
index a96bb1c..7e4e247 100644
--- a/cmdb-api/api/models/acl.py
+++ b/cmdb-api/api/models/acl.py
@@ -16,6 +16,7 @@ from api.lib.database import Model
 from api.lib.database import SoftDeleteMixin
 from api.lib.perm.acl.const import ACL_QUEUE
 from api.lib.perm.acl.const import OperateType
+from api.lib.utils import AESCrypto, VaultTransitCrypto
 
 
 class App(Model):
@@ -126,7 +127,7 @@ class User(CRUDModel, SoftDeleteMixin):
     catalog = db.Column(db.String(64))
     email = db.Column(db.String(100), unique=True, nullable=False)
     mobile = db.Column(db.String(14), unique=True)
-    _password = db.Column("password", db.String(80))
+    _password = db.Column("password", db.String(128))
     key = db.Column(db.String(32), nullable=False)
     secret = db.Column(db.String(32), nullable=False)
     date_joined = db.Column(db.DateTime, default=datetime.utcnow)
@@ -155,14 +156,21 @@ class User(CRUDModel, SoftDeleteMixin):
         return self._password
 
     def _set_password(self, password):
-        self._password = hashlib.md5(password.encode('utf-8')).hexdigest()
+        md5_password = hashlib.md5(password.encode('utf-8')).hexdigest()
+        if current_app.config.get("ENCRYPT_PASSWORD_TYPE") == 'VAULT':
+            self._password = VaultTransitCrypto.encrypt(md5_password)
+        else:
+            self._password = AESCrypto.encrypt(md5_password)
 
     password = db.synonym("_password", descriptor=property(_get_password, _set_password))
 
-    def check_password(self, password):
+    def check_password(self, md5_password):
         if self.password is None:
             return False
-        return self.password == password or self.password == hashlib.md5(password.encode('utf-8')).hexdigest()
+        if current_app.config.get("ENCRYPT_PASSWORD_TYPE") == 'VAULT':
+            return VaultTransitCrypto.decrypt(self.password) == md5_password
+        else:
+            return AESCrypto.decrypt(self.password) == md5_password
 
 
 class RoleQuery(BaseQuery):
diff --git a/cmdb-api/requirements.txt b/cmdb-api/requirements.txt
index 5df5a7b..3f7c824 100644
--- a/cmdb-api/requirements.txt
+++ b/cmdb-api/requirements.txt
@@ -18,6 +18,7 @@ Flask-RESTful==0.3.10
 Flask-SQLAlchemy==2.5.0
 future==0.18.3
 gunicorn==21.0.1
+hvac==1.2.1
 itsdangerous==2.1.2
 Jinja2==3.1.2
 jinja2schema==0.1.4
diff --git a/cmdb-api/settings.example.py b/cmdb-api/settings.example.py
index ed6b6ce..68cd02b 100644
--- a/cmdb-api/settings.example.py
+++ b/cmdb-api/settings.example.py
@@ -97,3 +97,7 @@ BOOL_TRUE = ['true', 'TRUE', 'True', True, '1', 1, "Yes", "YES", "yes", 'Y', 'y'
 
 # # messenger
 USE_MESSENGER = True
+
+ENCRYPT_PASSWORD_TYPE = "VAULT"  # "VAULT" or "AES"
+VAULT_URL = 'http://'  # 当 ENCRYPT_PASSWORD_TYPE 为 VAULT 时，需要配置
+VAULT_TOKEN = ''  # 当 ENCRYPT_PASSWORD_TYPE 为 VAULT 时，需要配置