[Update] 添加org

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

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])

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'))

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"))

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',
             ],
         },
     },

apps/orgs/admin.py

+from django.contrib import admin
+
+# Register your models here.

apps/orgs/apps.py

+from django.apps import AppConfig
+
+
+class OrgsConfig(AppConfig):
+    name = 'orgs'

apps/orgs/context_processor.py

+# -*- 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
+

apps/orgs/middleware.py

+# -*- 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

apps/orgs/mixins.py

+# -*- 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

apps/orgs/models.py

+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

apps/orgs/tests.py

+from django.test import TestCase
+
+# Create your tests here.

apps/orgs/utils.py

+# -*- 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))

apps/orgs/views.py

+from django.shortcuts import render
+
+# Create your views here.