diff --git a/gallery/__init__.py b/gallery/__init__.py index 15f14a9..bd76ae7 100644 --- a/gallery/__init__.py +++ b/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')) diff --git a/gallery/auth.py b/gallery/auth.py index 3e13334..b51c885 100644 --- a/gallery/auth.py +++ b/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 = [] + + username = request.form['username'].strip() + password = request.form['password'].strip() - return view(**kwargs) + user = db_session.query(db.Users).filter_by(username=username).first() - return wrapped_view + 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!') + if error: + abort(403) -@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') + login_user(user) - 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() - - 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!') @@ -91,75 +78,30 @@ def register(): error.append('Enter password again!') 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), - created_at=dt.utcnow()) - db_session.add(register_user) - db_session.commit() - except exc.IntegrityError: - return f'User {username} is already registered!' + 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() 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')) diff --git a/gallery/db.py b/gallery/db.py index bb85dc6..df1f389 100644 --- a/gallery/db.py +++ b/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 diff --git a/gallery/templates/groups/group.html b/gallery/templates/groups/group.html index 66df5da..9376051 100644 --- a/gallery/templates/groups/group.html +++ b/gallery/templates/groups/group.html @@ -75,7 +75,7 @@ {% else %}

*crickets chirping*

- {% if g.user %} + {% if current_user.is_authenticated %}

Add some images to the group!

{% else %}

Login to start managing this image group!

diff --git a/gallery/templates/groups/list.html b/gallery/templates/groups/list.html index 62cb3e6..ef726f8 100644 --- a/gallery/templates/groups/list.html +++ b/gallery/templates/groups/list.html @@ -43,7 +43,7 @@ {% else %}

*crickets chirping*

- {% if g.user %} + {% if current_user.is_authenticated %}

You can get started by creating a new image group!

{% else %}

Login to start seeing anything here!

diff --git a/gallery/templates/image.html b/gallery/templates/image.html index 7790967..c275ce3 100644 --- a/gallery/templates/image.html +++ b/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 @@
- {% if g.user.id == image.author_id %} + {% if current_user.id == image.author_id %}
{% if request.path == "/" %} - + {% endif %}
@@ -80,67 +80,67 @@ - {% if g.user %} - + {% if current_user.is_authenticated %} + {% endif %} - {% if g.user %} - - - - Profile - - - - - - - - Settings - - - + {% if current_user.is_authenticated %} + + + + Profile + + + + + + + + Settings + + + {% else %} - + {% endif %}
- {% if g.user %} -
- -
- -

Upload stuffs

-

May the world see your stuff 👀

-
- + {% if current_user.is_authenticated %} +
+ +
+ +

Upload stuffs

+

May the world see your stuff 👀

+ + - - - - - -
-
+ + + + + +
+
{% endif %}
diff --git a/gallery/templates/profile.html b/gallery/templates/profile.html index 59b620b..196abc3 100644 --- a/gallery/templates/profile.html +++ b/gallery/templates/profile.html @@ -8,7 +8,7 @@
diff --git a/gallery/views/api.py b/gallery/views/api.py index 2c50364..0bbacaa 100644 --- a/gallery/views/api.py +++ b/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': diff --git a/gallery/views/settings.py b/gallery/views/settings.py index 04da5e4..94926e9 100644 --- a/gallery/views/settings.py +++ b/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')