Authentication & Flask-Login

import functools
import os
import secrets
import sqlite3
import tempfile

from flask import Blueprint, Flask, current_app, g, redirect, request, session, url_for
from werkzeug.security import check_password_hash, generate_password_hash

flaskr_auth = Blueprint("auth", __name__, url_prefix="/auth")


def get_db() -> sqlite3.Connection:
    if "db" not in g:
        g.db = sqlite3.connect(current_app.config["DATABASE"])
        g.db.row_factory = sqlite3.Row
    return g.db


def close_db(error=None):
    db = g.pop("db", None)
    if db is not None:
        db.close()


@flaskr_auth.before_app_request
def load_logged_in_user():
    user_id = session.get("user_id")
    g.user = None
    if user_id is not None:
        g.user = get_db().execute("SELECT * FROM user WHERE id = ?", (user_id,)).fetchone()


def login_required(view):
    @functools.wraps(view)
    def wrapped_view(**kwargs):
        if g.user is None:
            return redirect(url_for("auth.login"))
        return view(**kwargs)

    return wrapped_view


@flaskr_auth.post("/register")
def register():
    username, password = request.form["username"], request.form["password"]
    if not username or len(password) < 12:
        return "Username required and password of at least 12 characters", 400
    db = get_db()
    try:
        db.execute("INSERT INTO user (username, password) VALUES (?, ?)", (username, generate_password_hash(password)))
        db.commit()
    except sqlite3.IntegrityError:
        return f"User {username} is already registered.", 400
    return redirect(url_for("auth.login"))


@flaskr_auth.route("/login", methods=("GET", "POST"))
def login():
    if request.method == "POST":
        user = get_db().execute("SELECT * FROM user WHERE username = ?", (request.form["username"],)).fetchone()
        if user is None or not check_password_hash(user["password"], request.form["password"]):
            return "Incorrect username or password.", 401
        session.clear()
        session["user_id"] = user["id"]
        return redirect(url_for("index"))
    return "login form"


@flaskr_auth.get("/logout")
def logout():
    session.clear()
    return redirect(url_for("index"))


def create_flaskr_app() -> Flask:
    app = Flask(__name__)
    app.config["SECRET_KEY"] = secrets.token_hex(32)
    app.config["DATABASE"] = os.path.join(tempfile.mkdtemp(), "flaskr.sqlite")
    app.teardown_appcontext(close_db)
    app.register_blueprint(flaskr_auth)

    @app.get("/")
    @login_required
    def index():
        return f"Hello {g.user['username']}"

    with app.app_context():
        get_db().execute("CREATE TABLE IF NOT EXISTS user (id INTEGER PRIMARY KEY, username TEXT UNIQUE NOT NULL, password TEXT NOT NULL)")
    return app


flaskr = create_flaskr_app()
client = flaskr.test_client()
user_secret = secrets.token_urlsafe(16)
assert client.get("/").status_code == 302  # anonymous: redirected to /auth/login
assert client.post("/auth/register", data={"username": "ana", "password": user_secret}).status_code == 302
assert client.post("/auth/login", data={"username": "ana", "password": user_secret[::-1]}).status_code == 401
assert client.post("/auth/login", data={"username": "ana", "password": user_secret}).status_code == 302
assert client.get("/").get_data(as_text=True) == "Hello ana"
client.get("/auth/logout")
assert client.get("/").status_code == 302

pip install flask-login flask-sqlalchemy flask-wtf

account = Blueprint("account", __name__, url_prefix="/account")


@account.get("/")
@login_required
def overview():
    return jsonify(email=current_user.email, fresh=login_fresh())


@account.route("/reauth", methods=["GET", "POST"])
@login_required
def reauth():
    if request.method == "POST":
        if current_user.check_password(request.form.get("password", "")):
            confirm_login()  # marks the session fresh again
            return redirect_after_login()
        flash("Wrong password.", "error")
    return render_template_string(PAGE, fields=NEW_PASSWORD)


@account.route("/password", methods=["GET", "POST"])
@fresh_login_required
def change_password():
    if request.method == "POST":
        problem = password_problem(request.form.get("password", ""))
        if problem:
            flash(problem, "error")
            return render_template_string(PAGE, fields=NEW_PASSWORD), 400
        user = current_user._get_current_object()  # the real User behind the proxy
        user.set_password(request.form["password"])  # rotates session_token
        db.session.commit()
        login_user(user)  # new get_id() for this browser; every other device is signed out
        flash("Password changed.", "success")
        return redirect(url_for("account.overview"))
    return render_template_string(PAGE, fields=NEW_PASSWORD)

api = Blueprint("api", __name__, url_prefix="/api")


def api_serializer() -> URLSafeTimedSerializer:
    return URLSafeTimedSerializer(current_app.config["SECRET_KEY"], salt="bookshelf-api-token")


def issue_api_token(user: User) -> str:
    return api_serializer().dumps(user.get_id())


@login_manager.request_loader
def load_user_from_bearer_token(req):
    scheme, _, token = req.headers.get("Authorization", "").partition(" ")
    if scheme.lower() != "bearer" or not token:
        return None
    try:
        user_id = api_serializer().loads(token, max_age=3600)  # one hour
    except BadSignature:  # SignatureExpired is a subclass
        return None
    return load_user(user_id)


@api.get("/me")
@login_required
def me():
    return jsonify(email=current_user.email)

from flask_login import session_protected, user_logged_in, user_logged_out, user_login_confirmed

audit_log = logging.getLogger("bookshelf.audit")
AUDIT: list[str] = []  # also kept in memory so the example can assert on it


@user_logged_in.connect
def audit_login(sender, user, **extra):
    AUDIT.append(f"login:{user.id}")
    audit_log.info("user %s signed in", user.id)


@user_logged_out.connect
def audit_logout(sender, user, **extra):
    AUDIT.append(f"logout:{user.id}")


@user_login_confirmed.connect
def audit_confirmed(sender, **extra):
    AUDIT.append("confirmed")


@session_protected.connect
def audit_protected(sender, **extra):
    AUDIT.append("protected")

OUTBOX: list[tuple[str, str]] = []  # stands in for the mail system (use Flask-Mail or a queue)


def reset_serializer() -> URLSafeTimedSerializer:
    return URLSafeTimedSerializer(current_app.config["SECRET_KEY"], salt="bookshelf-password-reset")


def fingerprint(user: User) -> str:
    return hashlib.sha256(user.password_hash.encode()).hexdigest()[:16]


def make_reset_token(user: User) -> str:
    return reset_serializer().dumps({"uid": user.id, "fp": fingerprint(user)})


def user_from_reset_token(token: str, max_age: int = 1800) -> User | None:
    try:
        data = reset_serializer().loads(token, max_age=max_age)  # raises when forged or older than 30 minutes
    except BadSignature:
        return None
    user = db.session.get(User, data["uid"])
    if user is None or not user.is_active or not secrets.compare_digest(data["fp"], fingerprint(user)):
        return None
    return user


@auth.route("/forgot", methods=["GET", "POST"])
def forgot():
    if request.method == "POST":
        user = find_user(request.form.get("email", ""))
        if user is not None and user.is_active:
            link = url_for("auth.reset", token=make_reset_token(user), _external=True)
            OUTBOX.append((user.email, link))  # queue the e-mail; never send it inline
        flash("If that address has an account, a reset link is on its way.", "message")  # same answer either way
        return redirect(url_for("auth.login"))
    return render_template_string(PAGE, fields=EMAIL_ONLY)


@auth.route("/reset/<token>", methods=["GET", "POST"])
def reset(token):
    user = user_from_reset_token(token)
    if user is None:
        abort(404)
    if request.method == "POST":
        problem = password_problem(request.form.get("password", ""))
        if problem is None:
            user.set_password(request.form["password"])  # new fingerprint and new session_token
            db.session.commit()
            flash("Password updated. Please sign in.", "success")
            return redirect(url_for("auth.login"))
        flash(problem, "error")
    response = make_response(render_template_string(PAGE, fields=NEW_PASSWORD))
    response.headers["Referrer-Policy"] = "no-referrer"  # the token is in the URL: do not leak it
    response.headers["Cache-Control"] = "no-store"
    return response

from werkzeug.security import check_password_hash, generate_password_hash  # already imported above; shown for clarity


def no_store_for_signed_in_users(response):
    if current_user.is_authenticated:
        response.headers["Cache-Control"] = "no-store"
    return response


def create_app(test_config: dict | None = None) -> Flask:
    app = Flask(__name__)
    app.config.from_mapping(
        SECRET_KEY=os.environ.get("BOOKSHELF_SECRET_KEY") or secrets.token_hex(32),
        SQLALCHEMY_DATABASE_URI=os.environ.get("BOOKSHELF_DATABASE_URL", "sqlite://"),
        PREFERRED_URL_SCHEME="https",
        SESSION_COOKIE_SECURE=True,
        SESSION_COOKIE_HTTPONLY=True,
        SESSION_COOKIE_SAMESITE="Lax",
        PERMANENT_SESSION_LIFETIME=timedelta(hours=8),
        REMEMBER_COOKIE_DURATION=timedelta(days=14),
        REMEMBER_COOKIE_SECURE=True,
        REMEMBER_COOKIE_HTTPONLY=True,
        REMEMBER_COOKIE_SAMESITE="Lax",
    )
    if test_config:
        app.config.from_mapping(test_config)
    db.init_app(app)
    login_manager.init_app(app)
    csrf.init_app(app)
    for blueprint in (main, auth, account, api):
        app.register_blueprint(blueprint)
    app.after_request(no_store_for_signed_in_users)
    with app.app_context():
        db.create_all()
    return app

TOKEN_RE = re.compile(r'name="csrf_token" value="([^"]+)"')


def form_token(client, url: str) -> str:
    return TOKEN_RE.search(client.get(url).get_data(as_text=True)).group(1)


def post(client, url: str, form_page: str | None = None, headers: dict | None = None, **data):
    data["csrf_token"] = form_token(client, form_page or url)
    # Over HTTPS, Flask-WTF also checks the Referer header, which a browser always sends (WTF_CSRF_SSL_STRICT).
    return client.post(url, data=data, headers={"Referer": "https://localhost/", **(headers or {})})


def make_user(app: Flask, email: str, password: str) -> None:
    with app.app_context():
        user = User(email=email, password_hash="")
        user.set_password(password)
        db.session.add(user)
        db.session.commit()


app = create_app()
PASSWORD = secrets.token_urlsafe(18)
make_user(app, "ana@example.com", PASSWORD)


def test_anonymous_and_login():
    client = app.test_client()
    response = client.get("/shelf")
    assert response.status_code == 302 and response.headers["Location"].startswith("/auth/login?next=")
    assert client.post("/auth/login", data={"email": "ana@example.com"}).status_code == 400  # no CSRF token
    bad = post(client, "/auth/login", email="ana@example.com", password="not the password")
    assert bad.status_code == 401
    good = post(client, "/auth/login", email="ana@example.com", password=PASSWORD)
    assert good.status_code == 302 and good.headers["Location"] == "/shelf"
    assert client.get("/shelf").get_data(as_text=True) == "Shelf of ana@example.com"
    assert "no-store" in client.get("/shelf").headers["Cache-Control"]
    out = post(client, "/auth/logout", form_page="/account/password")
    assert out.status_code == 302
    assert client.get("/shelf").status_code == 302


def test_open_redirect_is_refused():
    client = app.test_client()
    for evil in ("https://evil.example/", "//evil.example", "/\\evil.example", "javascript:alert(1)"):
        response = post(client, f"/auth/login?next={evil}", form_page="/auth/login", email="ana@example.com", password=PASSWORD)
        assert response.headers["Location"] == "/shelf", evil
        post(client, "/auth/logout", form_page="/account/password")
    local = post(client, "/auth/login?next=/account/", form_page="/auth/login", email="ana@example.com", password=PASSWORD)
    assert local.headers["Location"] == "/account/"


def test_api_gets_json_401_and_bearer_token():
    client = app.test_client()
    response = client.get("/api/me")
    assert response.status_code == 401 and response.json == {"error": "authentication required"}
    with app.app_context():
        token = issue_api_token(db.session.execute(db.select(User)).scalar_one())
    ok = client.get("/api/me", headers={"Authorization": f"Bearer {token}"})
    assert ok.status_code == 200 and ok.json == {"email": "ana@example.com"}
    forged = client.get("/api/me", headers={"Authorization": f"Bearer {token}x"})
    assert forged.status_code == 401


test_anonymous_and_login()
test_open_redirect_is_refused()
test_api_gets_json_401_and_bearer_token()

def test_remember_me_and_freshness():
    client = app.test_client()
    post(client, "/auth/login", email="ana@example.com", password=PASSWORD, remember="on")
    assert client.get_cookie("remember_token") is not None
    assert client.get("/account/").json["fresh"] is True
    client.delete_cookie("session")  # the browser session ends; only remember_token is left
    restored = client.get("/account/")
    assert restored.status_code == 200 and restored.json["fresh"] is False  # remembered, not fresh
    redirected = client.get("/account/password")
    assert redirected.status_code == 302 and redirected.headers["Location"].startswith("/account/reauth")
    assert post(client, "/account/reauth", password=PASSWORD[::-1]).status_code == 200  # form shown again
    confirmed = post(client, "/account/reauth", password=PASSWORD)
    assert confirmed.status_code == 302
    assert client.get("/account/").json["fresh"] is True
    assert client.get("/account/password").status_code == 200
    post(client, "/auth/logout", form_page="/account/password")
    assert client.get_cookie("remember_token") is None  # logout removes the remember cookie too
    assert client.get("/shelf").status_code == 302


def test_password_change_rotates_other_sessions():
    laptop, phone = app.test_client(), app.test_client()
    for browser in (laptop, phone):
        post(browser, "/auth/login", email="ana@example.com", password=PASSWORD)
    new_password = secrets.token_urlsafe(18)
    assert post(laptop, "/account/password", password=new_password).status_code == 302
    assert laptop.get("/shelf").status_code == 200  # this browser got a new identifier through login_user
    assert phone.get("/shelf").status_code == 302   # every other session is signed out
    assert post(phone, "/auth/login", email="ana@example.com", password=PASSWORD).status_code == 401  # old password
    assert post(phone, "/auth/login", email="ana@example.com", password=new_password).status_code == 302
    post(laptop, "/account/password", password=PASSWORD)  # restore for the next tests


test_remember_me_and_freshness()
test_password_change_rotates_other_sessions()