Atheneum/server/atheneum/service/role_service.py

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


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

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


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

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

    # def __hash__(self):

    def __init__(
            self,
            parent: Optional['RoleTree'],
            data: Role,
            **kwargs: dict) -> None:
        """Configure a RoleTree."""
        super().__init__(**kwargs)
        self.parent: Optional[RoleTree] = parent
        self.data: Role = data
        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] = []
        for child_role in children.keys():
            element = children[child_role]
            new_node = self(child_role)
            if isinstance(element, dict) and element:
                role_list += new_node.populate(element)
            self[child_role] = new_node
            role_list.append(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 [role_tree for role_tree in 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)


ROLES = RoleTree(None, Role.ADMIN)
ROLES.populate({
    Role.MODERATOR: {
        Role.USER: {
            Role.ANONYMOUS: None
        }
    },
    Role.AUDITOR: {
        Role.USER: None
    }
})