[Update] 修改gateway test connection (#2135)

apps/assets/api/domain.py

@@ -9,7 +9,6 @@ from django.views.generic.detail import SingleObjectMixin
 from common.utils import get_logger
 from common.permissions import IsOrgAdmin, IsAppUser, IsOrgAdminOrAppUser
 from ..models import Domain, Gateway
-from ..utils import test_gateway_connectability
 from .. import serializers
 
 
@@ -54,7 +53,7 @@ class GatewayTestConnectionApi(SingleObjectMixin, APIView):
 
     def get(self, request, *args, **kwargs):
         self.object = self.get_object(Gateway.objects.all())
-        ok, e = test_gateway_connectability(self.object)
+        ok, e = self.object.test_connective()
         if ok:
             return Response("ok")
         else:

apps/assets/forms/domain.py

@@ -36,6 +36,10 @@ class DomainForm(forms.ModelForm):
 
 
 class GatewayForm(PasswordAndKeyAuthForm, OrgModelForm):
+    protocol = forms.ChoiceField(
+        choices=[Gateway.PROTOCOL_CHOICES[0]],
+    )
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         password_field = self.fields.get('password')

apps/assets/models/domain.py

@@ -4,6 +4,8 @@
 import uuid
 import random
 
+import paramiko
+
 from django.db import models
 from django.utils.translation import ugettext_lazy as _
 
@@ -57,3 +59,37 @@ class Gateway(AssetUser):
     class Meta:
         unique_together = [('name', 'org_id')]
         verbose_name = _("Gateway")
+
+    def test_connective(self):
+        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) as e:
+            return False, str(e)
+
+        sock = proxy.get_transport().open_channel(
+            'direct-tcpip', ('127.0.0.1', self.port), ('127.0.0.1', 0)
+        )
+
+        try:
+            client.connect("127.0.0.1", port=self.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.AuthenticationException, TimeoutError) as e:
+            return False, str(e)
+        finally:
+            client.close()
+        return True, None

apps/assets/models/label.py

@@ -17,7 +17,8 @@ class Label(OrgModelMixin):
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
     name = models.CharField(max_length=128, verbose_name=_("Name"))
     value = models.CharField(max_length=128, verbose_name=_("Value"))
-    category = models.CharField(max_length=128, choices=CATEGORY_CHOICES, default=USER_CATEGORY, verbose_name=_("Category"))
+    category = models.CharField(max_length=128, choices=CATEGORY_CHOICES,
+                                default=USER_CATEGORY, verbose_name=_("Category"))
     is_active = models.BooleanField(default=True, verbose_name=_("Is active"))
     comment = models.TextField(blank=True, null=True, verbose_name=_("Comment"))
     date_created = models.DateTimeField(

apps/assets/serializers/domain.py

@@ -23,7 +23,6 @@ class DomainSerializer(serializers.ModelSerializer):
 
 
 class GatewaySerializer(serializers.ModelSerializer):
-
     class Meta:
         model = Gateway
         fields = [

apps/assets/utils.py

@@ -49,44 +49,3 @@ class LabelFilter:
             for kwargs in conditions:
                 queryset = queryset.filter(**kwargs)
         return queryset
-
-
-def test_gateway_connectability(gateway):
-    """
-    Test system cant connect his assets or not.
-    :param gateway:
-    :return:
-    """
-    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(gateway.ip, gateway.port,
-                      username=gateway.username,
-                      password=gateway.password,
-                      pkey=gateway.private_key_obj)
-    except(paramiko.AuthenticationException,
-           paramiko.BadAuthenticationType,
-           SSHException) as e:
-        return False, str(e)
-
-    sock = proxy.get_transport().open_channel(
-        'direct-tcpip', ('127.0.0.1', gateway.port), ('127.0.0.1', 0)
-    )
-
-    try:
-        client.connect("127.0.0.1", port=gateway.port,
-                       username=gateway.username,
-                       password=gateway.password,
-                       key_filename=gateway.private_key_file,
-                       sock=sock,
-                       timeout=5
-                       )
-    except (paramiko.SSHException, paramiko.ssh_exception.SSHException,
-            paramiko.AuthenticationException, TimeoutError) as e:
-        return False, str(e)
-    finally:
-        client.close()
-    return True, None