From 8ef75903253c440c6fd3e9edbd00e1f9bf8ce381 Mon Sep 17 00:00:00 2001
From: Marcelo Arocha <marceloarocha@gmail.com>
Date: Mon, 6 Nov 2023 14:09:23 -0300
Subject: [PATCH] admin route schema

---
 models/enums.py          |  1 +
 routes/authentication.py | 18 ++++++++++-
 services/auth_service.py | 64 ++++++++++++++++++++++++++++++++++------
 3 files changed, 73 insertions(+), 10 deletions(-)

diff --git a/models/enums.py b/models/enums.py
index 4f44bec5..de7f80c9 100644
--- a/models/enums.py
+++ b/models/enums.py
@@ -31,6 +31,7 @@ class RoleEnum(Enum):
     TRAINING = "training"
     DOCTOR = "doctor"
     READONLY = "readonly"
+    MULTI_SCHEMA = "multi-schema"
 
 
 class FeatureEnum(Enum):
diff --git a/routes/authentication.py b/routes/authentication.py
index 3ba3ae43..cffa7079 100644
--- a/routes/authentication.py
+++ b/routes/authentication.py
@@ -18,16 +18,32 @@
 app_auth = Blueprint("app_auth", __name__)
 
 
+@app_auth.route("/pre-auth", methods=["POST"])
+def pre_auth():
+    data = request.get_json()
+
+    email = data.get("email", None)
+    password = data.get("password", None)
+
+    try:
+        auth_data = auth_service.pre_auth(email, password)
+    except ValidationError as e:
+        return {"status": "error", "message": str(e), "code": e.code}, e.httpStatus
+
+    return auth_data, status.HTTP_200_OK
+
+
 @app_auth.route("/authenticate", methods=["POST"])
 def auth():
     data = request.get_json()
 
     email = data.get("email", None)
     password = data.get("password", None)
+    schema = data.get("schema", None)
     refresh_token = None
 
     try:
-        auth_data = auth_service.auth_local(email, password)
+        auth_data = auth_service.auth_local(email, password, force_schema=schema)
 
         refresh_token = auth_data["refresh_token"]
         # temp
diff --git a/services/auth_service.py b/services/auth_service.py
index 74f1f0f9..361d1f16 100644
--- a/services/auth_service.py
+++ b/services/auth_service.py
@@ -18,9 +18,19 @@
 from exception.validation_error import ValidationError
 
 
-def _auth_user(user, db_session):
+def _has_force_schema_permission(roles, schema):
+    if schema != "hsc_test":
+        return False
+
+    return (
+        RoleEnum.ADMIN.value in roles or RoleEnum.TRAINING.value in roles
+    ) and RoleEnum.MULTI_SCHEMA.value in roles
+
+
+def _auth_user(user, force_schema=None):
+    roles = user.config["roles"] if user.config and "roles" in user.config else []
+
     if Config.ENV == NoHarmENV.STAGING.value:
-        roles = user.config["roles"] if user.config and "roles" in user.config else []
         if (
             RoleEnum.SUPPORT.value not in roles
             and RoleEnum.ADMIN.value not in roles
@@ -34,11 +44,20 @@ def _auth_user(user, db_session):
                 status.HTTP_401_UNAUTHORIZED,
             )
 
-    claims = {"schema": user.schema, "config": user.config}
+    user_schema = user.schema
+    if force_schema and _has_force_schema_permission(roles, user.schema):
+        user_schema = force_schema
+
+    claims = {"schema": user_schema, "config": user.config}
     access_token = create_access_token(identity=user.id, additional_claims=claims)
     refresh_token = create_refresh_token(identity=user.id, additional_claims=claims)
 
-    notification = Notify.getNotification(schema=user.schema)
+    db_session = db.create_scoped_session()
+    db_session.connection(
+        execution_options={"schema_translate_map": {None: user_schema}}
+    )
+
+    notification = Notify.getNotification(schema=user_schema)
 
     if notification is not None:
         notificationMemory = (
@@ -55,7 +74,7 @@ def _auth_user(user, db_session):
 
     features = db_session.query(Memory).filter(Memory.kind == "features").first()
 
-    nameUrl = Memory.getNameUrl(user.schema)
+    nameUrl = Memory.getNameUrl(user_schema)
 
     hospitals = db_session.query(Hospital).order_by(asc(Hospital.name)).all()
     hospitalList = []
@@ -85,7 +104,7 @@ def _auth_user(user, db_session):
         "userName": user.name,
         "userId": user.id,
         "email": user.email,
-        "schema": user.schema,
+        "schema": user_schema,
         "roles": user.config["roles"] if user.config and "roles" in user.config else [],
         "userFeatures": user.config["features"]
         if user.config and "features" in user.config
@@ -109,7 +128,34 @@ def _auth_user(user, db_session):
     }
 
 
-def auth_local(email, password):
+def pre_auth(email, password):
+    user = User.authenticate(email, password)
+
+    if user is None:
+        raise ValidationError(
+            "Usuário inválido",
+            "errors.unauthorizedUser",
+            status.HTTP_400_BAD_REQUEST,
+        )
+
+    roles = user.config["roles"] if user.config and "roles" in user.config else []
+
+    if _has_force_schema_permission(roles, user.schema):
+        result = db.session.execute(
+            """
+            SELECT schema_name FROM information_schema.schemata where schema_name not in ('pg_catalog', 'information_schema') order by schema_name                   
+        """
+        )
+        schemas = []
+        for r in result:
+            schemas.append(r[0])
+
+        return {"maintainer": True, "schemas": schemas}
+
+    return {"maintainer": False, "schemas": []}
+
+
+def auth_local(email, password, force_schema=None):
     preCheckUser = User.query.filter_by(email=email).first()
 
     if preCheckUser is None:
@@ -155,7 +201,7 @@ def auth_local(email, password):
             status.HTTP_400_BAD_REQUEST,
         )
 
-    return _auth_user(user, db_session)
+    return _auth_user(user, force_schema=force_schema)
 
 
 def auth_provider(code, schema):
@@ -292,7 +338,7 @@ def auth_provider(code, schema):
                 status.HTTP_401_UNAUTHORIZED,
             )
 
-    return _auth_user(nh_user, db.session)
+    return _auth_user(nh_user)
 
 
 def _get_oauth_user(email, name, schema, oauth_config):