From fa0382fc5ea919fdb3063eeb0f4526b260ae8b97 Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Thu, 1 Dec 2022 18:22:41 +0800
Subject: [PATCH] perf: gateway manager

---
 apps/assets/models/domain.py      | 89 ++++++++++---------------------
 apps/assets/serializers/domain.py | 10 ----
 2 files changed, 27 insertions(+), 72 deletions(-)

diff --git a/apps/assets/models/domain.py b/apps/assets/models/domain.py
index daa08a343..9580b17ba 100644
--- a/apps/assets/models/domain.py
+++ b/apps/assets/models/domain.py
@@ -9,12 +9,13 @@ from django.utils.translation import ugettext_lazy as _
 
 from common.utils import get_logger, lazyproperty
 from orgs.mixins.models import OrgModelMixin
-from assets.models import Host
+from assets.models import Host, Platform
 from assets.const import GATEWAY_NAME
+from orgs.mixins.models import OrgManager
 
 logger = get_logger(__file__)
 
-__all__ = ['Domain']
+__all__ = ['Domain', 'Gateway']
 
 
 class Domain(OrgModelMixin):
@@ -33,10 +34,7 @@ class Domain(OrgModelMixin):
 
     @classmethod
     def get_gateway_queryset(cls):
-        queryset = Host.objects.filter(
-            platform__name=GATEWAY_NAME
-        )
-        return queryset
+        return Gateway.objects.all()
 
     @lazyproperty
     def gateways(self):
@@ -55,63 +53,30 @@ class Domain(OrgModelMixin):
             return random.choice(self.gateways)
 
 
+class GatewayManager(OrgManager):
+    def get_queryset(self):
+        queryset = super().get_queryset()
+        queryset = queryset.filter(platform__name=GATEWAY_NAME)
+        return queryset
+
+    def bulk_create(self, objs, batch_size=None, ignore_conflicts=False):
+        platform = Gateway().default_platform
+        for obj in objs:
+            obj.platform_id = platform.id
+        return super().bulk_create(objs, batch_size, ignore_conflicts)
+
+
 class Gateway(Host):
+    objects = GatewayManager()
+
     class Meta:
         proxy = True
 
-    def test_connective(self, local_port=None):
-        if local_port is None:
-            local_port = self.port
-
-        client = paramiko.SSHClient()
-        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        proxy = paramiko.SSHClient()
-        proxy.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-
-        try:
-            proxy.connect(self.ip, port=self.port,
-                          username=self.username,
-                          password=self.password,
-                          pkey=self.private_key_obj)
-        except(paramiko.AuthenticationException,
-               paramiko.BadAuthenticationType,
-               paramiko.SSHException,
-               paramiko.ChannelException,
-               paramiko.ssh_exception.NoValidConnectionsError,
-               socket.gaierror) as e:
-            err = str(e)
-            if err.startswith('[Errno None] Unable to connect to port'):
-                err = _('Unable to connect to port {port} on {ip}')
-                err = err.format(port=self.port, ip=self.ip)
-            elif err == 'Authentication failed.':
-                err = _('Authentication failed')
-            elif err == 'Connect failed':
-                err = _('Connect failed')
-            self.is_connective = False
-            return False, err
-
-        try:
-            sock = proxy.get_transport().open_channel(
-                'direct-tcpip', ('127.0.0.1', local_port), ('127.0.0.1', 0)
-            )
-            client.connect("127.0.0.1", port=local_port,
-                           username=self.username,
-                           password=self.password,
-                           key_filename=self.private_key_file,
-                           sock=sock,
-                           timeout=5)
-        except (paramiko.SSHException,
-                paramiko.ssh_exception.SSHException,
-                paramiko.ChannelException,
-                paramiko.AuthenticationException,
-                TimeoutError) as e:
-
-            err = getattr(e, 'text', str(e))
-            if err == 'Connect failed':
-                err = _('Connect failed')
-            self.is_connective = False
-            return False, err
-        finally:
-            client.close()
-        self.is_connective = True
-        return True, None
+    @lazyproperty
+    def default_platform(self):
+        return Platform.objects.get(name=GATEWAY_NAME, internal=True)
+
+    def save(self, *args, **kwargs):
+        platform = self.default_platform
+        self.platform_id = platform.id
+        return super().save(*args, **kwargs)
diff --git a/apps/assets/serializers/domain.py b/apps/assets/serializers/domain.py
index 3e2b4889f..f17aa19ee 100644
--- a/apps/assets/serializers/domain.py
+++ b/apps/assets/serializers/domain.py
@@ -93,15 +93,6 @@ class GatewaySerializer(BulkOrgResourceModelSerializer, WritableNestedModelSeria
         validated_data.pop('passphrase', None)
         return username, password, private_key
 
-    @staticmethod
-    def generate_default_data():
-        platform = Platform.objects.get(name=GATEWAY_NAME, internal=True)
-        # node = Node.objects.all().order_by('date_created').first()
-        data = {
-            'platform': platform,
-        }
-        return data
-
     @staticmethod
     def create_accounts(instance, username, password, private_key):
         account_name = f'{instance.name}-{_("Gateway")}'
@@ -135,7 +126,6 @@ class GatewaySerializer(BulkOrgResourceModelSerializer, WritableNestedModelSeria
 
     def create(self, validated_data):
         auth_fields = self.clean_auth_fields(validated_data)
-        validated_data.update(self.generate_default_data())
         instance = super().create(validated_data)
         self.create_accounts(instance, *auth_fields)
         return instance