Source code for src.auth

import os
import random
import re
from typing import Any

from src.apputils import PathUtils

from flask import Flask, abort
from flask_login import LoginManager, UserMixin, current_user
from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField
from wtforms.validators import DataRequired
from functools import wraps
import sqlite3
import logging

logger = logging.getLogger(__name__)

[docs] def text_slug(text: str) -> str: text = text.lower() text = re.sub(r'[^a-z0-9\s-]', '', text) text = re.sub(r'[\s-]+', '-', text) return text.strip('-')
ADMIN_COLOR = "#b000f0" DELETED_COLOR = "#585858" LOGIN_TABLE_HEADERS = [ "id", "un", "pw", "isadmin", "color" ] USER_PALATTE = [ "#ff9500", "#00d5ff", "#00f0a0", "#f03000", "#f00098", "#f0d000" ]
[docs] def display_user(id: str) -> str: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM logins WHERE id=?", (id,)) rows = cursor.fetchall() if len(rows) > 0: user = (*rows[0]) else: user = ("deleted", "Deleted User", "", 0, DELETED_COLOR) return f""" <style> .user-pill {{ font-family: monospace; display: inline-flex; align-items: center; gap: 8px; padding: 6px 12px; border-radius: 999px; background: #444; font-size: 14px; }} .avatar {{ width: 24px; height: 24px; border-radius: 50%; text-align: center; font-size: 20px; line-height: 24px; }} .name {{ white-space: nowrap; }} </style> <div class="user-pill"> <div class="avatar" style="background-color: {user.color};">{user.un[0].upper()}</div> <span class="name">{user.un}</span> </div> """
[docs] def generate_random_color(isadmin: bool) -> str: return ADMIN_COLOR if isadmin else USER_PALATTE[random.randint(0, len(USER_PALATTE) - 1)]
[docs] def require_admin(f): """Ensures that the user is logged into an admin account on the modifiee endpoint""" @wraps(f) def decorated(*args, **kwargs): if not current_user.is_authenticated: abort(401) # unauthorized if not current_user.is_admin: abort(403) # forbidden return f(*args, **kwargs) return decorated
[docs] class (UserMixin): """ A user """ def __init__(self, id, un, pwd, isadmin, color): super().__init__() self.id = id self.un = un self.pwd = pwd self.is_admin = (True if isadmin == 1 else False) self.color = color
[docs] class LoginForm(FlaskForm): """Form used on the login page to log in""" username = StringField("Username", validators=[DataRequired()]) password = PasswordField("Password", validators=[DataRequired()]) submit = SubmitField("Log in")
[docs] class CreateAccount(FlaskForm): """Form used on the login page to log in""" username = StringField("Username", validators=[DataRequired()], render_kw={ "placeholder": "Enter username (ie. 'Frodo Baggins')..." }) password = PasswordField("Password", validators=[DataRequired()], render_kw={ "placeholder": "Enter password..." }) submit = SubmitField("Create Account")
login_manager = LoginManager() login_manager.login_view = "login"
[docs] @login_manager.user_loader def get_user_from_db(uid) -> | None: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM logins WHERE id=?", (uid,)) rows = cursor.fetchall() if len(rows) > 0: return (*rows[0]) else: return None
[docs] def get_user_is_admin(user: UserMixin): if isinstance(user, ): return user.is_admin return False
[docs] def get_can_delete_user(user_uid, uid_to_delete) -> bool: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM logins WHERE isadmin=1") rows = cursor.fetchall() if uid_to_delete not in [x[0] for x in rows]: logger.info(f"Deletion of {uid_to_delete} authorized: not an admin") return True elif uid_to_delete == user_uid: logger.info(f"Deletion of {uid_to_delete} {( "authorized: admins can delete themselves" if len(rows) > 1 else "prohibited: user is last admin" )}") return len(rows) > 1 else: return False # else no
[docs] def override_get_admin() -> | None: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM logins WHERE isadmin=1") rows = cursor.fetchall() if len(rows) > 0: return (*rows[0]) else: return None
[docs] def get_users() -> list[tuple]: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM logins") return cursor.fetchall()
[docs] def get_user_from_db_unpw(un, pwd) -> | None: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT * FROM logins WHERE un=? AND pw=?", (un, pwd)) rows = cursor.fetchall() if len(rows) > 0: return (*rows[0]) else: return None
[docs] def delete_user(my_uid, uid) -> bool: if not get_can_delete_user(my_uid, uid): return False with sqlite3.connect(PathUtils.file_set.login_db) as conn: conn.execute("DELETE FROM logins WHERE id=?", (uid,)) return True
[docs] def init_loginm_app(app: Flask) -> None: """Registers a login manager for admin to the flask app""" login_manager.init_app(app)
[docs] def add_user_to_db(username, password, is_admin=False) -> None: with sqlite3.connect(PathUtils.file_set.login_db) as conn: conn.execute("INSERT OR REPLACE INTO logins (id, un, pw, isadmin, color) VALUES (?, ?, ?, ?, ?)", ( text_slug(username), username, password, is_admin, generate_random_color(is_admin) ))
[docs] def get_password_is_admin(password) -> Any | bool: with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT isadmin FROM logins WHERE pw=?", (password,)) rows = cursor.fetchall() if len(rows) > 0: return rows[0] == 1 else: return False
[docs] def get_table_exists() -> bool: if not os.path.exists(PathUtils.file_set.login_db): return False with sqlite3.connect(PathUtils.file_set.login_db) as conn: cursor = conn.cursor() cursor.execute("SELECT name FROM sqlite_master WHERE type='table'") tablenames = cursor.fetchall() if len(tablenames) <= 0: return False cursor.execute("SELECT isadmin FROM logins") return any([x == 1 for x in cursor.fetchall()]) # make sure there's an admin
[docs] def generate_login_db(un, pwd) -> None: if get_table_exists(): return with sqlite3.connect(PathUtils.file_set.login_db) as conn: conn.executescript(f""" CREATE TABLE IF NOT EXISTS logins ( id TEXT PRIMARY KEY, un TEXT, pw TEXT, isadmin INT, color TEXT ); """) add_user_to_db(un, pwd, True)