geovisio 2.5.0__py3-none-any.whl → 2.7.0__py3-none-any.whl

geovisio/utils/tokens.py

@@ -1,11 +1,9 @@
 from geovisio import errors
-from geovisio.utils import auth
+from geovisio.utils import auth, db
 from geovisio.web.tokens import _decode_jwt_token, _generate_jwt_token
-
-
-import psycopg
 from authlib.jose.errors import BadSignatureError
 from flask import current_app
+from flask_babel import gettext as _
 from psycopg.rows import dict_row
 
 
@@ -14,7 +12,7 @@ import logging
 
 class InvalidTokenException(errors.InvalidAPIUsage):
     def __init__(self, details, status_code=401):
-        msg = f"Token not valid"
+        msg = "Token not valid"
         super().__init__(msg, status_code=status_code, payload={"details": {"error": details}})
 
 
@@ -39,39 +37,37 @@ def get_account_from_jwt_token(jwt_token: str) -> auth.Account:
     """
     try:
         decoded = _decode_jwt_token(jwt_token)
-    except BadSignatureError as e:
+    except BadSignatureError:
         logging.exception("invalid signature of jwt token")
-        raise InvalidTokenException("JWT token signature does not match")
+        raise InvalidTokenException(_("JWT token signature does not match"))
     token_id = decoded["sub"]
 
-    with psycopg.connect(current_app.config["DB_URL"], row_factory=dict_row) as conn:
-        with conn.cursor() as cursor:
-            # check token existence
-            records = cursor.execute(
-                """
-                SELECT
-                    t.account_id AS id, a.name, a.oauth_provider, a.oauth_id
-                FROM tokens t
-                LEFT OUTER JOIN accounts a ON t.account_id = a.id
-                WHERE t.id = %(token)s
-            """,
-                {"token": token_id},
-            ).fetchone()
-            if not records:
-                raise InvalidTokenException("Token does not exist anymore", status_code=403)
-
-            if not records["id"]:
-                raise InvalidTokenException(
-                    "Token not yet claimed, this token cannot be used yet. Either claim this token or generate a new one", status_code=403
-                )
-
-            return auth.Account(
-                id=str(records["id"]),
-                name=records["name"],
-                oauth_provider=records["oauth_provider"],
-                oauth_id=records["oauth_id"],
+    with db.cursor(current_app, row_factory=dict_row) as cursor:
+        # check token existence
+        records = cursor.execute(
+            """SELECT
+    t.account_id AS id, a.name, a.oauth_provider, a.oauth_id, a.role
+FROM tokens t
+LEFT OUTER JOIN accounts a ON t.account_id = a.id
+WHERE t.id = %(token)s""",
+            {"token": token_id},
+        ).fetchone()
+        if not records:
+            raise InvalidTokenException(_("Token does not exist anymore"), status_code=403)
+
+        if not records["id"]:
+            raise InvalidTokenException(
+                _("Token not yet claimed, this token cannot be used yet. Either claim this token or generate a new one"), status_code=403
             )
 
+        return auth.Account(
+            id=str(records["id"]),
+            name=records["name"],
+            oauth_provider=records["oauth_provider"],
+            oauth_id=records["oauth_id"],
+            role=auth.AccountRole[records["role"]],
+        )
+
 
 def get_default_account_jwt_token() -> str:
     """
@@ -80,18 +76,17 @@ def get_default_account_jwt_token() -> str:
     Note: do not expose this method externally, only an instance administrator should be able to get the default account JWT token!
     """
 
-    with psycopg.connect(current_app.config["DB_URL"], row_factory=dict_row) as conn:
-        with conn.cursor() as cursor:
-            # check token existence
-            records = cursor.execute(
-                """
+    with db.cursor(current_app, row_factory=dict_row) as cursor:
+        # check token existence
+        records = cursor.execute(
+            """
                 SELECT t.id AS id
                 FROM tokens t
                 JOIN accounts a ON t.account_id = a.id
                 WHERE a.is_default
             """
-            ).fetchone()
-            if not records:
-                raise Exception("Default account has no associated token")
+        ).fetchone()
+        if not records:
+            raise Exception("Default account has no associated token")
 
-            return _generate_jwt_token(records["id"])
+        return _generate_jwt_token(records["id"])