From b8b509c639120a8982792c8661dddd6c234a372f Mon Sep 17 00:00:00 2001
From: Drew Short <warrick@sothr.com>
Date: Mon, 7 Oct 2019 21:34:23 -0500
Subject: [PATCH] Refactor: Improve roles

---
 server/corvus/model/user_model.py     |  4 +-
 server/corvus/service/role_service.py | 53 ++++++++++++++++++---------
 server/manage.py                      |  6 ++-
 3 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/server/corvus/model/user_model.py b/server/corvus/model/user_model.py
index 49e3f8b..98abd87 100644
--- a/server/corvus/model/user_model.py
+++ b/server/corvus/model/user_model.py
@@ -10,8 +10,8 @@ class User(db.Model):  # pylint: disable=too-few-public-methods
 
     __tablename__ = 'user'
 
-    ROLE_USER = 'USER'
-    ROLE_ADMIN = 'ADMIN'
+    ROLE_USER = Role.USER.value
+    ROLE_ADMIN = Role.ADMIN.value
 
     id = db.Column(db.Integer, primary_key=True)
     name = db.Column(db.Unicode(60), unique=True, nullable=False)
diff --git a/server/corvus/service/role_service.py b/server/corvus/service/role_service.py
index 89ffd66..5b2c740 100644
--- a/server/corvus/service/role_service.py
+++ b/server/corvus/service/role_service.py
@@ -7,6 +7,7 @@ from typing import Optional, List, Set, Dict
 class Role(Enum):
     """User role definitions."""
 
+    OWNER = 'OWNER'
     ADMIN = 'ADMIN'
     AUDITOR = 'AUDITOR'
     MODERATOR = 'MODERATOR'
@@ -14,39 +15,48 @@ class Role(Enum):
     ANONYMOUS = 'ANONYMOUS'
     NONE = 'NONE'
 
+    def __str__(self):
+        return self.value
+
 
 class RoleTree(defaultdict):
     """Simple tree structure to handle hierarchy."""
 
-    def __call__(self, data: Role) -> 'RoleTree':
+    def __call__(self, data: Role, power: int) -> 'RoleTree':
         """Handle direct calls to the tree."""
-        return RoleTree(self, data)
+        return RoleTree(self, data, power)
 
-    # def __hash__(self):
+    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] = []
+        role_list: List[RoleTree] = [self]
         for child_role in children.keys():
             element = children[child_role]
-            new_node = self(child_role)
+            new_node = self(child_role, self.power + 1)
             if isinstance(element, dict) and element:
-                role_list += new_node.populate(element)
+                role_list.extend(new_node.populate(element))
+            else:
+                role_list.append(new_node)
             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] = []
@@ -94,15 +104,22 @@ class RoleTree(defaultdict):
             roles.extend(role_tree.get_children_roles())
         return set(roles)
 
-
-ROLES = RoleTree(None, Role.ADMIN)
-ROLES.populate({
-    Role.MODERATOR: {
-        Role.USER: {
-            Role.ANONYMOUS: None
+    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
+            }
         }
-    },
-    Role.AUDITOR: {
-        Role.USER: None
-    }
-})
+    }),
+    key=lambda rt: rt.power)
diff --git a/server/manage.py b/server/manage.py
index 0e08875..28c8db3 100644
--- a/server/manage.py
+++ b/server/manage.py
@@ -11,6 +11,7 @@ from click import Context
 from corvus import corvus
 from corvus.model import User
 from corvus.service import user_service
+from corvus.service.role_service import ROLE_LIST
 
 logging.basicConfig()
 
@@ -54,7 +55,10 @@ def delete_user(name: str):
 @click.option('--role',
               default=User.ROLE_USER,
               envvar='ROLE',
-              help='Role to assign to the user. default=[USER]')
+              help='Role to assign to the user. '
+                   + 'default=[USER] acceptable values = ['
+                   + ','.join(sorted(set(map(lambda rt: str(rt), ROLE_LIST))))
+                   + ']')
 def register_user(
         name: str,
         role: str,