jumpserver/apps/assets/backends/manager.py

# -*- 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):
        authbook = AuthBook(**kwargs)
        authbook.save()
        return authbook

    def __getattr__(self, item):
        return getattr(self._queryset, item)