portal.py

import threading
from functools import wraps

from flask import Flask, render_template, request, redirect, make_response

import db
import jwt_auth

app = Flask(__name__)

COOKIE_NAME = "portal_token"


def current_user():
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return None
    username = jwt_auth.verify_token(token)
    if username is None:
        return None
    return db.get_user(username)


def login_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        user = current_user()
        if user is None:
            return redirect("/portal/login")
        return f(user, *args, **kwargs)
    return wrapper


@app.route("/portal/login", methods=["GET", "POST"])
def login():
    error = None
    if request.method == "POST":
        user = db.check_login(request.form["username"], request.form["password"])
        if user is None:
            error = "invalid username or password"
        else:
            remember = bool(request.form.get("remember"))
            token = jwt_auth.make_token(user["username"], remember)

            resp = make_response(redirect("/portal"))
            max_age = jwt_auth.REMEMBER_DAYS * 86400 if remember else jwt_auth.NORMAL_HOURS * 3600
            resp.set_cookie(COOKIE_NAME, token, httponly=True,
                            max_age=max_age, samesite="Lax")
            return resp
    return render_template("portal_login.html", error=error)


@app.route("/portal/logout")
def logout():
    resp = make_response(redirect("/portal/login"))
    resp.delete_cookie(COOKIE_NAME)
    return resp


@app.route("/portal")
@login_required
def account(user):
    pct = (user["used_bytes"] * 100 // user["quota_bytes"]) if user["quota_bytes"] else 100
    remaining = max(0, user["quota_bytes"] - user["used_bytes"])
    return render_template(
        "portal_account.html", user=user, pct=min(pct, 100),
        remaining=remaining, human=db.human,
        purchases=db.get_purchases(user["username"]))


@app.route("/portal/buy", methods=["POST"])
@login_required
def buy(user):
    mb = int(request.form["mb"])
    if mb > 0:
        db.add_quota(user["username"], mb * 1024 ** 2)
        db.log_purchase(user["username"], mb)
    return redirect("/portal")


def start(port=5001):
    t = threading.Thread(
        target=lambda: app.run(host="0.0.0.0", port=port,
                               debug=False, use_reloader=False),
        daemon=True)
    t.start()
    print(f"user portal on http://0.0.0.0:{port}/portal")
    return t


if __name__ == "__main__":
    db.init_db()
    app.run(host="0.0.0.0", port=5001, debug=True)