From a763576be8345fc8e97b59bf05d35d3c8f93f2a5 Mon Sep 17 00:00:00 2001
From: nannan00 <17491932+nannan00@users.noreply.github.com>
Date: Mon, 22 Jan 2024 15:47:46 +0800
Subject: [PATCH] feat(cmd): add register idp plugin (#1540)

---
 .../bkuser/apps/idp/management/__init__.py    | 10 +++
 .../apps/idp/management/commands/__init__.py  | 10 +++
 .../commands/register_idp_plugin.py           | 68 +++++++++++++++++++
 src/idp-plugins/idp_plugins/constants.py      |  3 +
 src/idp-plugins/idp_plugins/models.py         |  8 +++
 5 files changed, 99 insertions(+)
 create mode 100644 src/bk-user/bkuser/apps/idp/management/__init__.py
 create mode 100644 src/bk-user/bkuser/apps/idp/management/commands/__init__.py
 create mode 100644 src/bk-user/bkuser/apps/idp/management/commands/register_idp_plugin.py

diff --git a/src/bk-user/bkuser/apps/idp/management/__init__.py b/src/bk-user/bkuser/apps/idp/management/__init__.py
new file mode 100644
index 000000000..1060b7bf4
--- /dev/null
+++ b/src/bk-user/bkuser/apps/idp/management/__init__.py
@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-用户管理(Bk-User) available.
+Copyright (C) 2017-2021 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 http://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.
+"""
diff --git a/src/bk-user/bkuser/apps/idp/management/commands/__init__.py b/src/bk-user/bkuser/apps/idp/management/commands/__init__.py
new file mode 100644
index 000000000..1060b7bf4
--- /dev/null
+++ b/src/bk-user/bkuser/apps/idp/management/commands/__init__.py
@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-用户管理(Bk-User) available.
+Copyright (C) 2017-2021 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 http://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.
+"""
diff --git a/src/bk-user/bkuser/apps/idp/management/commands/register_idp_plugin.py b/src/bk-user/bkuser/apps/idp/management/commands/register_idp_plugin.py
new file mode 100644
index 000000000..03011e6ec
--- /dev/null
+++ b/src/bk-user/bkuser/apps/idp/management/commands/register_idp_plugin.py
@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+"""
+TencentBlueKing is pleased to support the open source community by making 蓝鲸智云-用户管理(Bk-User) available.
+Copyright (C) 2017-2021 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 http://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 os
+
+from django.conf import settings
+from django.core.management.base import BaseCommand
+from django.utils.module_loading import import_string
+
+from bkuser.apps.idp.models import IdpPlugin
+from bkuser.idp_plugins.constants import CUSTOM_PLUGIN_ID_PREFIX, MAX_LOGO_SIZE
+from bkuser.utils.base64 import load_image_as_base64
+
+
+class Command(BaseCommand):
+    """向数据库中写入认证源插件信息"""
+
+    def add_arguments(self, parser):
+        parser.add_argument("--dir_name", dest="dir_name", required=True, help="插件目录名称")
+
+    def handle(self, dir_name, *args, **options):
+        plugin_base_dir = settings.BASE_DIR / "bkuser" / "idp_plugins"
+        # 1. 检查指定的插件目录是否存在
+        if not os.path.isdir(plugin_base_dir / dir_name):
+            raise RuntimeError(f"plugin directory [{dir_name}] not found in bkuser/idp_plugins!")
+
+        # 2. 检查自定义插件是否配置 Metadata
+        try:
+            metadata = import_string(f"bkuser.idp_plugins.{dir_name}.METADATA")
+        except ImportError:
+            raise RuntimeError("custom idp plugin must set metadata!")
+
+        # 3. 确保自定义插件的 ID 符合规范
+        if not metadata.id.startswith(CUSTOM_PLUGIN_ID_PREFIX):
+            raise RuntimeError(f"custom idp plugin's id must start with `{CUSTOM_PLUGIN_ID_PREFIX}`")
+
+        logo_path = plugin_base_dir / f"{dir_name}/logo.png"
+
+        # 4. 如果发现有 logo，还需要检查下尺寸大小，避免有性能问题
+        if os.path.exists(logo_path) and os.path.getsize(logo_path) > MAX_LOGO_SIZE:
+            raise RuntimeError(f"plugin logo size must be less than {MAX_LOGO_SIZE/1024}KB!")
+
+        # 5. 尝试获取下 logo，取不到就用默认的
+        try:
+            logo = load_image_as_base64(logo_path)
+        except Exception:
+            self.stdout.write("failed to load plugin logo, use default logo...")
+            logo = ""
+
+        # 6. 如果同名插件已存在，更新，否则创建
+        IdpPlugin.objects.update_or_create(
+            id=metadata.id,
+            defaults={
+                "name": metadata.name,
+                "description": metadata.description,
+                "logo": logo,
+            },
+        )
+
+        # 7. 注册到 DB 成功要给提示
+        self.stdout.write(f"register idp plugin [{metadata.id}] into database successfully.")
diff --git a/src/idp-plugins/idp_plugins/constants.py b/src/idp-plugins/idp_plugins/constants.py
index 9c10b5674..26e8e46cb 100644
--- a/src/idp-plugins/idp_plugins/constants.py
+++ b/src/idp-plugins/idp_plugins/constants.py
@@ -14,6 +14,9 @@
 # 非内置插件，必须以指定前缀开头
 CUSTOM_PLUGIN_ID_PREFIX = "custom_"
 
+# Logo 限制 64KB，过大的 logo 会导致性能下降，1024 * 1024 约 37KB
+MAX_LOGO_SIZE = 64 * 1024
+
 
 class PluginTypeEnum(str, StructuredEnum):
     """认证源插件类型枚举"""
diff --git a/src/idp-plugins/idp_plugins/models.py b/src/idp-plugins/idp_plugins/models.py
index c3382ea07..a6192bef0 100644
--- a/src/idp-plugins/idp_plugins/models.py
+++ b/src/idp-plugins/idp_plugins/models.py
@@ -13,6 +13,14 @@
 from .constants import AllowedHttpMethodEnum
 
 
+class PluginMetadata(BaseModel):
+    """插件基本信息"""
+
+    id: str
+    name: str
+    description: str
+
+
 class DispatchConfigItem(BaseModel):
     action: str = Field(pattern="^[a-zA-Z0-9][a-zA-Z0-9_-]{1,30}[a-zA-Z0-9]$")
     http_method: AllowedHttpMethodEnum