functioning user login and register system

app/__init__.py

@@ -1,5 +1,15 @@
 from flask import Flask
+from flask_login import LoginManager
+from flask_wtf.csrf import CSRFProtect
+from flask_bcrypt import Bcrypt
+from config import Config
 
 app = Flask(__name__)
+app.config['SECRET_KEY'] = Config.secretKey
+csrf = CSRFProtect(app)
+bcrypt = Bcrypt(app)
+loginManager = LoginManager(app)
+loginManager.login_view = 'login'
+loginManager.login_message_category = 'info'
 
 from app import routes

app/models.py

@@ -0,0 +1,23 @@
+from flask_login import UserMixin
+from config import Config
+
+class User(UserMixin):
+    def __init__(self, user, userid, password, idpass, level):
+        self.user = user
+        self.password = password
+        self.userid = userid
+        self.idpass = idpass
+        self.level = level
+
+    def __repr__(self):
+        return f"User('{self.user}', '{self.userid}', '{self.password}', '{self.idpass}', '{self.level}')"
+
+    def get_id(self):
+        return str(self.userid)
+    
+    def get(userid):
+        userData = Config.users.find_one({"userid": userid})
+        if not userData:
+            return None
+        else:
+            return User(userData["user"], userData["userid"], userData["password"], userData["idpass"], userData["level"])

app/routes.py

@@ -1,6 +1,11 @@
-from app import app, worker
+from app import app, worker, bcrypt, loginManager
+from app.models import User
 from config import Config, Errors
-from flask import render_template, request, send_file, redirect
+from flask import render_template, request, send_file, redirect, flash
+from flask_login import login_user, current_user, logout_user, login_required
+from flask_wtf import FlaskForm
+from wtforms import StringField, PasswordField, SubmitField, BooleanField
+from wtforms.validators import DataRequired, Length, EqualTo
 from werkzeug.datastructures import FileStorage
 from werkzeug.utils import secure_filename
 from io import BytesIO
@@ -9,11 +14,34 @@ import io
 import random
 import magic
 
+class RegistrationForm(FlaskForm):
+    username = StringField('Username', validators=[DataRequired(), Length(min=2, max=16)])
+    password = PasswordField('Password', validators=[DataRequired(), Length(min=8, max=32)])
+    password2 = PasswordField('Confirm Password', validators=[DataRequired(), EqualTo('password')])
+    tnc = BooleanField('I agree to the Terms and Conditions', validators=[DataRequired()])
+    submit = SubmitField('Register')
+
+    def validate_username(self, username):
+        user = Config.users.find_one({"username": username.data})
+        if user:
+            raise ValueError("That username is taken. Try another.")
+        
+class LoginForm(FlaskForm):
+    username = StringField('Username', validators=[DataRequired(), Length(min=2, max=16)])
+    password = PasswordField('Password', validators=[DataRequired(), Length(min=8, max=32)])
+    submit = SubmitField('Login')
+
+@loginManager.user_loader
+def load_user(userid):
+    user = User.get(userid)
+    return user
+
 @app.route('/', methods=["GET", "POST"])
 def index():
     
     # Check for a GET or POST request
     if request.method == "GET":
+        print(current_user.is_authenticated)
         return render_template('index.html')
 
     elif request.method == "POST":
@@ -95,4 +123,52 @@ def getInfo(id):
 
 @app.route('/teapot')
 def teapot():
-    return 'I\'m a teapot. 418.', 418
+    return 'I\'m a teapot. 418.', 418
+
+@app.route('/register', methods=["GET", "POST"])
+def register():
+    if current_user.is_authenticated:
+        return redirect("/")
+    else:
+        if request.method == "GET":
+            return render_template("register.html", form=RegistrationForm())
+        elif request.method == "POST":
+            username = request.form.get("username")
+            password = request.form.get("password")
+
+            res = worker.registerUser(username, password)
+
+            if res == True:
+                flash("Successfully registered!", "success")
+                return redirect("/login")
+            else:
+                flash("Something went wrong, sorry.", "danger")
+                return redirect("/register")
+
+@app.route('/login', methods=["GET", "POST"])
+def login():
+    if current_user.is_authenticated:
+        return redirect("/")
+    else:
+        if request.method == "GET":
+            return render_template("login.html", form=LoginForm())
+        elif request.method == "POST":
+            username = request.form.get("username")
+            password = request.form.get("password")
+
+            userid = Config.users.find_one({"user": username})["userid"]
+            user = User.get(userid)
+
+            if user and bcrypt.check_password_hash(user.password, password):
+                login_user(user)
+                print(current_user.is_authenticated)
+                flash("Successfully logged in!", "success")
+                return redirect("/")
+            else:
+                flash("Incorrect username or password.", "danger")
+                return redirect("/login")
+            
+@app.route('/logout')
+def logout():
+    logout_user()
+    return redirect("/")

app/templates/base.html

@@ -0,0 +1,12 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta ref="stylesheet" href="{{ url_for('static', filename='css/style.css') }}">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>xygt.cc - {{ title }}</title>
+</head>
+<body>
+    {% block content %}{% endblock %}
+</body>
+</html>

app/templates/index.html

@@ -1 +1,12 @@
-Under Construction - No frontend available yet.
+{% extends 'base.html' %}
+
+{% block content %}
+    <h1>Home</h1>
+    <!-- Check if the user is logged in -->
+    {% if current_user.is_authenticated %}
+        <p>Welcome, {{ current_user.user }}!</p>
+        <p><a href="{{ url_for('logout') }}">Logout</a></p>
+    {% else %}
+        <p><a href="{{ url_for('login') }}">Login</a></p>
+    {% endif %}
+{% endblock %}

app/templates/login.html

@@ -0,0 +1,21 @@
+{% extends 'base.html' %}
+
+{% block content %}
+    <h1>Login</h1>
+    <form action="" method="post">
+        {{ form.hidden_tag() }}
+        {% for field in form %}
+            <p>
+                {{ field.label }}<br>
+                {{ field() }}
+                {% if field.errors %}
+                    <ul class=errors>
+                        {% for category, message in messages %}
+                            <li class="{{ category }}">{{ message }}</li>
+                        {% endfor %}
+                    </ul>
+                {% endif %}
+            </p>
+        {% endfor %}
+    </form>
+{% endblock %}

app/templates/register.html

@@ -0,0 +1,21 @@
+{% extends 'base.html' %}
+
+{% block content %}
+    <h1>Register</h1>
+    <form action="" method="post">
+        {{ form.hidden_tag() }}
+        {% for field in form %}
+            <p>
+                {{ field.label }}<br>
+                {{ field() }}
+                {% if field.errors %}
+                    <ul class=errors>
+                    {% for error in field.errors %}
+                        <li>{{ error }}</li>
+                    {% endfor %}
+                    </ul>
+                {% endif %}
+            </p>
+        {% endfor %}
+    </form>
+{% endblock %}

app/worker.py

@@ -1,4 +1,6 @@
 from config import disallowedMimeTypes, Errors, Config
+from app.models import User
+from app import bcrypt
 import secrets
 import datetime
 import random
@@ -67,7 +69,6 @@ def shortenURL(url, ip, userid, id, retention):
         while True:                 # Loop to find an available file ID
             id = randomHex()        # Prevent conflicts if 2 of the same get made
             if Config.files.find_one({'id': id}) is None:
-                filename = id
                 break
 
     if userid == None:
@@ -108,3 +109,17 @@ def idInfo(id):
 def randomHex():
     hexRand = ''.join(secrets.choice('0123456789abcdef') for _ in range(6))
     return hexRand
+
+def registerUser(username, password):
+    # Initialise some values
+    try:
+        level = 1
+        userid = randomHex()
+        idpass = bcrypt.generate_password_hash(randomHex()).decode("utf-8")
+        password = bcrypt.generate_password_hash(password).decode("utf-8")
+        user = User(username, userid, password, idpass, level)
+        Config.users.insert_one(user.__dict__)
+
+        return True
+    except:
+        return False

config.py

@@ -1,5 +1,6 @@
 import os
 from pymongo import MongoClient
+import random
 
 class Config:
     # MongoDB init stuff
@@ -16,6 +17,7 @@ class Config:
     minretention = 7
     fileDir = "./data"
     ipLogEnabled = False
+    secretKey = "CHANGEINPRODUCTION"
 
 quotes = {
     "Anon /g/": "Biometrics are shit, you lose a limb and you're fucked.",

requirements.txt

@@ -1,4 +1,7 @@
 flask
 flask-dotenv
 pymongo
-python-magic
+python-magic
+flask-login
+flask-bcrypt
+flask-wtf