mirror of https://github.com/jumpserver/jumpserver
163 lines
4.8 KiB
Python
163 lines
4.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
#
|
||
from itertools import chain, groupby
|
||
from django.core.exceptions import MultipleObjectsReturned, ObjectDoesNotExist
|
||
|
||
from orgs.utils import current_org
|
||
from common.utils import get_logger, lazyproperty
|
||
from common.struct import QuerySetChain
|
||
|
||
from ..models import AssetUser, AuthBook
|
||
from .db import (
|
||
AuthbookBackend, SystemUserBackend, AdminUserBackend,
|
||
DynamicSystemUserBackend
|
||
)
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class NotSupportError(Exception):
|
||
pass
|
||
|
||
|
||
class AssetUserQueryset:
|
||
ObjectDoesNotExist = ObjectDoesNotExist
|
||
MultipleObjectsReturned = MultipleObjectsReturned
|
||
|
||
def __init__(self, backends=()):
|
||
self.backends = backends
|
||
self._distinct_queryset = None
|
||
|
||
def backends_queryset(self):
|
||
return [b.get_queryset() for b in self.backends]
|
||
|
||
@lazyproperty
|
||
def backends_counts(self):
|
||
return [b.count() for b in self.backends]
|
||
|
||
def filter(self, hostname=None, ip=None, username=None,
|
||
assets=None, asset=None, node=None,
|
||
id=None, prefer_id=None, prefer=None, id__in=None):
|
||
if not assets and asset:
|
||
assets = [asset]
|
||
|
||
kwargs = dict(
|
||
hostname=hostname, ip=ip, username=username,
|
||
assets=assets, node=node, prefer=prefer, prefer_id=prefer_id,
|
||
id__in=id__in, union_id=id,
|
||
)
|
||
logger.debug("Filter: {}".format(kwargs))
|
||
backends = []
|
||
for backend in self.backends:
|
||
clone = backend.filter(**kwargs)
|
||
backends.append(clone)
|
||
return self._clone(backends)
|
||
|
||
def _clone(self, backends=None):
|
||
if backends is None:
|
||
backends = self.backends
|
||
return self.__class__(backends)
|
||
|
||
def search(self, item):
|
||
backends = []
|
||
for backend in self.backends:
|
||
new = backend.search(item)
|
||
backends.append(new)
|
||
return self._clone(backends)
|
||
|
||
def distinct(self):
|
||
logger.debug("Distinct asset user queryset")
|
||
queryset_chain = chain(*(backend.get_queryset() for backend in self.backends))
|
||
queryset_sorted = sorted(
|
||
queryset_chain,
|
||
key=lambda item: (item["asset_username"], item["score"]),
|
||
reverse=True,
|
||
)
|
||
results = groupby(queryset_sorted, key=lambda item: item["asset_username"])
|
||
final = [next(result[1]) for result in results]
|
||
self._distinct_queryset = final
|
||
return self
|
||
|
||
def get(self, latest=False, **kwargs):
|
||
queryset = self.filter(**kwargs)
|
||
if latest:
|
||
queryset = queryset.distinct()
|
||
queryset = list(queryset)
|
||
count = len(queryset)
|
||
if count == 1:
|
||
data = queryset[0]
|
||
return data
|
||
elif count > 1:
|
||
msg = 'Should return 1 record, but get {}'.format(count)
|
||
raise MultipleObjectsReturned(msg)
|
||
else:
|
||
msg = 'No record found(org is {})'.format(current_org.name)
|
||
raise ObjectDoesNotExist(msg)
|
||
|
||
def get_latest(self, **kwargs):
|
||
return self.get(latest=True, **kwargs)
|
||
|
||
@staticmethod
|
||
def to_asset_user(data):
|
||
obj = AssetUser()
|
||
for k, v in data.items():
|
||
setattr(obj, k, v)
|
||
return obj
|
||
|
||
@property
|
||
def queryset(self):
|
||
if self._distinct_queryset is not None:
|
||
return self._distinct_queryset
|
||
return QuerySetChain(self.backends_queryset())
|
||
|
||
def count(self):
|
||
if self._distinct_queryset is not None:
|
||
return len(self._distinct_queryset)
|
||
else:
|
||
return sum(self.backends_counts)
|
||
|
||
def __getitem__(self, ndx):
|
||
return self.queryset.__getitem__(ndx)
|
||
|
||
def __iter__(self):
|
||
self._data = iter(self.queryset)
|
||
return self
|
||
|
||
def __next__(self):
|
||
return self.to_asset_user(next(self._data))
|
||
|
||
|
||
class AssetUserManager:
|
||
support_backends = (
|
||
('db', AuthbookBackend),
|
||
('system_user', SystemUserBackend),
|
||
('admin_user', AdminUserBackend),
|
||
('system_user_dynamic', DynamicSystemUserBackend),
|
||
)
|
||
|
||
def __init__(self):
|
||
self.backends = [backend() for name, backend in self.support_backends]
|
||
self._queryset = AssetUserQueryset(self.backends)
|
||
|
||
def all(self):
|
||
return self._queryset
|
||
|
||
def delete(self, obj):
|
||
name_backends_map = dict(self.support_backends)
|
||
backend_name = obj.backend
|
||
backend_cls = name_backends_map.get(backend_name)
|
||
union_id = obj.union_id
|
||
if backend_cls:
|
||
backend_cls().delete(union_id)
|
||
else:
|
||
raise ObjectDoesNotExist("Not backend found")
|
||
|
||
@staticmethod
|
||
def create(**kwargs):
|
||
# 使用create方法创建AuthBook对象，解决并发创建问题（添加锁机制）
|
||
authbook = AuthBook.create(**kwargs)
|
||
return authbook
|
||
|
||
def __getattr__(self, item):
|
||
return getattr(self._queryset, item)
|