Switch to Flask_Login for user/session management

gallery/__init__.py

@@ -11,6 +11,7 @@ import logging
 from flask_compress import Compress
 from flask_caching import Cache
 from flask_assets import Environment, Bundle
+from flask_login import LoginManager
 from flask import Flask, render_template, abort
 from werkzeug.exceptions import HTTPException
 
@@ -19,18 +20,26 @@ import platformdirs
 from dotenv import load_dotenv
 from yaml import safe_load
 
+# Import database
+from sqlalchemy.orm import sessionmaker
+from gallery import db
+
 
 USER_DIR = platformdirs.user_config_dir('onlylegs')
 
+db_session = sessionmaker(bind=db.engine)
+db_session = db_session()
+login_manager = LoginManager()
+assets = Environment()
+cache = Cache(config={'CACHE_TYPE': 'SimpleCache', 'CACHE_DEFAULT_TIMEOUT': 300})
+compress = Compress()
+
 
 def create_app(test_config=None):
     """
     Create and configure the main app
     """
     app = Flask(__name__, instance_path=os.path.join(USER_DIR, 'instance'))
-    assets = Environment()
-    cache = Cache(config={'CACHE_TYPE': 'SimpleCache', 'CACHE_DEFAULT_TIMEOUT': 300})
-    compress = Compress()
 
     # Get environment variables
     load_dotenv(os.path.join(USER_DIR, '.env'))
@@ -56,6 +65,13 @@ def create_app(test_config=None):
     else:
         app.config.from_mapping(test_config)
 
+    login_manager.init_app(app)
+    login_manager.login_view = 'gallery.index'
+
+    @login_manager.user_loader
+    def load_user(user_id):
+        return db_session.query(db.Users).filter_by(id=user_id).first()
+
     # Load JS assets
     # TODO: disable caching for sass files as it makes it hard to work on when it is enabled
     assets.register('js_pre', Bundle('js/pre/*.js', output='gen/pre_packed.js'))

gallery/auth.py

@@ -3,17 +3,15 @@ OnlyLegs - Authentication
 User registration, login and logout and locking access to pages behind a login
 """
 import re
-import uuid
 import logging
 from datetime import datetime as dt
 
-import functools
-from flask import Blueprint, flash, g, redirect, request, session, url_for, abort, jsonify
+from flask import Blueprint, flash, redirect, request, url_for, abort, jsonify
 from werkzeug.security import check_password_hash, generate_password_hash
 
-from sqlalchemy.orm import sessionmaker
-from sqlalchemy import exc
+from flask_login import login_user, logout_user, login_required
 
+from sqlalchemy.orm import sessionmaker
 from gallery import db
 
 
@@ -22,42 +20,30 @@ db_session = sessionmaker(bind=db.engine)
 db_session = db_session()
 
 
-def login_required(view):
+@blueprint.route('/login', methods=['POST'])
+def login():
     """
-    Decorator to check if a user is logged in before accessing a page
+    Log in a registered user by adding the user id to the session
     """
-    @functools.wraps(view)
-    def wrapped_view(**kwargs):
-        if g.user is None or session.get('uuid') is None:
-            logging.error('Authentication failed')
-            session.clear()
-            return redirect(url_for('gallery.index'))
+    error = []
     
-        return view(**kwargs)
+    username = request.form['username'].strip()
+    password = request.form['password'].strip()
 
-    return wrapped_view
+    user = db_session.query(db.Users).filter_by(username=username).first()
 
+    if not user and not check_password_hash(user.password, password):
+        logging.error('Login attempt from %s', username, request.remote_addr)
+        error.append('Username or Password is incorrect!')
 
-@blueprint.before_app_request
-def load_logged_in_user():
-    """
-    Runs before every request and checks if a user is logged in
-    """
-    user_id = session.get('user_id')
-    user_uuid = session.get('uuid')
+    if error:
+        abort(403)
 
-    if user_id is None or user_uuid is None:
-        g.user = None
-        session.clear()
-    else:
-        is_alive = db_session.query(db.Sessions).filter_by(session_uuid=user_uuid).first()
+    login_user(user)
 
-        if is_alive is None:
-            logging.info('Session expired')
-            flash(['Session expired!', '3'])
-            session.clear()
-        else:
-            g.user = db_session.query(db.Users).filter_by(id=user_id).first()
+    logging.info('User %s logged in from %s', username, request.remote_addr)
+    flash(['Logged in successfully!', '4'])
+    return 'gwa gwa'
 
 
 @blueprint.route('/register', methods=['POST'])
@@ -65,17 +51,18 @@ def register():
     """
     Register a new user
     """
+    error = []
+    
     # Thanks Fennec for reminding me to strip out the whitespace lol
     username = request.form['username'].strip()
     email = request.form['email'].strip()
     password = request.form['password'].strip()
     password_repeat = request.form['password-repeat'].strip()
 
-    error = []
-
     email_regex = re.compile(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b')
     username_regex = re.compile(r'\b[A-Za-z0-9._-]+\b')
 
+    # Validate the form
     if not username or not username_regex.match(username):
         error.append('Username is invalid!')
 
@@ -92,74 +79,29 @@ def register():
     elif password_repeat != password:
         error.append('Passwords do not match!')
         
+    user_exists = db_session.query(db.Users).filter_by(username=username).first()
+    if user_exists:
+        error.append('User already exists!')
+
+    # If there are errors, return them
     if error:
         return jsonify(error)
 
-    try:
-        register_user = db.Users(username=username,
-                                 email=email,
-                                 password=generate_password_hash(password),
+    register_user = db.Users(username=username, email=email,
+                             password=generate_password_hash(password, method='sha256'),
                              created_at=dt.utcnow())
     db_session.add(register_user)
     db_session.commit()
-    except exc.IntegrityError:
-        return f'User {username} is already registered!'
 
     logging.info('User %s registered', username)
     return 'gwa gwa'
 
 
-@blueprint.route('/login', methods=['POST'])
-def login():
-    """
-    Log in a registered user by adding the user id to the session
-    """
-    username = request.form['username'].strip()
-    password = request.form['password'].strip()
-
-    user = db_session.query(db.Users).filter_by(username=username).first()
-    error = []
-
-    if user is None:
-        logging.error('User %s does not exist. Login attempt from %s',
-                      username, request.remote_addr)
-        error.append('Username or Password is incorrect!')
-    elif not check_password_hash(user.password, password):
-        logging.error('User %s entered wrong password. Login attempt from %s',
-                      username, request.remote_addr)
-        error.append('Username or Password is incorrect!')
-
-    if error:
-        abort(403)
-
-    try:
-        session.clear()
-        session['user_id'] = user.id
-        session['uuid'] = str(uuid.uuid4())
-
-        session_query = db.Sessions(user_id=user.id,
-                                    session_uuid=session.get('uuid'),
-                                    ip_address=request.remote_addr,
-                                    user_agent=request.user_agent.string,
-                                    active=True,
-                                    created_at=dt.utcnow())
-
-        db_session.add(session_query)
-        db_session.commit()
-    except Exception as err:
-        logging.error('User %s could not be logged in: %s', username, err)
-        abort(500)
-
-    logging.info('User %s logged in from %s', username, request.remote_addr)
-    flash(['Logged in successfully!', '4'])
-    return 'gwa gwa'
-
-
 @blueprint.route('/logout')
+@login_required
 def logout():
     """
     Clear the current session, including the stored user id
     """
-    logging.info('User (%s) %s logged out', session.get('user_id'), g.user.username)
-    session.clear()
+    logout_user()
     return redirect(url_for('gallery.index'))

gallery/db.py

@@ -8,6 +8,8 @@ from sqlalchemy import (
     create_engine, Column, Integer, String, Boolean, DateTime, ForeignKey, PickleType)
 from sqlalchemy.orm import declarative_base, relationship
 
+from flask_login import UserMixin
+
 
 USER_DIR = platformdirs.user_config_dir('onlylegs')
 DB_PATH = os.path.join(USER_DIR, 'gallery.sqlite')
@@ -18,7 +20,7 @@ engine = create_engine(f'sqlite:///{DB_PATH}', echo=False)
 base = declarative_base()
 
 
-class Users (base):  # pylint: disable=too-few-public-methods, C0103
+class Users (base, UserMixin):  # pylint: disable=too-few-public-methods, C0103
     """
     User table
     Joins with post, groups, session and log

gallery/templates/groups/group.html

@@ -75,7 +75,7 @@
     {% else %}
         <div class="big-text">
             <h1>*crickets chirping*</h1>
-            {% if g.user %}
+            {% if current_user.is_authenticated %}
                 <p>Add some images to the group!</p>
             {% else %}
                 <p>Login to start managing this image group!</p>

gallery/templates/groups/list.html

@@ -43,7 +43,7 @@
     {% else %}
         <div class="big-text">
             <h1>*crickets chirping*</h1>
-            {% if g.user %}
+            {% if current_user.is_authenticated %}
                 <p>You can get started by creating a new image group!</p>
             {% else %}
                 <p>Login to start seeing anything here!</p>

gallery/templates/image.html

@@ -33,7 +33,7 @@
             }
         }
 
-        {% if g.user.id == image.author_id %}
+        {% if current_user.id == image.author_id %}
             cancelBtn = document.createElement('button');
             cancelBtn.classList.add('btn-block');
             cancelBtn.innerHTML = 'nuuuuuuuu';
@@ -144,7 +144,7 @@
                     </span>
                 </a>
             </div>
-            {% if g.user.id == image.author_id %}
+            {% if current_user.id == image.author_id %}
                 <div>
                     <button class="pill-item pill__critical" onclick="imageDelete()">
                         <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" fill="currentColor" viewBox="0 0 256 256"><path d="M216,48H176V40a24,24,0,0,0-24-24H104A24,24,0,0,0,80,40v8H40a8,8,0,0,0,0,16h8V208a16,16,0,0,0,16,16H192a16,16,0,0,0,16-16V64h8a8,8,0,0,0,0-16ZM112,168a8,8,0,0,1-16,0V104a8,8,0,0,1,16,0Zm48,0a8,8,0,0,1-16,0V104a8,8,0,0,1,16,0Zm0-120H96V40a8,8,0,0,1,8-8h48a8,8,0,0,1,8,8Z"></path></svg>

gallery/templates/layout.html

@@ -80,7 +80,7 @@
                 </span>
             </a>
 
-            {% if g.user %}
+            {% if current_user.is_authenticated %}
             <button class="navigation-item {% block nav_upload %}{% endblock %}" onclick="toggleUploadTab()">
                 <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" fill="currentColor" viewBox="0 0 256 256"><path d="M74.34,77.66a8,8,0,0,1,0-11.32l48-48a8,8,0,0,1,11.32,0l48,48a8,8,0,0,1-11.32,11.32L136,43.31V128a8,8,0,0,1-16,0V43.31L85.66,77.66A8,8,0,0,1,74.34,77.66ZM240,136v64a16,16,0,0,1-16,16H32a16,16,0,0,1-16-16V136a16,16,0,0,1,16-16h68a4,4,0,0,1,4,4v3.46c0,13.45,11,24.79,24.46,24.54A24,24,0,0,0,152,128v-4a4,4,0,0,1,4-4h68A16,16,0,0,1,240,136Zm-40,32a12,12,0,1,0-12,12A12,12,0,0,0,200,168Z"></path></svg>
                 <span class="tool-tip">
@@ -92,7 +92,7 @@
 
             <span class="navigation-spacer"></span>
 
-            {% if g.user %}
+            {% if current_user.is_authenticated %}
             <a href="{{url_for('gallery.profile')}}" class="navigation-item {% block nav_profile %}{% endblock %}">
                 <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" fill="currentColor" viewBox="0 0 256 256"><path d="M231.73,221.94A8,8,0,0,1,224,232H160A8,8,0,0,1,152.27,222a40,40,0,0,1,17.11-23.33,32,32,0,1,1,45.24,0A40,40,0,0,1,231.73,221.94ZM216,72H130.67L102.93,51.2a16.12,16.12,0,0,0-9.6-3.2H40A16,16,0,0,0,24,64V200a16,16,0,0,0,16,16h80a8,8,0,0,0,0-16H40V64H93.33l27.74,20.8a16.12,16.12,0,0,0,9.6,3.2H216v32a8,8,0,0,0,16,0V88A16,16,0,0,0,216,72Z"></path></svg>
                 <span class="tool-tip">
@@ -119,7 +119,7 @@
             {% endif %}
         </div>
 
-        {% if g.user %}
+        {% if current_user.is_authenticated %}
         <div class="upload-panel">
             <span class="click-off" onclick="closeUploadTab()"></span>
             <div class="container">

gallery/templates/profile.html

@@ -8,7 +8,7 @@
 
         <div class="banner-content">
             <h1>Profile</h1>
-            <p>Hello {{ g.user['username'] }}</p>
+            <p>Hello {{ current_user.username }}</p>
         </div>
     </div>
 

gallery/views/api.py

@@ -8,13 +8,14 @@ import logging
 from datetime import datetime as dt
 import platformdirs
 
-from flask import Blueprint, send_from_directory, abort, flash, jsonify, request, g, current_app
+from flask import Blueprint, send_from_directory, abort, flash, jsonify, request, current_app
 from werkzeug.utils import secure_filename
 
+from flask_login import login_required, current_user
+
 from colorthief import ColorThief
 
 from sqlalchemy.orm import sessionmaker
-from gallery.auth import login_required
 
 from gallery import db
 from gallery.utils import metadata as mt
@@ -83,7 +84,7 @@ def upload():
     img_colors = ColorThief(img_path).get_palette(color_count=3)  # Get color palette
 
     # Save to database
-    query = db.Posts(author_id=g.user.id,
+    query = db.Posts(author_id=current_user.id,
                      created_at=dt.utcnow(),
                      file_name=img_name+'.'+img_ext,
                      file_type=img_ext,
@@ -109,7 +110,7 @@ def delete_image(image_id):
     # Check if image exists and if user is allowed to delete it (author)
     if img is None:
         abort(404)
-    if img.author_id != g.user.id:
+    if img.author_id != current_user.id:
         abort(403)
 
     # Delete file
@@ -148,7 +149,7 @@ def create_group():
     """
     new_group = db.Groups(name=request.form['name'],
                           description=request.form['description'],
-                          author_id=g.user.id,
+                          author_id=current_user.id,
                           created_at=dt.utcnow())
 
     db_session.add(new_group)
@@ -170,7 +171,7 @@ def modify_group():
 
     if group is None:
         abort(404)
-    elif group.author_id != g.user.id:
+    elif group.author_id != current_user.id:
         abort(403)
 
     if request.form['action'] == 'add':

gallery/views/settings.py

@@ -2,9 +2,7 @@
 OnlyLegs - Settings page
 """
 from flask import Blueprint, render_template
-
-from gallery.auth import login_required
-
+from flask_login import login_required
 
 blueprint = Blueprint('settings', __name__, url_prefix='/settings')