From a35e0c5efa6631604a44498adaecffbc197b1917 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Thu, 22 Sep 2022 16:39:41 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E5=88=9B=E5=BB=BA=E5=86=85=E7=BD=AE=20?=
 =?UTF-8?q?platforms?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/assets/const/host.py  | 14 ++++--------
 apps/assets/const/types.py | 45 +++++++++++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/apps/assets/const/host.py b/apps/assets/const/host.py
index 40b712c67..387a2e7e7 100644
--- a/apps/assets/const/host.py
+++ b/apps/assets/const/host.py
@@ -16,14 +16,8 @@ class HostTypes(BaseType):
                 'domain_enabled': True,
                 'su_enabled': True,
                 'su_methods': [
-                    {
-                      'name': 'sudo su',
-                      'id': 'sudo su'
-                    },
-                    {
-                        'name': 'su -',
-                        'id': 'su -'
-                    }
+                    {'name': 'sudo su', 'id': 'sudo su'},
+                    {'name': 'su -', 'id': 'su -'}
                 ],
             },
             cls.WINDOWS: {
@@ -72,10 +66,10 @@ class HostTypes(BaseType):
             ],
             cls.WINDOWS: [
                 {'name': 'Windows'},
-                {'name': 'Windows-TLS', 'protocol_settings': {
+                {'name': 'Windows-TLS', 'protocols_setting': {
                     'rdp': {'security': 'tls'},
                 }},
-                {'name': 'Windows-RDP', 'protocol_settings': {
+                {'name': 'Windows-RDP', 'protocols_setting': {
                     'rdp': {'security': 'rdp'},
                 }}
             ],
diff --git a/apps/assets/const/types.py b/apps/assets/const/types.py
index 265e01e92..f0eece68b 100644
--- a/apps/assets/const/types.py
+++ b/apps/assets/const/types.py
@@ -1,3 +1,5 @@
+from copy import deepcopy
+
 from common.db.models import IncludesTextChoicesMeta, ChoicesMixin
 from common.tree import TreeNode
 
@@ -161,14 +163,47 @@ class AllTypes(ChoicesMixin, metaclass=IncludesTextChoicesMeta):
 
     @classmethod
     def create_or_update_internal_platforms(cls):
-        from assets.models import Platform
+        from assets.models import Platform, PlatformAutomation, PlatformProtocol
+        print("Create internal platforms")
         for category, type_cls in cls.category_types():
+            print("## Category: {}".format(category.label))
             data = type_cls.internal_platforms()
+
             for tp, platform_datas in data.items():
+                print("  >> Type: {}".format(tp.label))
+                default_platform_data = cls.get_type_default_platform(category, tp)
+                default_automation = default_platform_data.pop('automation', {})
+                default_protocols = default_platform_data.pop('protocols', [])
+
                 for d in platform_datas:
-                    platform_data = {
-                        **d, 'category': category.value, 'type': tp.value,
-                        'internal': True, 'charset': 'utf-8'
-                    }
+                    name = d['name']
+                    print("    - Platform: {}".format(name))
+                    _automation = d.pop('automation', {})
+                    _protocols = d.pop('_protocols', [])
+                    _protocols_setting = d.pop('protocols_setting', {})
 
+                    protocols_data = deepcopy(default_protocols)
+                    if _protocols:
+                        protocols_data = [p for p in protocols_data if p['name'] in _protocols]
+                    for p in protocols_data:
+                        p['setting'] = {**_protocols_setting.get(p['name'], {}), **p.get('setting', {})}
+
+                    platform_data = {**default_platform_data, **d}
+                    automation_data = {**default_automation, **_automation}
+                    platform, created = Platform.objects.update_or_create(
+                        defaults=platform_data, name=platform_data['name']
+                    )
+
+                    if not platform.automation:
+                        automation = PlatformAutomation.objects.create()
+                        platform.automation = automation
+                        platform.save()
+                    else:
+                        automation = platform.automation
+                    for k, v in automation_data.items():
+                        setattr(automation, k, v)
+                    automation.save()
+
+                    platform.protocols.all().delete()
+                    [PlatformProtocol.objects.create(**p, platform=platform) for p in protocols_data]