diff --git a/apps/assets/models/asset.py b/apps/assets/models/asset.py index c344fa82f..6de5c2059 100644 --- a/apps/assets/models/asset.py +++ b/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) diff --git a/apps/common/mixins.py b/apps/common/mixins.py index 243ee93c6..29e0ad876 100644 --- a/apps/common/mixins.py +++ b/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 diff --git a/apps/jumpserver/views.py b/apps/jumpserver/views.py index c6a2dc4f3..1ad2bb20f 100644 --- a/apps/jumpserver/views.py +++ b/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 diff --git a/apps/orgs/middleware.py b/apps/orgs/middleware.py index 527797857..04a2d6f7a 100644 --- a/apps/orgs/middleware.py +++ b/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 diff --git a/apps/orgs/mixins.py b/apps/orgs/mixins.py index 5bcf75f2e..f5d59e686 100644 --- a/apps/orgs/mixins.py +++ b/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) diff --git a/apps/orgs/models.py b/apps/orgs/models.py index 3f43c6bc6..afdf2ef10 100644 --- a/apps/orgs/models.py +++ b/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): diff --git a/apps/orgs/urls/views_urls.py b/apps/orgs/urls/views_urls.py index a0f21cf32..e3408e741 100644 --- a/apps/orgs/urls/views_urls.py +++ b/apps/orgs/urls/views_urls.py @@ -8,6 +8,6 @@ app_name = 'orgs' urlpatterns = [ - url(r'^(?P.*)/switch/$', views.SwitchOrgView.as_view(), name='org-switch') + url(r'^(?P.*)/switch/$', views.SwitchOrgView.as_view(), name='org-switch'), + url(r'^switch-a-org/$', views.SwitchToAOrgView.as_view(), name='switch-a-org') ] - diff --git a/apps/orgs/views.py b/apps/orgs/views.py index ec2136716..374f34812 100644 --- a/apps/orgs/views.py +++ b/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})) diff --git a/apps/templates/_nav.html b/apps/templates/_nav.html index 60f38b92c..ed2e67921 100644 --- a/apps/templates/_nav.html +++ b/apps/templates/_nav.html @@ -1,5 +1,5 @@ {% load i18n %} -{% if ADMIN_ORGS %} +{% if ADMIN_ORGS and ADMIN_ORGS|length > 1 %}
  • @@ -17,8 +17,8 @@ {% endif %}
  • - {% trans 'Dashboard' %} + {% trans 'Dashboard' %} +
  • @@ -96,4 +96,11 @@ {% trans 'Settings' %} -
  • \ No newline at end of file + + + \ No newline at end of file diff --git a/apps/users/api.py b/apps/users/api.py index c23112384..10a12ee66 100644 --- a/apps/users/api.py +++ b/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,) diff --git a/apps/users/models/user.py b/apps/users/models/user.py index d69151f55..30054f137 100644 --- a/apps/users/models/user.py +++ b/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):