diff --git a/apps/assets/forms/asset.py b/apps/assets/forms/asset.py index 516ce4afc..39ad7713e 100644 --- a/apps/assets/forms/asset.py +++ b/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(), diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index 6de5c2059..073ad3b32 100644 --- a/apps/assets/models/asset.py +++ b/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 diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py index 74156daa6..7b5b9dd48 100644 --- a/apps/assets/models/node.py +++ b/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() + + diff --git a/apps/orgs/mixins.py b/apps/orgs/mixins.py index d9f55d301..d7d4b0d23 100644 --- a/apps/orgs/mixins.py +++ b/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):