[Update] 完成基本框架

apps/assets/models/asset.py

@@ -13,7 +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
+from orgs.mixins import OrgModelMixin,OrgManager
 
 __all__ = ['Asset']
 logger = logging.getLogger(__name__)
@@ -37,7 +37,7 @@ def default_node():
         return None
 
 
-class AssetQuerySet(OrgQuerySet):
+class AssetQuerySet(models.QuerySet):
     def active(self):
         return self.filter(is_active=True)
 

apps/common/mixins.py

@@ -120,7 +120,7 @@ class AdminUserRequiredMixin(UserPassesTestMixin):
     def test_func(self):
         if not self.request.user.is_authenticated:
             return False
-        elif not self.request.user.is_superuser:
+        elif not self.request.user:
             self.raise_exception = True
             return False
         return True

apps/jumpserver/views.py

@@ -10,9 +10,10 @@ from django.shortcuts import redirect
 from users.models import User
 from assets.models import Asset
 from terminal.models import Session
+from orgs.mixins import OrgViewGenericMixin
 
 
-class IndexView(LoginRequiredMixin, TemplateView):
+class IndexView(LoginRequiredMixin, OrgViewGenericMixin, TemplateView):
     template_name = 'index.html'
 
     session_week = None

apps/orgs/middleware.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 #
 
-from .utils import get_org_from_request
+from .utils import get_org_from_request, set_current_org
 
 
 class OrgMiddleware:
@@ -11,5 +11,6 @@ class OrgMiddleware:
     def __call__(self, request):
         org = get_org_from_request(request)
         request.current_org = org
+        set_current_org(org)
         response = self.get_response(request)
         return response

apps/orgs/mixins.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 #
 from django.db import models
+from django.shortcuts import redirect
+from django.contrib.auth import get_user_model
 
 from common.utils import get_logger
 from .utils import get_current_org, get_model_by_db_table
@@ -8,24 +10,27 @@ 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)
+__all__ = ['OrgManager', 'OrgViewGenericMixin', 'OrgModelMixin']
 
 
-class OrgManager(OrgQuerySet.as_manager().__class__):
+class OrgManager(models.Manager):
     def get_queryset(self):
         current_org = get_current_org()
+        user_model = get_user_model()
         kwargs = {}
 
+        print("Get queryset ")
+        print(self.model)
+        print(current_org)
+
         if not current_org:
             kwargs['id'] = None
-        elif current_org.is_real:
+        elif issubclass(self.model, user_model):
+            kwargs['orgs'] = current_org
+        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)
 
@@ -55,11 +60,25 @@ class OrgModelMixin(models.Model):
 
     def save(self, force_insert=False, force_update=False, using=None,
              update_fields=None):
+        user_model = get_user_model()
         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)
+        instance = super().save(
+            force_insert=force_insert, force_update=force_update,
+            using=using, update_fields=update_fields
+        )
+        if isinstance(instance, user_model):
+            instance.orgs.add(current_org)
+        return instance
 
     class Meta:
         abstract = True
+
+
+class OrgViewGenericMixin:
+    def dispatch(self, request, *args, **kwargs):
+        current_org = get_current_org()
+        if not current_org:
+            return redirect('orgs:switch-a-org')
+        return super().dispatch(request, *args, **kwargs)

apps/orgs/models.py

@@ -63,7 +63,7 @@ class Organization(models.Model):
         pass
 
     def is_real(self):
-        return len(str(self.id)) == 32
+        return len(str(self.id)) == 36
 
     @classmethod
     def get_user_admin_orgs(cls, user):

apps/orgs/urls/views_urls.py

@@ -8,6 +8,6 @@ app_name = 'orgs'
 
 
 urlpatterns = [
-    url(r'^(?P<pk>.*)/switch/$', views.SwitchOrgView.as_view(), name='org-switch')
+    url(r'^(?P<pk>.*)/switch/$', views.SwitchOrgView.as_view(), name='org-switch'),
+    url(r'^switch-a-org/$', views.SwitchToAOrgView.as_view(), name='switch-a-org')
 ]
-

apps/orgs/views.py

@@ -1,6 +1,7 @@
-from django.shortcuts import redirect
+from django.shortcuts import redirect, reverse
+from django.http import HttpResponseForbidden
 
-from django.views.generic import DetailView
+from django.views.generic import DetailView, View
 
 from .models import Organization
 
@@ -14,3 +15,16 @@ class SwitchOrgView(DetailView):
         self.object = Organization.get_instance(pk)
         request.session['oid'] = self.object.id.__str__()
         return redirect('index')
+
+
+class SwitchToAOrgView(View):
+    def get(self, request, *args, **kwargs):
+        admin_orgs = Organization.get_user_admin_orgs(request.user)
+        if not admin_orgs:
+            return HttpResponseForbidden()
+        default_org = Organization.default()
+        if default_org in admin_orgs:
+            redirect_org = default_org
+        else:
+            redirect_org = admin_orgs[0]
+        return redirect(reverse('orgs:org-switch', kwargs={'pk': redirect_org.id}))

apps/templates/_nav.html

@@ -1,5 +1,5 @@
 {% load i18n %}
-{% if ADMIN_ORGS %}
+{% if ADMIN_ORGS and ADMIN_ORGS|length > 1 %}
 <li id="org">
     <a class="dropdown-toggle" data-toggle="dropdown" aria-expanded="false">
         <i class="fa fa-star" style="width: 14px"></i> <span class="nav-label">{{ CURRENT_ORG.name }}</span>
@@ -7,8 +7,8 @@
     </a>
     <ul class="dropdown-menu">
         {% for org in ADMIN_ORGS %}
-            {% if org.is_default %}
-                <li><a class="org-dropdown" href="{% url 'orgs:org-switch' pk=org.id %}" data-id="{{ org.id }}">{{ org.name }}</a></li>
+            {% if org.id != CURRENT_ORG.id %}
+            <li><a class="org-dropdown" href="{% url 'orgs:org-switch' pk=org.id %}" data-id="{{ org.id }}">{{ org.name }}</a></li>
             {% endif %}
         {% endfor %}
     </ul>
@@ -17,8 +17,8 @@
 {% endif %}
 <li id="index">
     <a href="{% url 'index' %}">
-        <i class="fa fa-dashboard" style="width: 14px"></i> <span class="nav-label">{% trans 'Dashboard' %}</span><span
-            class="label label-info pull-right"></span>
+        <i class="fa fa-dashboard" style="width: 14px"></i> <span class="nav-label">{% trans 'Dashboard' %}</span>
+        <span class="label label-info pull-right"></span>
     </a>
 </li>
 <li id="users">
@@ -96,4 +96,11 @@
     <a href="{% url 'settings:basic-setting' %}">
         <i class="fa fa-gears"></i> <span class="nav-label">{% trans 'Settings' %}</span><span class="label label-info pull-right"></span>
     </a>
-</li>
+</li>
+
+<script>
+$(document).ready(function () {
+    var current_org = '{{ CURRENT_ORG.name }}';
+    console.log(current_org);
+})
+</script>

apps/users/api.py

@@ -20,6 +20,7 @@ from .permissions import IsSuperUser, IsValidUser, IsCurrentUserOrReadOnly, \
     IsSuperUserOrAppUser
 from .utils import check_user_valid, generate_token, get_login_ip, \
     check_otp_code, set_user_login_failed_count_to_cache, is_block_login
+from orgs.utils import get_current_org
 from common.mixins import IDInFilterMixin
 from common.utils import get_logger
 
@@ -33,6 +34,15 @@ class UserViewSet(IDInFilterMixin, BulkModelViewSet):
     permission_classes = (IsSuperUser,)
     filter_fields = ('username', 'email', 'name', 'id')
 
+    def get_queryset(self):
+        queryset = super().get_queryset()
+        current_org = get_current_org()
+        if current_org.is_real():
+            queryset = queryset.filter(orgs=current_org)
+        elif current_org.is_default():
+            queryset = queryset.filter(orgs=None)
+        return queryset
+
     def get_permissions(self):
         if self.action == "retrieve":
             self.permission_classes = (IsSuperUserOrAppUser,)

apps/users/models/user.py

@@ -6,7 +6,7 @@ from collections import OrderedDict
 
 from django.conf import settings
 from django.contrib.auth.hashers import make_password
-from django.contrib.auth.models import AbstractUser
+from django.contrib.auth.models import AbstractUser, UserManager
 from django.core import signing
 from django.db import models
 from django.utils.translation import ugettext_lazy as _
@@ -15,6 +15,7 @@ from django.shortcuts import reverse
 
 from common.utils import get_signer, date_expired_default
 from common.models import Setting
+from orgs.utils import get_current_org
 
 
 __all__ = ['User']
@@ -186,6 +187,18 @@ class User(AbstractUser):
         else:
             self.role = 'User'
 
+    @property
+    def admin_orgs(self):
+        from orgs.models import Organization
+        return Organization.get_user_admin_orgs(self)
+
+    @property
+    def is_org_admin(self):
+        if self.is_superuser or self.admin_orgs:
+            return True
+        else:
+            return False
+
     @property
     def is_app(self):
         return self.role == 'App'
@@ -207,8 +220,11 @@ class User(AbstractUser):
         if self.username == 'admin':
             self.role = 'Admin'
             self.is_active = True
-
-        super().save(*args, **kwargs)
+        instance = super().save(*args, **kwargs)
+        current_org = get_current_org()
+        if current_org and current_org.is_real():
+            instance.orgs.add(current_org)
+        return instance
 
     @property
     def private_token(self):