lapp/auth.py

from flask import Blueprint, render_template, request, redirect, url_for, flash, make_response
from flask_login import login_user, logout_user, login_required, current_user
from models import db
from models import User
from log import Log
import pickle
from cryptography.fernet import Fernet
import datetime


auth_bp = Blueprint("auth", __name__)

# =====================================================
# SECRET KEY (only for this program)
# =====================================================
# Generate once using: Fernet.generate_key()
_SECRET_KEY = b'YFK7QCyTzhyLO4vqrnRxvDAI5uu8mXEYrInEjbRoQgs='
fernet = Fernet(_SECRET_KEY)

# =====================================================
# ENCRYPT
# =====================================================
def encrypt_object(obj) -> bytes:
    """
    Encrypt any Python object and return encrypted bytes.
    """
    serialized = pickle.dumps(obj)
    encrypted = fernet.encrypt(serialized)
    return encrypted

# =====================================================
# DECRYPT
# =====================================================
def decrypt_object(encrypted_data: bytes):
    """
    Decrypt bytes back into the original Python object.
    """
    decrypted = fernet.decrypt(encrypted_data)
    obj = pickle.loads(decrypted)
    return obj

def create_cookie(userid):
    cookie_data = {
        "user": userid,
        "datetime": datetime.datetime.now()
    }
    return encrypt_object(cookie_data).decode('utf-8')

def is_cookie_ok_to_autologin(cookie_content):
    user = 0
    autologin = False
    try:
        cookie_data = decrypt_object(cookie_content)
        user = cookie_data["user"]
        dt = cookie_data["datetime"]
        lastweek = datetime.datetime.now() - datetime.timedelta(days=7)
        autologin = dt > lastweek
    except:
        Log.info("Cookie error")
    if autologin:
        return user
    return 0

#. @auth_bp.route("/", methods=["GET", "POST"])
@auth_bp.route("/login", methods=["GET", "POST"])
def login():
    username = request.cookies.get('username')

    userid = is_cookie_ok_to_autologin(username)
    if userid>0:
        user = User.query.filter_by(id=userid).first()
        if user and user.is_approved and user.group_id>0:
            login_user(user)
            return redirect(url_for("lists.home"))
    
    if request.method == "POST":
        user = User.query.filter_by(name=request.form["name"]).first()
        if user and user.check_password(request.form["password"]):
            if not user.is_approved:
                return render_template("pending.html")
            if user.group_id>0:
                login_user(user)
                return redirect(url_for("lists.home"))
        flash("Invalid key credentials.")
        Log.info(f"Authorization failure for user '{request.form["name"]}' ")

    
    resp = make_response(render_template("login.html"))
    resp.set_cookie('username', '')
    return resp
    

@auth_bp.route("/register", methods=["GET", "POST"])
def register():
    if request.method == "POST":
        name = request.form["name"].strip()
        group = request.form["group"]
        password = request.form["password"]

        existing_user = User.query.filter_by(name=name).first()
        if existing_user:
            return render_template("register.html", error="That name is already registered.")

        user = User(name=name)
        group_id = user.set_group_id(group)
        if group_id == 0:
            return render_template("register.html", error="Unknown registration key.")

        user.set_password(password)
        user.group_id = group_id
        db.session.add(user)
        db.session.commit()
        return render_template("pending.html")
    return render_template("register.html")

@auth_bp.route("/logout")
@login_required
def logout():
    logout_user()
    resp = make_response(redirect(url_for("auth.login")))
    resp.set_cookie('username', '')
    return resp


@auth_bp.route("/change_pwd", methods=["GET", "POST"])
@login_required
def change_pwd():
    if request.method == "POST":
        pwd_old = request.form["password_old"]
        pwd_new1 = request.form["password_new1"]
        pwd_new2 = request.form["password_new2"]
        if current_user.check_password(pwd_old):
            if pwd_new1 == pwd_new2:
                current_user.set_password(pwd_new2)
                db.session.commit()
                return redirect(url_for("lists.home"))
            flash("New Passwords are not identical")
        flash("Invalid credentials")
        Log.info(f"Authorization failure for user '{request.form["name"]}' ")

        return render_template(url_for("lists.home"))
    return render_template("change_pwd.html")