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