Corvus/server/corvus/service/role_service.py

"""Role service for Corvus."""
from collections import defaultdict
from enum import Enum
from typing import Optional, List, Set, Dict


class Role(Enum):
    """User role definitions."""

    OWNER = 'OWNER'
    ADMIN = 'ADMIN'
    AUDITOR = 'AUDITOR'
    MODERATOR = 'MODERATOR'
    USER = 'USER'
    ANONYMOUS = 'ANONYMOUS'
    NONE = 'NONE'

    def __str__(self):
        return self.value


class RoleTree(defaultdict):
    """Simple tree structure to handle hierarchy."""

    def __call__(self, data: Role, power: int) -> 'RoleTree':
        """Handle direct calls to the tree."""
        return RoleTree(self, data, power)

    parent: Optional['RoleTree']
    data: Role
    power: int
    roles: Dict[Role, List['RoleTree']]

    def __init__(
            self,
            parent: Optional['RoleTree'],
            data: Role,
            power: int = None,
            **kwargs: dict) -> None:
        """Configure a RoleTree."""
        super().__init__(**kwargs)
        self.parent: Optional[RoleTree] = parent
        self.data: Role = data
        self.power: int = power if power is not None else 1
        self.default_factory = self  # type: ignore
        self.roles: Dict[Role, List[RoleTree]] = {data: [self]}

    def populate(
            self, children: Dict[Role, Optional[dict]]) -> List['RoleTree']:
        """Populate a RoleTree from a dictionary of a Role hierarchy."""
        role_list: List[RoleTree] = [self]
        for child_role in children.keys():
            element = children[child_role]
            new_node = self(child_role, self.power + 1)
            if isinstance(element, dict) and element:
                role_list.extend(new_node.populate(element))
            else:
                role_list.append(new_node)
            self[child_role] = new_node
        for role_tree in role_list:
            if role_tree.data not in self.roles.keys():
                self.roles[role_tree.data] = []
            self.roles[role_tree.data].append(role_tree)
        return role_list

    def find_role(self, request_role: Role) -> List['RoleTree']:
        """Identify all instances of a role."""
        try:
            return self.roles[request_role]
        except KeyError:
            return []

    def get_parent_roles(self) -> List[Role]:
        """Return all the roles from self to the highest parent."""
        if self.parent is not None:
            return [self.data] + self.parent.get_parent_roles()
        return [self.data]

    def get_children_roles(self) -> List[Role]:
        """Return all the roles from self to the lowest child."""
        if self.roles and (
                len(self.roles.keys()) > 1 or len(self.roles[self.data]) > 1):
            child_roles = [self.data]
            for role in self.roles.keys():
                for role_tree in self.roles[role]:
                    if role_tree.data != self.data:
                        child_roles.extend(role_tree.get_children_roles())
            return child_roles
        return [self.data]

    def find_roles_in_hierarchy(self, request_role: Role) -> Set[Role]:
        """Find a set of all roles that fall within the hierarchy."""
        roles: List[Role] = []
        role_trees = self.find_role(request_role)
        for role_tree in role_trees:
            roles.extend(role_tree.get_parent_roles())
        return set(roles)

    def find_children_roles(self, request_role: Role) -> Set[Role]:
        """Find all children roles, including this role."""
        roles: List[Role] = []
        role_trees = self.find_role(request_role)
        for role_tree in role_trees:
            roles.extend(role_tree.get_children_roles())
        return set(roles)

    def __str__(self):
        return str(self.data)


ROLES = RoleTree(None, Role.OWNER, 0)
ROLE_LIST = sorted(
    ROLES.populate({
        Role.ADMIN: {
            Role.MODERATOR: {
                Role.USER: {
                    Role.ANONYMOUS: None
                }
            },
            Role.AUDITOR: {
                Role.USER: None
            }
        }
    }),
    key=lambda rt: rt.power)