[Update] 优化授权协议, 支持vnc (#2220)

* [Update] 优化授权协议, 支持vnc

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'),
+        ),
+    ]

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)

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'

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