[Update] 修改资产

apps/assets/forms/asset.py

@@ -3,14 +3,17 @@
 from django import forms
 from django.utils.translation import gettext_lazy as _
 
-from ..models import Asset, AdminUser
 from common.utils import get_logger
+from orgs.mixins import OrgModelForm
+
+from ..models import Asset, AdminUser
+
 
 logger = get_logger(__file__)
 __all__ = ['AssetCreateForm', 'AssetUpdateForm', 'AssetBulkUpdateForm']
 
 
-class AssetCreateForm(forms.ModelForm):
+class AssetCreateForm(OrgModelForm):
     class Meta:
         model = Asset
         fields = [
@@ -50,7 +53,7 @@ class AssetCreateForm(forms.ModelForm):
         }
 
 
-class AssetUpdateForm(forms.ModelForm):
+class AssetUpdateForm(OrgModelForm):
     class Meta:
         model = Asset
         fields = [
@@ -90,7 +93,7 @@ class AssetUpdateForm(forms.ModelForm):
         }
 
 
-class AssetBulkUpdateForm(forms.ModelForm):
+class AssetBulkUpdateForm(OrgModelForm):
     assets = forms.ModelMultipleChoiceField(
         required=True, help_text='* required',
         label=_('Select assets'), queryset=Asset.objects.all(),

apps/assets/models/asset.py

@@ -71,16 +71,11 @@ class Asset(OrgModelMixin):
     )
 
     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, unique=True,
-                                verbose_name=_('Hostname'))
-    protocol = models.CharField(max_length=128, default=SSH_PROTOCOL,
-                                choices=PROTOCOL_CHOICES,
-                                verbose_name=_('Protocol'))
+    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'))
     port = models.IntegerField(default=22, verbose_name=_('Port'))
-    platform = models.CharField(max_length=128, choices=PLATFORM_CHOICES,
-                                default='Linux', verbose_name=_('Platform'))
+    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)
@@ -94,11 +89,8 @@ class Asset(OrgModelMixin):
                                    null=True, verbose_name=_("Admin user"))
 
     # Some information
-    public_ip = models.GenericIPAddressField(max_length=32, blank=True,
-                                             null=True,
-                                             verbose_name=_('Public IP'))
-    number = models.CharField(max_length=32, null=True, blank=True,
-                              verbose_name=_('Asset number'))
+    public_ip = models.GenericIPAddressField(max_length=32, blank=True, null=True, verbose_name=_('Public IP'))
+    number = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Asset number'))
 
     # Collect
     vendor = models.CharField(max_length=64, null=True, blank=True,
@@ -233,7 +225,7 @@ class Asset(OrgModelMixin):
         return data
 
     class Meta:
-        unique_together = ('ip', 'port')
+        unique_together = ('org', 'hostname')
         verbose_name = _("Asset")
 
     @classmethod

apps/assets/models/node.py

@@ -7,7 +7,8 @@ from django.db.models import Q
 from django.utils.translation import ugettext_lazy as _
 
 from orgs.mixins import OrgModelMixin
-from common.utils import with_cache
+from orgs.utils import get_current_org, set_current_org
+from orgs.models import Organization
 
 __all__ = ['Node']
 
@@ -119,7 +120,11 @@ class Node(OrgModelMixin):
         return self.get_all_assets().valid()
 
     def is_root(self):
-        return self.key == '0'
+        root = self.__class__.root()
+        if self == root:
+            return True
+        else:
+            return False
 
     @property
     def parent(self):
@@ -148,8 +153,8 @@ class Node(OrgModelMixin):
 
     def get_ancestor(self, with_self=False):
         if self.is_root():
-            ancestor = self.__class__.objects.filter(key='0')
-            return ancestor
+            root = self.__class__.root()
+            return [root]
 
         _key = self.key.split(':')
         if not with_self:
@@ -163,11 +168,26 @@ class Node(OrgModelMixin):
         ).order_by('key')
         return ancestor
 
+    @classmethod
+    def create_root_node(cls):
+        with transaction.atomic():
+            org = get_current_org()
+            set_current_org(Organization.root())
+            org_nodes_roots = cls.objects.filter(key__regex=r'^[0-9]+$')
+            org_nodes_roots_keys = org_nodes_roots.values_list('key', flat=True)
+            max_value = max([int(k) for k in org_nodes_roots_keys]) if org_nodes_roots_keys else 0
+            set_current_org(org)
+            root = cls.objects.create(key=max_value+1, value=org.name)
+            return root
+
     @classmethod
     def root(cls):
-        obj, created = cls.objects.get_or_create(
-            key='0', defaults={"key": '0', 'value': "ROOT"}
-        )
-        print(obj)
-        return obj
+        root = cls.objects.filter(key__regex=r'^[0-9]+$')
+        print("GET ROOT NODE")
+        if len(root) == 1:
+            return root.get()
+        else:
+            return cls.create_root_node()
+
+
 

apps/orgs/mixins.py

@@ -18,28 +18,17 @@ __all__ = [
 
 
 class OrgManager(models.Manager):
-    def __init__(self, *args, **kwargs):
-        print("INit manager")
-        super().__init__(*args, **kwargs)
-
     def get_queryset(self):
-        print("GET CURR")
         current_org = get_current_org()
         kwargs = {}
 
-        print("Get queryset ")
-        print(current_org)
-
-        print(self.model)
         if not current_org:
-            pass
+            kwargs['id'] = None
         elif current_org.is_real():
             kwargs['org'] = current_org
         elif current_org.is_default():
             kwargs['org'] = None
         queryset = super().get_queryset().filter(**kwargs)
-        print(kwargs)
-        print(queryset)
         return queryset
 
     def all(self):