diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index 7a2b3fe57..c344fa82f 100644 --- a/apps/assets/models/asset.py +++ b/apps/assets/models/asset.py @@ -13,6 +13,7 @@ from django.core.cache import cache from ..const import ASSET_ADMIN_CONN_CACHE_KEY from .user import AdminUser, SystemUser +from orgs.mixins import OrgModelMixin, OrgQuerySet, OrgManager __all__ = ['Asset'] logger = logging.getLogger(__name__) @@ -36,7 +37,7 @@ def default_node(): return None -class AssetQuerySet(models.QuerySet): +class AssetQuerySet(OrgQuerySet): def active(self): return self.filter(is_active=True) @@ -44,12 +45,11 @@ class AssetQuerySet(models.QuerySet): return self.active() -class AssetManager(models.Manager): - def get_queryset(self): - return AssetQuerySet(self.model, using=self._db) +class AssetManager(OrgManager): + pass -class Asset(models.Model): +class Asset(OrgModelMixin): # Important PLATFORM_CHOICES = ( ('Linux', 'Linux'), diff --git a/apps/assets/models/base.py b/apps/assets/models/base.py index 908e6b647..3698f3a05 100644 --- a/apps/assets/models/base.py +++ b/apps/assets/models/base.py @@ -11,12 +11,13 @@ from django.conf import settings from common.utils import get_signer, ssh_key_string_to_obj, ssh_key_gen from common.validators import alphanumeric +from orgs.mixins import OrgModelMixin from .utils import private_key_validator signer = get_signer() -class AssetUser(models.Model): +class AssetUser(OrgModelMixin): id = models.UUIDField(default=uuid.uuid4, primary_key=True) name = models.CharField(max_length=128, unique=True, verbose_name=_('Name')) username = models.CharField(max_length=32, blank=True, verbose_name=_('Username'), validators=[alphanumeric]) diff --git a/apps/assets/models/domain.py b/apps/assets/models/domain.py index 6f29a0381..67283bd80 100644 --- a/apps/assets/models/domain.py +++ b/apps/assets/models/domain.py @@ -7,12 +7,13 @@ import random from django.db import models from django.utils.translation import ugettext_lazy as _ +from orgs.mixins import OrgModelMixin from .base import AssetUser __all__ = ['Domain', 'Gateway'] -class Domain(models.Model): +class Domain(OrgModelMixin): id = models.UUIDField(default=uuid.uuid4, primary_key=True) name = models.CharField(max_length=128, unique=True, verbose_name=_('Name')) comment = models.TextField(blank=True, verbose_name=_('Comment')) diff --git a/apps/assets/models/node.py b/apps/assets/models/node.py index 4f4f9ad8b..74156daa6 100644 --- a/apps/assets/models/node.py +++ b/apps/assets/models/node.py @@ -5,12 +5,14 @@ import uuid from django.db import models, transaction 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 __all__ = ['Node'] -class Node(models.Model): +class Node(OrgModelMixin): id = models.UUIDField(default=uuid.uuid4, primary_key=True) key = models.CharField(unique=True, max_length=64, verbose_name=_("Key")) # '1:1:1:1' value = models.CharField(max_length=128, verbose_name=_("Value")) diff --git a/apps/jumpserver/settings.py b/apps/jumpserver/settings.py index c0a536276..77a37839d 100644 --- a/apps/jumpserver/settings.py +++ b/apps/jumpserver/settings.py @@ -56,6 +56,7 @@ ALLOWED_HOSTS = CONFIG.ALLOWED_HOSTS or [] # Application definition INSTALLED_APPS = [ + 'orgs.apps.OrgsConfig', 'users.apps.UsersConfig', 'assets.apps.AssetsConfig', 'perms.apps.PermsConfig', @@ -87,6 +88,7 @@ MIDDLEWARE = [ 'django.middleware.clickjacking.XFrameOptionsMiddleware', 'jumpserver.middleware.TimezoneMiddleware', 'jumpserver.middleware.DemoMiddleware', + 'orgs.middleware.OrgMiddleware', ] ROOT_URLCONF = 'jumpserver.urls' @@ -107,6 +109,7 @@ TEMPLATES = [ 'django.template.context_processors.static', 'django.template.context_processors.request', 'django.template.context_processors.media', + 'orgs.context_processor.org_processor', ], }, }, diff --git a/apps/orgs/__init__.py b/apps/orgs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/orgs/admin.py b/apps/orgs/admin.py new file mode 100644 index 000000000..8c38f3f3d --- /dev/null +++ b/apps/orgs/admin.py @@ -0,0 +1,3 @@ +from django.contrib import admin + +# Register your models here. diff --git a/apps/orgs/apps.py b/apps/orgs/apps.py new file mode 100644 index 000000000..fdd6dc1b6 --- /dev/null +++ b/apps/orgs/apps.py @@ -0,0 +1,5 @@ +from django.apps import AppConfig + + +class OrgsConfig(AppConfig): + name = 'orgs' diff --git a/apps/orgs/context_processor.py b/apps/orgs/context_processor.py new file mode 100644 index 000000000..aa89d3a54 --- /dev/null +++ b/apps/orgs/context_processor.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +# + +from .utils import get_current_org +from .models import Organization + + +def org_processor(request): + context = { + 'ADMIN_ORGS': Organization.get_user_admin_orgs(request.user), + 'CURRENT_ORG': get_current_org(), + } + return context + diff --git a/apps/orgs/middleware.py b/apps/orgs/middleware.py new file mode 100644 index 000000000..527797857 --- /dev/null +++ b/apps/orgs/middleware.py @@ -0,0 +1,15 @@ +# -*- coding: utf-8 -*- +# + +from .utils import get_org_from_request + + +class OrgMiddleware: + def __init__(self, get_response): + self.get_response = get_response + + def __call__(self, request): + org = get_org_from_request(request) + request.current_org = org + response = self.get_response(request) + return response diff --git a/apps/orgs/mixins.py b/apps/orgs/mixins.py new file mode 100644 index 000000000..5bcf75f2e --- /dev/null +++ b/apps/orgs/mixins.py @@ -0,0 +1,65 @@ +# -*- coding: utf-8 -*- +# +from django.db import models + +from common.utils import get_logger +from .utils import get_current_org, get_model_by_db_table + +logger = get_logger(__file__) + + +class OrgQuerySet(models.QuerySet): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + +class OrgManager(OrgQuerySet.as_manager().__class__): + def get_queryset(self): + current_org = get_current_org() + kwargs = {} + + if not current_org: + kwargs['id'] = None + elif current_org.is_real: + kwargs['org'] = current_org + elif current_org.is_default(): + kwargs['org'] = None + print("GET QUWRYSET ") + print(kwargs) + return super().get_queryset().filter(**kwargs) + + +class OrgModelMixin(models.Model): + org = models.ForeignKey('orgs.Organization', on_delete=models.PROTECT, null=True) + + objects = OrgManager() + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def _do_update(self, base_qs, using, pk_val, values, update_fields, forced_update): + current_org = get_current_org() + + if current_org and current_org.is_real(): + kwargs = {'org': current_org} + base_qs = base_qs.filter(**kwargs) + else: + logger.warn( + 'Attempting to update %s instance "%s" without a current tenant ' + 'set. This may cause issues in a partitioned environment. ' + 'Recommend calling set_current_org() before performing this ' + 'operation.', self._meta.model.__name__, self + ) + return super()._do_update(base_qs, using, pk_val, values, update_fields, forced_update) + + def save(self, force_insert=False, force_update=False, using=None, + update_fields=None): + current_org = get_current_org() + if current_org and not current_org.is_real(): + self.org = current_org + return super().save(force_insert=force_insert, force_update=force_update, + using=using, update_fields=update_fields) + + class Meta: + abstract = True diff --git a/apps/orgs/models.py b/apps/orgs/models.py new file mode 100644 index 000000000..3f43c6bc6 --- /dev/null +++ b/apps/orgs/models.py @@ -0,0 +1,92 @@ +import uuid + +from django.db import models +from django.core.cache import cache +from django.utils.translation import ugettext_lazy as _ + + +class Organization(models.Model): + id = models.UUIDField(default=uuid.uuid4, primary_key=True) + name = models.CharField(max_length=128, verbose_name=_("Name")) + users = models.ManyToManyField('users.User', related_name='orgs', blank=True) + admins = models.ManyToManyField('users.User', related_name='admin_orgs', blank=True) + created_by = models.CharField(max_length=32, null=True, blank=True, verbose_name=_('Created by')) + date_created = models.DateTimeField(auto_now_add=True, null=True, blank=True, verbose_name=_('Date created')) + comment = models.TextField(max_length=128, default='', blank=True, verbose_name=_('Comment')) + + CACHE_PREFIX = 'JMS_ORG_{}' + ROOT_ID = 'ROOT' + DEFAULT_ID = 'DEFAULT' + + def __str__(self): + return self.name + + def set_to_cache(self): + key = self.CACHE_PREFIX.format(self.id) + cache.set(key, self, 3600) + + def expire_cache(self): + key = self.CACHE_PREFIX.format(self.id) + cache.set(key, self, 0) + + @classmethod + def get_instance_from_cache(cls, oid): + key = cls.CACHE_PREFIX.format(oid) + return cache.get(key, None) + + @classmethod + def get_instance(cls, oid, default=True): + cached = cls.get_instance_from_cache(oid) + if cached: + return cached + + if oid == cls.DEFAULT_ID: + return cls.default() + elif oid == cls.ROOT_ID: + return cls.root() + + try: + org = cls.objects.get(id=oid) + org.set_to_cache() + except cls.DoesNotExist: + org = cls.default() if default else None + return org + + def get_org_users(self): + from users.models import User + if self.is_default(): + return User.objects.all() + else: + return self.users.all() + + def get_org_admins(self): + pass + + def is_real(self): + return len(str(self.id)) == 32 + + @classmethod + def get_user_admin_orgs(cls, user): + admin_orgs = [] + if user.is_anonymous: + return admin_orgs + elif user.is_superuser: + admin_orgs = list(cls.objects.all()) + admin_orgs.append(cls.default()) + elif user.is_org_admin: + admin_orgs = user.admin_orgs.all() + return admin_orgs + + @classmethod + def default(cls): + return cls(id=cls.DEFAULT_ID, name="Default") + + @classmethod + def root(cls): + return cls(id=cls.ROOT_ID, name='Root') + + def is_default(self): + if self.id is self.DEFAULT_ID: + return True + else: + return False diff --git a/apps/orgs/tests.py b/apps/orgs/tests.py new file mode 100644 index 000000000..7ce503c2d --- /dev/null +++ b/apps/orgs/tests.py @@ -0,0 +1,3 @@ +from django.test import TestCase + +# Create your tests here. diff --git a/apps/orgs/utils.py b/apps/orgs/utils.py new file mode 100644 index 000000000..ee82fc01e --- /dev/null +++ b/apps/orgs/utils.py @@ -0,0 +1,44 @@ +# -*- coding: utf-8 -*- +# +import re +from django.apps import apps + +from .models import Organization + +try: + from threading import local +except ImportError: + from django.utils._threading_local import local + +_thread_locals = local() + + +def get_org_from_request(request): + oid = request.session.get("oid") + org = Organization.get_instance(oid) + return org + + +def get_current_request(): + return getattr(_thread_locals, 'request', None) + + +def get_current_org(): + return getattr(_thread_locals, 'current_org', None) + + +def get_current_user(): + return getattr(_thread_locals, 'user', None) + + +def set_current_org(org): + setattr(_thread_locals, 'current_org', org) + + +def get_model_by_db_table(db_table): + for model in apps.get_models(): + if model._meta.db_table == db_table: + return model + else: + # here you can do fallback logic if no model with db_table found + raise ValueError('No model found with db_table {}!'.format(db_table)) diff --git a/apps/orgs/views.py b/apps/orgs/views.py new file mode 100644 index 000000000..91ea44a21 --- /dev/null +++ b/apps/orgs/views.py @@ -0,0 +1,3 @@ +from django.shortcuts import render + +# Create your views here.