新增密码仓库(KMS)的初步支持 (#2912)

archery/settings.py

@@ -71,6 +71,7 @@
         ],
     ),
     CURRENT_AUDITOR=(str, "sql.utils.workflow_audit:AuditV2"),
+    PASSWORD_MIXIN_PATH=(str, "sql.plugins.password:DummyMixin"),
 )
 
 # SECURITY WARNING: keep the secret key used in production secret!
@@ -113,6 +114,8 @@
 
 CURRENT_AUDITOR = env("CURRENT_AUDITOR")
 
+PASSWORD_MIXIN_PATH = env("PASSWORD_MIXIN_PATH")
+
 # Application definition
 INSTALLED_APPS = (
     "django.contrib.admin",

sql/engines/__init__.py

@@ -24,8 +24,7 @@ def __init__(self, instance: Instance = None):
             self.instance_name = instance.instance_name
             self.host = instance.host
             self.port = int(instance.port)
-            self.user = instance.user
-            self.password = instance.password
+            self.user, self.password = self.instance.get_username_password()
             self.db_name = instance.db_name
             self.mode = instance.mode
 

sql/models.py

@@ -1,15 +1,34 @@
 # -*- coding: UTF-8 -*-
+import importlib
+import logging
 from typing import Optional
 
 from django.db import models
 from django.contrib.auth.models import AbstractUser
 from mirage import fields
 from django.utils.translation import gettext as _
+from django.conf import settings
 from mirage.crypto import Crypto
 
 from common.utils.const import WorkflowStatus, WorkflowType, WorkflowAction
 
 
+logger = logging.getLogger("default")
+file, _class = settings.PASSWORD_MIXIN_PATH.split(":")
+
+try:
+    password_module = importlib.import_module(file)
+    PasswordMixin = getattr(password_module, _class)
+except (ImportError, AttributeError) as e:
+    logger.error(
+        f"failed to import password minxin {settings.PASSWORD_MIXIN_PATH}, {str(e)}"
+    )
+    logger.error(f"falling back to dummy mixin")
+    from sql.plugins.password import DummyMixin
+
+    PasswordMixin = DummyMixin
+
+
 class ResourceGroup(models.Model):
     """
     资源组
@@ -179,7 +198,7 @@ class Meta:
         verbose_name_plural = "隧道配置"
 
 
-class Instance(models.Model):
+class Instance(models.Model, PasswordMixin):
     """
     各个线上实例配置
     """

sql/plugins/password.py

@@ -0,0 +1,63 @@
+"""
+本文件存放一些和密码相关的插件
+
+插件提供两个方法, 获取用户名和密码, 用于连接数据库
+插件在使用时, 会和 Instance 一起使用, 可以用 self 对象获取 instance 的各种信息
+
+包括名字, 类型, 等等. 通过这些信息, 插件可以获取数据库的用户名和密码
+"""
+
+import time
+import requests
+
+
+class DummyMixin:
+    """mixin 模板, 用于提供一些基础的方法, 给其他 mixin 继承
+    默认从schema 中直接提取 username 和 password
+    """
+
+    def get_username_password(self):
+        return self.user, self.password
+
+
+password_cache = {
+    "instance_name": {
+        "username": "username",
+        "password": "password",
+        "expires_at": "1740557906.15272",
+    }
+}
+
+
+class VaultMixin(DummyMixin):
+    """
+    和 sqlinstance 搭配使用
+    从 vault 中获取用户名和密码, 调用的是 localhost 8000 端口的 vault 服务
+    不使用任何 token, 适合 vault-proxy 部署方式, 如需其他部署方式, 可继承后修改配置
+    使用的是 static secret, 如需其他获取方式, 可继承后修改配置
+    """
+
+    vault_server = "localhost:8200"
+    vault_token = ""
+
+    def get_username_password(self):
+        if self.instance_name in password_cache:
+            if password_cache[self.instance_name]["expires_at"] > time.time():
+                return (
+                    password_cache[self.instance_name]["username"],
+                    password_cache[self.instance_name]["password"],
+                )
+
+        vault_role = f"{self.instance_name}-archery-rw"
+        response = requests.get(
+            f"http://{self.vault_server}/v1/database/static-creds/{vault_role}",
+            headers={"X-Vault-Token": self.vault_token},
+        )
+        response.raise_for_status()
+        data = response.json()["data"]
+        password_cache[self.instance_name] = {
+            "username": data["username"],
+            "password": data["password"],
+            "expires_at": time.time() + data["ttl"] - 60,
+        }
+        return data["username"], data["password"]

sql/plugins/tests.py

@@ -1,20 +1,22 @@
 # -*- coding: UTF-8 -*-
-""" 
+"""
 @author: hhyo
 @license: Apache Licence
 @file: tests.py
 @time: 2019/03/04
 """
 import json
 from django.test import Client, TestCase
-from unittest.mock import patch, ANY
+from unittest.mock import patch, ANY, Mock
+from pytest_mock import MockerFixture
 from django.contrib.auth import get_user_model
 
 from sql.plugins.my2sql import My2SQL
 from sql.plugins.schemasync import SchemaSync
 from sql.plugins.soar import Soar
 from sql.plugins.sqladvisor import SQLAdvisor
 from sql.plugins.pt_archiver import PtArchiver
+from sql.plugins.password import VaultMixin
 
 from common.config import SysConfig
 
@@ -334,3 +336,25 @@ def test_rewrite(self, _subprocess):
         # 异常测试
         with self.assertRaises(RuntimeError):
             self.soar.rewrite(sql, "unknown")
+
+
+def test_password_mixin(mocker: MockerFixture):
+    from sql.plugins.password import requests
+
+    class MockReponse(Mock):
+        def json(self):
+            return {"data": {"username": "test", "password": "test", "ttl": 360}}
+
+    mocker.patch.object(requests, "get", return_value=MockReponse())
+
+    class DummyInstance:
+        instance_name = "dummy"
+
+    class CompondInstance(DummyInstance, VaultMixin):
+        pass
+
+    instance = CompondInstance()
+    username, password = instance.get_username_password()
+    assert username == "test"
+    assert password == "test"
+    assert requests.get.call_count == 1

sql/test_model.py

@@ -1,9 +1,17 @@
 """models.py 的补充测试"""
 
+from django.conf import settings
 from sql.models import InstanceTag
 
 
 def test_instance_tag_str():
     i = InstanceTag(tag_name="test")
 
     assert str(i) == "test"
+
+
+def test_password_mixin_import_error():
+    settings.PASSWORD_MIXIN_PATH = "sql.not_found:ErrorMixin"
+    from sql.models import PasswordMixin
+
+    assert PasswordMixin.__name__ == "DummyMixin"