From fe7c3c29ad1fd23e1662daac4826e4d7b3d25f96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=80=81=E5=B9=BF?= <ibuler@qq.com>
Date: Wed, 19 Dec 2018 17:03:10 +0800
Subject: [PATCH] =?UTF-8?q?[Update]=20=E4=BC=98=E5=8C=96=E6=8E=88=E6=9D=83?=
 =?UTF-8?q?=E5=8D=8F=E8=AE=AE,=20=E6=94=AF=E6=8C=81vnc=20(#2220)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* [Update] 优化授权协议, 支持vnc
---
 .../migrations/0024_auto_20181219_1614.py     | 23 +++++++++++++++++++
 apps/assets/models/asset.py                   | 16 +++++++------
 apps/assets/models/user.py                    |  2 ++
 apps/perms/utils.py                           |  8 +++++--
 4 files changed, 40 insertions(+), 9 deletions(-)
 create mode 100644 apps/assets/migrations/0024_auto_20181219_1614.py

diff --git a/apps/assets/migrations/0024_auto_20181219_1614.py b/apps/assets/migrations/0024_auto_20181219_1614.py
new file mode 100644
index 000000000..5e6a6c4e3
--- /dev/null
+++ b/apps/assets/migrations/0024_auto_20181219_1614.py
@@ -0,0 +1,23 @@
+# Generated by Django 2.1.4 on 2018-12-19 08:14
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('assets', '0023_auto_20181016_1650'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='asset',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'ssh'), ('rdp', 'rdp'), ('telnet', 'telnet (beta)'), ('vnc', 'vnc')], default='ssh', max_length=128, verbose_name='Protocol'),
+        ),
+        migrations.AlterField(
+            model_name='systemuser',
+            name='protocol',
+            field=models.CharField(choices=[('ssh', 'ssh'), ('rdp', 'rdp'), ('telnet', 'telnet (beta)'), ('vnc', 'vnc')], default='ssh', max_length=16, verbose_name='Protocol'),
+        ),
+    ]
diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py
index ccff0fff5..9ab5aae86 100644
--- a/apps/assets/models/asset.py
+++ b/apps/assets/models/asset.py
@@ -59,19 +59,21 @@ class Asset(OrgModelMixin):
         ('Other', 'Other'),
     )
 
-    SSH_PROTOCOL = 'ssh'
-    RDP_PROTOCOL = 'rdp'
-    TELNET_PROTOCOL = 'telnet'
+    PROTOCOL_SSH = 'ssh'
+    PROTOCOL_RDP = 'rdp'
+    PROTOCOL_TELNET = 'telnet'
+    PROTOCOL_VNC = 'vnc'
     PROTOCOL_CHOICES = (
-        (SSH_PROTOCOL, 'ssh'),
-        (RDP_PROTOCOL, 'rdp'),
-        (TELNET_PROTOCOL, 'telnet (beta)'),
+        (PROTOCOL_SSH, 'ssh'),
+        (PROTOCOL_RDP, 'rdp'),
+        (PROTOCOL_TELNET, 'telnet (beta)'),
+        (PROTOCOL_VNC, 'vnc'),
     )
 
     id = models.UUIDField(default=uuid.uuid4, primary_key=True)
     ip = models.GenericIPAddressField(max_length=32, verbose_name=_('IP'), db_index=True)
     hostname = models.CharField(max_length=128, verbose_name=_('Hostname'))
-    protocol = models.CharField(max_length=128, default=SSH_PROTOCOL, choices=PROTOCOL_CHOICES, verbose_name=_('Protocol'))
+    protocol = models.CharField(max_length=128, default=PROTOCOL_SSH, choices=PROTOCOL_CHOICES, verbose_name=_('Protocol'))
     port = models.IntegerField(default=22, verbose_name=_('Port'))
     platform = models.CharField(max_length=128, choices=PLATFORM_CHOICES, default='Linux', verbose_name=_('Platform'))
     domain = models.ForeignKey("assets.Domain", null=True, blank=True, related_name='assets', verbose_name=_("Domain"), on_delete=models.SET_NULL)
diff --git a/apps/assets/models/user.py b/apps/assets/models/user.py
index 147687471..1e661d631 100644
--- a/apps/assets/models/user.py
+++ b/apps/assets/models/user.py
@@ -115,10 +115,12 @@ class SystemUser(AssetUser):
     PROTOCOL_SSH = 'ssh'
     PROTOCOL_RDP = 'rdp'
     PROTOCOL_TELNET = 'telnet'
+    PROTOCOL_VNC = 'vnc'
     PROTOCOL_CHOICES = (
         (PROTOCOL_SSH, 'ssh'),
         (PROTOCOL_RDP, 'rdp'),
         (PROTOCOL_TELNET, 'telnet (beta)'),
+        (PROTOCOL_VNC, 'vnc'),
     )
 
     LOGIN_AUTO = 'auto'
diff --git a/apps/perms/utils.py b/apps/perms/utils.py
index 8b79175aa..ef8e3d09a 100644
--- a/apps/perms/utils.py
+++ b/apps/perms/utils.py
@@ -136,7 +136,9 @@ class AssetPermissionUtil:
         permissions = self.permissions.prefetch_related('assets', 'system_users')
         for perm in permissions:
             for asset in perm.assets.all().valid().prefetch_related('nodes'):
-                assets[asset].update(perm.system_users.all())
+                assets[asset].update(
+                    perm.system_users.filter(protocol=asset.protocol)
+                )
         return assets
 
     def get_assets(self):
@@ -147,7 +149,9 @@ class AssetPermissionUtil:
         for node, system_users in nodes.items():
             _assets = node.get_all_assets().valid().prefetch_related('nodes')
             for asset in _assets:
-                assets[asset].update(system_users)
+                assets[asset].update(
+                    [s for s in system_users if s.protocol == asset.protocol]
+                )
         self._assets = assets
         return self._assets