Switch to SQLAlchemy for database management

Add Favicon

gallery/auth.py

@@ -2,7 +2,10 @@ import functools
 from flask import Blueprint, flash, g, redirect, request, session, url_for, abort, jsonify, current_app
 from werkzeug.security import check_password_hash, generate_password_hash
 
-from gallery.db import get_db
+from gallery import db
+from sqlalchemy.orm import sessionmaker
+db_session = sessionmaker(bind=db.engine)
+db_session = db_session()
 
 from .logger import logger
 
@@ -12,45 +15,26 @@ import uuid
 blueprint = Blueprint('auth', __name__, url_prefix='/auth')
 
 
-# def add_log(code, note=None):    
-#     code = int(code)
-#     note = str(note)
-    
-#     user_id = session.get('user_id')
-#     user_ip = request.remote_addr
-#     db = get_db()
-    
-#     db.execute(
-#         'INSERT INTO logs (ip, user_id, code, note)'
-#         ' VALUES (?, ?, ?, ?)',
-#         (user_ip, user_id, code, note)
-#     )
-#     db.commit()
-
-
 @blueprint.before_app_request
 def load_logged_in_user():
     user_id = session.get('user_id')
     user_uuid = session.get('uuid')
 
     if user_id is None or user_uuid is None:
-        # This is not needed as the user is not logged in anyway, also spams the logs
+        # This is not needed as the user is not logged in anyway, also spams the server logs with useless data
         #add_log(103, 'Auth error before app request')
         g.user = None
         session.clear()
     else:
-        db = get_db()
-        is_alive = db.execute('SELECT * FROM devices WHERE session_uuid = ?',
-                              (session.get('uuid'), )).fetchone()
+        is_alive = db_session.query(db.sessions).filter_by(session_uuid=user_uuid).first()
 
         if is_alive is None:
             logger.add(103, 'Session expired')
             flash(['Session expired!', '3'])
             session.clear()
         else:
-            g.user = db.execute('SELECT * FROM users WHERE id = ?',
-                                (user_id, )).fetchone()
-
+            g.user = db_session.query(db.users).filter_by(id=user_id).first()
+            
 
 @blueprint.route('/register', methods=['POST'])
 def register():
@@ -58,7 +42,6 @@ def register():
     email = request.form['email']
     password = request.form['password']
     password_repeat = request.form['password-repeat']
-    db = get_db()
     error = []
 
     if not username:
@@ -82,12 +65,10 @@ def register():
 
     if not error:
         try:
-            db.execute(
-                'INSERT INTO users (username, email, password) VALUES (?, ?, ?)',
-                (username, email, generate_password_hash(password)),
-            )
-            db.commit()
-        except db.IntegrityError:
+            tr = db.users(username, email, generate_password_hash(password))
+            db_session.add(tr)
+            db_session.commit()
+        except Exception as e:
             error.append(f"User {username} is already registered!")
         else:
             logger.add(103, f"User {username} registered")
@@ -100,27 +81,24 @@ def register():
 def login():
     username = request.form['username']
     password = request.form['password']
-    db = get_db()
     error = None
-    user = db.execute('SELECT * FROM users WHERE username = ?',
-                      (username, )).fetchone()
+    user = db_session.query(db.users).filter_by(username=username).first()
 
     if user is None:
         logger.add(101, f"User {username} does not exist from {request.remote_addr}")
         abort(403)
-    elif not check_password_hash(user['password'], password):
+    elif not check_password_hash(user.password, password):
         logger.add(102, f"User {username} password error from {request.remote_addr}")
         abort(403)
 
     try:
         session.clear()
-        session['user_id'] = user['id']
+        session['user_id'] = user.id
         session['uuid'] = str(uuid.uuid4())
-
-        db.execute(
-            'INSERT INTO devices (user_id, session_uuid, ip) VALUES (?, ?, ?)',
-            (user['id'], session.get('uuid'), request.remote_addr))
-        db.commit()
+        
+        tr = db.sessions(user.id, session.get('uuid'), request.remote_addr, request.user_agent.string, 1)
+        db_session.add(tr)
+        db_session.commit()
     except error as err:
         logger.add(105, f"User {username} auth error: {err}")
         abort(500)
@@ -135,7 +113,7 @@ def login():
 
 @blueprint.route('/logout')
 def logout():
-    logger.add(103, f"User {g.user['username']} - id: {g.user['id']} logged out")
+    logger.add(103, f"User {g.user.username} - id: {g.user.id} logged out")
     session.clear()
     return redirect(url_for('index'))