Refactoring, adding delete endpoints, using email validation pypi

Author: npalaska

lib/pbench/cli/server/shell.py

@@ -1,11 +1,12 @@
 #!/bin/env python3
 
-from pbench.server.api import create_app
+from pbench.server.api import App
 import subprocess
 
 
 def main():
-    app = create_app()
+    flask_app = App()
+    app = flask_app.create_app()
     port = app.config["PORT"]
     host = app.config["BIND_HOST"]
     workers = app.config["WORKERS"]
@@ -18,6 +19,6 @@ def main():
             "/run/pbench-server/gunicorn.pid",
             "--bind",
             f"{str(host)}:{str(port)}",
-            "pbench.server.api:create_app()",
+            "pbench.server.api:App.create_app()",
         ]
     )

lib/pbench/server/api/__init__.py

@@ -13,6 +13,7 @@
 import tempfile
 from flask_bcrypt import Bcrypt
 from flask_sqlalchemy import SQLAlchemy
+from flask_migrate import Migrate
 
 from flask import request, jsonify, Flask, make_response
 from flask_restful import Resource, abort, Api
@@ -22,113 +23,138 @@
 from pbench.server import PbenchServerConfig
 from pbench.common.exceptions import BadConfig
 from pbench.server.utils import filesize_bytes
-from pbench.server.api.resources.pbench_users import RegisterUser, LoginAPI, LogoutAPI, GetUser
+from pbench.server.api.resources.pbench_users import (
+    RegisterUser,
+    LoginAPI,
+    LogoutAPI,
+    GetUser,
+    DeleteUser,
+)
 
 ALLOWED_EXTENSIONS = {"xz"}
 
-app = None
-bcrypt = None
-db = None
 
-""""
-A storage engine to save revoked tokens. if speed is the primary concern, redis is a good option. but if data 
-persistence is more important, postgres can be another good option too. Here I am using an in memory store to get 
-something working.
-"""
-blacklist = set()
-
-
-def allowed_file(filename):
-    """Check if the file has the correct extension."""
-    try:
-        fn = filename.rsplit(".", 1)[1].lower()
-    except IndexError:
-        return False
-    allowed = "." in filename and fn in ALLOWED_EXTENSIONS
-    return allowed
-
-
-def register_endpoints(api, app):
-    """Register flask endpoints with the corresponding resource classes
-    to make the APIs active."""
-
-    api.add_resource(
-        Upload, f"{app.config['REST_URI']}/upload/ctrl/<string:controller>"
-    )
-    api.add_resource(HostInfo, f"{app.config['REST_URI']}/host_info")
-    api.add_resource(Elasticsearch, f"{app.config['REST_URI']}/elasticsearch")
-    api.add_resource(GraphQL, f"{app.config['REST_URI']}/graphql")
-    api.add_resource(RegisterUser, f"{app.config['REST_URI']}/register")
-    api.add_resource(LoginAPI, f"{app.config['REST_URI']}/login")
-    api.add_resource(LogoutAPI, f"{app.config['REST_URI']}/logout")
-    api.add_resource(GetUser, f"{app.config['REST_URI']}/user")
-
-
-def create_app():
-    """Create Flask app with defined resource endpoints."""
-
-    global app, db, bcrypt, blacklist
-
-    cfg_name = os.environ.get("_PBENCH_SERVER_CONFIG")
-    if not cfg_name:
-        print(
-            f"{__name__}: ERROR: No config file specified; set"
-            " _PBENCH_SERVER_CONFIG",
-            file=sys.stderr,
+class App:
+    app = None
+    bcrypt = None
+    db = None
+    blacklist = None
+
+    @staticmethod
+    def allowed_file(filename):
+        """Check if the file has the correct extension."""
+        try:
+            fn = filename.rsplit(".", 1)[1].lower()
+        except IndexError:
+            return False
+        allowed = "." in filename and fn in ALLOWED_EXTENSIONS
+        return allowed
+
+    @staticmethod
+    def register_endpoints(api, app, db, bcrypt, blacklist):
+        """Register flask endpoints with the corresponding resource classes
+        to make the APIs active."""
+
+        api.add_resource(
+            Upload, f"{app.config['REST_URI']}/upload/ctrl/<string:controller>"
+        )
+        api.add_resource(HostInfo, f"{app.config['REST_URI']}/host_info")
+        api.add_resource(Elasticsearch, f"{app.config['REST_URI']}/elasticsearch")
+        api.add_resource(GraphQL, f"{app.config['REST_URI']}/graphql")
+        api.add_resource(
+            RegisterUser, f"{app.config['REST_URI']}/user", resource_class_args=(app, db)
+        )
+        api.add_resource(
+            LoginAPI, f"{app.config['REST_URI']}/session", resource_class_args=(app, bcrypt)
         )
-        sys.exit(1)
-
-    try:
-        config = PbenchServerConfig(cfg_name)
-    except BadConfig as e:
-        print(f"{__name__}: {e} (config file {cfg_name})", file=sys.stderr)
-        sys.exit(1)
-
-    app = Flask(__name__)
-    api = Api(app)
-
-    app.logger = get_pbench_logger(__name__, config)
-    app.config_server = config.conf["pbench-server"]
-    app.config_elasticsearch = config.conf["elasticsearch"]
-    app.config_graphql = config.conf["graphql"]
-
-    prdp = app.config_server.get("pbench-receive-dir-prefix")
-    if not prdp:
-        app.logger.error("Missing config variable for pbench-receive-dir-prefix")
-        sys.exit(1)
-    try:
-        upload_directory = Path(f"{prdp}-002").resolve(strict=True)
-    except FileNotFoundError:
-        app.logger.exception("pbench-receive-dir-prefix does not exist on the host")
-        sys.exit(1)
-    except Exception:
-        app.logger.exception(
-            "Exception occurred during setting up the upload directory on the host"
+        api.add_resource(
+            LogoutAPI,
+            f"{app.config['REST_URI']}/session",
+            resource_class_args=(app, blacklist),
+        )
+        api.add_resource(
+            GetUser,
+            f"{app.config['REST_URI']}/user/<string:username>",
+            resource_class_args=(app),
+        )
+        api.add_resource(
+            DeleteUser,
+            f"{app.config['REST_URI']}/user/<string:username>",
+            resource_class_args=(app, db),
+        )
+
+    def create_app(self):
+        """Create Flask app with defined resource endpoints."""
+
+        cfg_name = os.environ.get("_PBENCH_SERVER_CONFIG")
+        if not cfg_name:
+            print(
+                f"{__name__}: ERROR: No config file specified; set"
+                " _PBENCH_SERVER_CONFIG",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+        try:
+            config = PbenchServerConfig(cfg_name)
+        except BadConfig as e:
+            print(f"{__name__}: {e} (config file {cfg_name})", file=sys.stderr)
+            sys.exit(1)
+
+        app = Flask(__name__)
+        api = Api(app)
+
+        app.logger = get_pbench_logger(__name__, config)
+        app.config_server = config.conf["pbench-server"]
+        app.config_elasticsearch = config.conf["elasticsearch"]
+        app.config_graphql = config.conf["graphql"]
+
+        prdp = app.config_server.get("pbench-receive-dir-prefix")
+        if not prdp:
+            app.logger.error("Missing config variable for pbench-receive-dir-prefix")
+            sys.exit(1)
+        try:
+            upload_directory = Path(f"{prdp}-002").resolve(strict=True)
+        except FileNotFoundError:
+            app.logger.exception("pbench-receive-dir-prefix does not exist on the host")
+            sys.exit(1)
+        except Exception:
+            app.logger.exception(
+                "Exception occurred during setting up the upload directory on the host"
+            )
+            sys.exit(1)
+        else:
+            app.upload_directory = upload_directory
+
+        app.config["PORT"] = app.config_server.get("rest_port")
+        app.config["VERSION"] = app.config_server.get("rest_version")
+        app.config["MAX_CONTENT_LENGTH"] = filesize_bytes(
+            app.config_server.get("rest_max_content_length")
         )
-        sys.exit(1)
-    else:
-        app.upload_directory = upload_directory
+        app.config["REST_URI"] = app.config_server.get("rest_uri")
+        app.config["LOG"] = app.config_server.get("rest_log")
+        app.config["BIND_HOST"] = app.config_server.get("bind_host")
+        app.config["WORKERS"] = app.config_server.get("workers")
 
-    app.config["PORT"] = app.config_server.get("rest_port")
-    app.config["VERSION"] = app.config_server.get("rest_version")
-    app.config["MAX_CONTENT_LENGTH"] = filesize_bytes(
-        app.config_server.get("rest_max_content_length")
-    )
-    app.config["REST_URI"] = app.config_server.get("rest_uri")
-    app.config["BIND_HOST"] = app.config_server.get("bind_host")
-    app.config["WORKERS"] = app.config_server.get("workers")
+        # postgresql specific configuration
+        app.config["DB_SECRET_KEY"] = os.getenv("SECRET_KEY", "my_precious")
+        app.config["BCRYPT_LOG_ROUNDS"] = app.config_server.get("bycrypt_log_rounds")
+        app.config["POSTGRES_DB_URI"] = app.config_server.get("postgres_db_uri")
 
-    # postgresql specific configuration
-    app.config["DB_SECRET_KEY"] = os.getenv("SECRET_KEY", "my_precious")
-    app.config["BCRYPT_LOG_ROUNDS"] = app.config_server.get("bycrypt_log_rounds")
-    app.config["POSTGRES_DB_URI"] = app.config_server.get("postgres_db_uri")
+        bcrypt = Bcrypt(app)
+        db = SQLAlchemy(app)
+        migrate = Migrate(db)
 
-    bcrypt = Bcrypt(app)
-    db = SQLAlchemy(app)
+        """"
+        A storage engine to save revoked tokens. if speed is the primary concern, redis is a good option. but if data 
+        persistence is more important, postgres can be another good option too. Here I am using an in memory store to get 
+        something working.
+        """
+        blacklist = set()
 
-    register_endpoints(api, app)
+        self.register_endpoints(api, app, db, bcrypt, blacklist)
 
-    return app
+        return app
 
 
 def _build_cors_preflight_response():
@@ -327,7 +353,7 @@ def put(self, controller):
         md5sum = request.headers.get("Content-MD5")
 
         app.logger.debug("Receiving file: {}", filename)
-        if not allowed_file(filename):
+        if not App.allowed_file(filename):
             app.logger.debug(
                 f"Tarfile upload: Bad file extension received for file {filename}"
             )

lib/pbench/server/api/resources/db_models.py

@@ -1,52 +1,52 @@
 import jwt
 import datetime
 from dateutil import parser
-from pbench.server.api import app, bcrypt, db, blacklist
+from pbench.server.api import App
 
 
-class UserModel(db.Model):
+class UserModel(App.db.Model):
     """ User Model for storing user related details """
+
     __tablename__ = "users"
 
-    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
-    username = db.Column(db.String(255), unique=True, nullable=False)
-    firstName = db.Column(db.String(255), unique=False, nullable=False)
-    lastName = db.Column(db.String(255), unique=False, nullable=False)
-    password = db.Column(db.String(255), nullable=False)
-    registered_on = db.Column(db.DateTime, nullable=False)
-    email = db.Column(db.String(255), unique=False, nullable=False)
+    id = App.db.Column(App.db.Integer, primary_key=True, autoincrement=True)
+    username = App.db.Column(App.db.String(255), unique=True, nullable=False)
+    firstName = App.db.Column(App.db.String(255), unique=False, nullable=False)
+    lastName = App.db.Column(App.db.String(255), unique=False, nullable=False)
+    password = App.db.Column(App.db.String(255), nullable=False)
+    registered_on = App.db.Column(App.db.DateTime, nullable=False)
+    email = App.db.Column(App.db.String(255), unique=False, nullable=False)
 
     def __init__(self, username, firstName, lastName, email, password):
         self.username = username
         self.firstName = firstName
         self.lastName = lastName
-        self.password = bcrypt.generate_password_hash(
-            password, app.config.get('BCRYPT_LOG_ROUNDS')
+        self.password = App.bcrypt.generate_password_hash(
+            password, App.app.config.get("BCRYPT_LOG_ROUNDS")
         ).decode()
         self.email = email
         self.registered_on = datetime.datetime.now()
 
     def __str__(self):
-        return f"New user, id: {self.id}, username: {self.username}"
+        return f"User, id: {self.id}, username: {self.username}"
 
-    def encode_auth_token(self, user_id):
+    def encode_auth_token(self, app, user_id):
         """
         Generates the Auth Token
         :return: string
         """
         try:
             payload = {
-                'iat': datetime.datetime.utcnow(),
-                'exp': datetime.datetime.utcnow() + datetime.timedelta(minutes=5),
-                'sub': user_id
+                "iat": datetime.datetime.utcnow(),
+                "exp": datetime.datetime.utcnow() + datetime.timedelta(minutes=5),
+                "sub": user_id,
             }
-            return jwt.encode(
-                payload,
-                app.config.get('SECRET_KEY'),
-                algorithm='HS256'
-            )
+            return jwt.encode(payload, app.config.get("SECRET_KEY"), algorithm="HS256")
         except Exception as e:
-            return e
+            app.logger.info(
+                f"Some exception occurred while encoding the auth token for user_id {user_id}"
+            )
+            return None
 
     @staticmethod
     def decode_auth_token(auth_token):
@@ -56,27 +56,28 @@ def decode_auth_token(auth_token):
         :return: integer|string
         """
         try:
-            payload = jwt.decode(auth_token, app.config.get('SECRET_KEY'))
-            if auth_token in blacklist:
-                return 'Token blacklisted. Please log in again.'
-            return payload['sub']
+            payload = jwt.decode(auth_token, App.app.config.get("SECRET_KEY"))
+            if auth_token in App.blacklist:
+                return "Token blacklisted. Please log in again."
+            return payload["sub"]
         except jwt.ExpiredSignatureError:
-            return 'Signature expired. Please log in again.'
+            return "Signature expired. Please log in again."
         except jwt.InvalidTokenError:
-            return 'Invalid token. Please log in again.'
+            return "Invalid token. Please log in again."
 
     # TODO: Add password recovery mechanism
 
 
-class URLModel(db.Model):
+class URLModel(App.db.Model):
     """ User Model for storing user related details """
+
     __tablename__ = "urls"
 
-    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
-    created = db.Column(db.DateTime, nullable=False)
-    updated = db.Column(db.DateTime, nullable=False)
-    config = db.Column(db.String(255), unique=False, nullable=False)
-    description = db.Column(db.String(255), nullable=False)
+    id = App.db.Column(App.db.Integer, primary_key=True, autoincrement=True)
+    created = App.db.Column(App.db.DateTime, nullable=False)
+    updated = App.db.Column(App.db.DateTime, nullable=False)
+    config = App.db.Column(App.db.String(255), unique=False, nullable=False)
+    description = App.db.Column(App.db.String(255), nullable=False)
 
     def __init__(self, created, config, description):
         self.created = parser.parse(created)
@@ -85,4 +86,4 @@ def __init__(self, created, config, description):
         self.description = description
 
     def __str__(self):
-        return f"New user, id: {self.id}, username: {self.username}"
+        return f"New user, id: {self.id}, username: {self.username}"