Initial implementation of userDB and associated APIs in server

Author: npalaska

lib/pbench/server/api/__init__.py

@@ -11,6 +11,8 @@
 import requests
 import humanize
 import tempfile
+from flask_bcrypt import Bcrypt
+from flask_sqlalchemy import SQLAlchemy
 
 from flask import request, jsonify, Flask, make_response
 from flask_restful import Resource, abort, Api
@@ -24,7 +26,8 @@
 ALLOWED_EXTENSIONS = {"xz"}
 
 app = None
-
+bcrypt = None
+db = None
 
 def allowed_file(filename):
     """Check if the file has the correct extension."""
@@ -51,7 +54,7 @@ def register_endpoints(api, app):
 def create_app():
     """Create Flask app with defined resource endpoints."""
 
-    global app
+    global app, db, bcrypt
 
     cfg_name = os.environ.get("_PBENCH_SERVER_CONFIG")
     if not cfg_name:
@@ -103,6 +106,14 @@ def create_app():
     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")
+
+    bcrypt = Bcrypt(app)
+    db = SQLAlchemy(app)
+
     register_endpoints(api, app)
 
     return app

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

@@ -0,0 +1,83 @@
+import jwt
+import datetime
+from dateutil import parser
+from pbench.server.api import app, bcrypt, db
+
+
+class UserModel(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)
+    # ToDo: decide if we want the email field
+
+    def __init__(self, username, firstName, lastName, password):
+        self.username = username
+        self.firstName = firstName
+        self.lastName = lastName
+        self.password = bcrypt.generate_password_hash(
+            password, app.config.get('BCRYPT_LOG_ROUNDS')
+        ).decode()
+        self.registered_on = datetime.datetime.now()
+
+    def __str__(self):
+        return f"New user, id: {self.id}, username: {self.username}"
+
+    def encode_auth_token(self, 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
+            }
+            return jwt.encode(
+                payload,
+                app.config.get('SECRET_KEY'),
+                algorithm='HS256'
+            )
+        except Exception as e:
+            return e
+
+    @staticmethod
+    def decode_auth_token(auth_token):
+        """
+        Validates the auth token
+        :param auth_token:
+        :return: integer|string
+        """
+        try:
+            payload = jwt.decode(auth_token, app.config.get('SECRET_KEY'))
+            return payload['sub']
+        except jwt.ExpiredSignatureError:
+            return 'Signature expired. Please log in again.'
+        except jwt.InvalidTokenError:
+            return 'Invalid token. Please log in again.'
+
+
+class URLModel(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)
+
+    def __init__(self, created, config, description):
+        self.created = parser.parse(created)
+        self.updated = datetime.datetime.now()
+        self.config = config
+        self.description = description
+
+    def __str__(self):
+        return f"New user, id: {self.id}, username: {self.username}"

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

@@ -0,0 +1,93 @@
+
+from flask import request, jsonify, make_response
+from flask_restful import Resource, abort
+from pbench.server.api import app, bcrypt, db
+from pbench.server.api.resources.db_models import UserModel
+
+
+class RegisterUser(Resource):
+    """
+    Abstracted pbench API for registering a new user
+    """
+
+    def post(self):
+        # get the post data
+        post_data = request.get_json()
+        # check if user already exists
+        user = UserModel.query.filter_by(email=post_data.get('username')).first()
+        if not user:
+            try:
+                user = UserModel(
+                    username=post_data.get('username'),
+                    password=post_data.get('password'),
+                    firstName=post_data.get('firstName'),
+                    lastName=post_data.get('lastName'),
+                )
+
+                # insert the user
+                db.session.add(user)
+                db.session.commit()
+                app.logger.info(f"New user registered, id: {self.id}, username: {self.username}")
+
+                # generate the auth token
+                auth_token = user.encode_auth_token(user.id)
+                response_object = {
+                    'status': 'success',
+                    'status_code': 201,
+                    'message': 'Successfully registered.',
+                    'auth_token': auth_token.decode()
+                }
+                return make_response(jsonify(response_object))
+
+            except Exception as e:
+                response_object = {
+                    'status': 'fail',
+                    'status_code': 401,
+                    'message': 'Some error occurred. Please try again.'
+                }
+                return make_response(jsonify(response_object))
+        else:
+            response_object = {
+                'status': 'fail',
+                'status_code': 202,
+                'message': 'User already exists. Please Log in.',
+            }
+            return make_response(jsonify(response_object))
+
+
+class LoginAPI(Resource):
+    """
+    Abstracted pbench API for User Login
+    """
+    def post(self):
+        # get the post data
+        post_data = request.get_json()
+        try:
+            # fetch the user data
+            user = UserModel.query.filter_by(username=post_data.get('username')).first()
+
+            if user and bcrypt.check_password_hash(user.password, post_data.get('password')):
+                auth_token = user.encode_auth_token(user.id)
+                if auth_token:
+                    response_object = {
+                        'status': 'success',
+                        'status_code': 200,
+                        'message': 'Successfully logged in.',
+                        'auth_token': auth_token.decode()
+                    }
+                    return make_response(jsonify(response_object))
+            else:
+                response_object = {
+                    'status': 'fail',
+                    'status_code': 404,
+                    'message': 'User does not exist.'
+                }
+                return make_response(jsonify(response_object))
+        except Exception as e:
+            app.logger.info(f"Exception occurred during user login {e}")
+            response_object = {
+                'status': 'fail',
+                'status_code': 500,
+                'message': 'Try again'
+            }
+            return make_response(jsonify(response_object))