S081 · AuthN vs AuthZ, Sessions & Password Storage

#!/usr/bin/env python3
"""auth_lab.py — a working authN/authZ implementation.
 
Run:
  pip install flask bcrypt PyJWT
  python auth_lab.py
 
Then:
  curl -X POST localhost:5000/signup -d 'email=a@a.com&password=hunter2'
  curl -X POST localhost:5000/login  -d 'email=a@a.com&password=hunter2' -c cookies.txt
  curl localhost:5000/me -b cookies.txt
  curl localhost:5000/admin -b cookies.txt   # should 403 (not admin)
"""
from __future__ import annotations
 
import secrets
import time
import uuid
from dataclasses import dataclass, field
from typing import Callable
 
import bcrypt
import jwt
from flask import Flask, jsonify, request, g
 
# --- Fake DBs (in-memory, threading unsafe — demo only) ---
USERS: dict[str, dict] = {}
SESSIONS: dict[str, dict] = {}   # session_id -> {user_id, expires_at}
JWT_SECRET = secrets.token_bytes(32)
JWT_ALG = "HS256"
 
app = Flask(__name__)
 
# --- Helpers ---
 
def hash_password(pw: str) -> bytes:
    """bcrypt with cost=12. Tune cost so hashing takes ~250ms on your hardware."""
    return bcrypt.hashpw(pw.encode(), bcrypt.gensalt(rounds=12))
 
 
def verify_password(pw: str, stored_hash: bytes) -> bool:
    """Constant-time compare via bcrypt's own routine."""
    return bcrypt.checkpw(pw.encode(), stored_hash)
 
 
def issue_session(user_id: str) -> str:
    sid = secrets.token_urlsafe(32)
    SESSIONS[sid] = {
        "user_id": user_id,
        "expires_at": time.time() + 3600,  # 1 hour
    }
    return sid
 
 
def issue_jwt(user_id: str, roles: list[str]) -> str:
    payload = {
        "sub": user_id,
        "roles": roles,
        "iat": int(time.time()),
        "exp": int(time.time()) + 900,   # 15 min
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALG)
 
 
# --- Middleware ---
 
def require_auth(fn: Callable) -> Callable:
    """Middleware: reject if no valid session OR JWT. Populate g.user."""
    def wrapper(*args, **kwargs):
        # Try session cookie first
        sid = request.cookies.get("sid")
        if sid and sid in SESSIONS:
            s = SESSIONS[sid]
            if s["expires_at"] > time.time():
                g.user = USERS.get(s["user_id"])
                if g.user:
                    return fn(*args, **kwargs)
        # Try JWT bearer
        auth = request.headers.get("Authorization", "")
        if auth.startswith("Bearer "):
            token = auth[7:]
            try:
                payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALG])
                g.user = USERS.get(payload["sub"])
                if g.user:
                    g.user["roles"] = payload.get("roles", [])
                    return fn(*args, **kwargs)
            except jwt.PyJWTError:
                pass
        return jsonify({"error": "unauthorized"}), 401
    wrapper.__name__ = fn.__name__
    return wrapper
 
 
def require_role(role: str) -> Callable:
    """Middleware: require an authenticated user with the given role."""
    def deco(fn: Callable) -> Callable:
        def wrapper(*args, **kwargs):
            user = getattr(g, "user", None)
            if not user or role not in user.get("roles", []):
                return jsonify({"error": "forbidden"}), 403
            return fn(*args, **kwargs)
        wrapper.__name__ = fn.__name__
        return wrapper
    return deco
 
 
# --- Rate limiting: simple per-email counter ---
LOGIN_ATTEMPTS: dict[str, list[float]] = {}
MAX_ATTEMPTS = 5
WINDOW = 300  # 5 min
 
 
def rate_limit_login(email: str) -> bool:
    now = time.time()
    attempts = LOGIN_ATTEMPTS.setdefault(email, [])
    # Prune old
    LOGIN_ATTEMPTS[email] = [t for t in attempts if now - t < WINDOW]
    if len(LOGIN_ATTEMPTS[email]) >= MAX_ATTEMPTS:
        return False
    LOGIN_ATTEMPTS[email].append(now)
    return True
 
 
# --- Routes ---
 
@app.route("/signup", methods=["POST"])
def signup():
    email = request.form.get("email", "").lower().strip()
    pw = request.form.get("password", "")
    if not email or len(pw) < 8:
        return jsonify({"error": "email required, password >= 8 chars"}), 400
    if email in USERS:
        return jsonify({"error": "user exists"}), 409
    user_id = str(uuid.uuid4())
    USERS[email] = {
        "id": user_id,
        "email": email,
        "hash": hash_password(pw),
        "roles": ["user"],  # default
    }
    return jsonify({"ok": True, "user_id": user_id}), 201
 
 
@app.route("/login", methods=["POST"])
def login():
    email = request.form.get("email", "").lower().strip()
    pw = request.form.get("password", "")
    # Rate-limit even before doing the expensive hash
    if not rate_limit_login(email):
        return jsonify({"error": "too many attempts, try again later"}), 429
 
    user = USERS.get(email)
    # CRITICAL: hash a dummy password if user doesn't exist, so timing
    # doesn't leak whether email is registered.
    dummy_hash = b"$2b$12$0000000000000000000000000000000000000000000000000000"
    stored = user["hash"] if user else dummy_hash
    ok = verify_password(pw, stored)
    if not user or not ok:
        return jsonify({"error": "invalid credentials"}), 401
 
    # Issue session cookie
    sid = issue_session(email)
    resp = jsonify({"ok": True, "roles": user["roles"]})
    resp.set_cookie("sid", sid, httponly=True, samesite="Strict",
                    secure=False)  # TRUE in production over HTTPS
    return resp
 
 
@app.route("/login/jwt", methods=["POST"])
def login_jwt():
    email = request.form.get("email", "").lower().strip()
    pw = request.form.get("password", "")
    if not rate_limit_login(email):
        return jsonify({"error": "rate limit"}), 429
    user = USERS.get(email)
    dummy_hash = b"$2b$12$0000000000000000000000000000000000000000000000000000"
    if not user or not verify_password(pw, user["hash"] if user else dummy_hash):
        return jsonify({"error": "invalid"}), 401
    token = issue_jwt(user["id"], user["roles"])
    return jsonify({"access_token": token, "expires_in": 900})
 
 
@app.route("/logout", methods=["POST"])
@require_auth
def logout():
    sid = request.cookies.get("sid")
    if sid:
        SESSIONS.pop(sid, None)
    resp = jsonify({"ok": True})
    resp.set_cookie("sid", "", expires=0)
    return resp
 
 
@app.route("/me")
@require_auth
def me():
    return jsonify({"email": g.user["email"], "roles": g.user["roles"]})
 
 
@app.route("/admin")
@require_auth
@require_role("admin")
def admin():
    return jsonify({"secret": "you're an admin"})
 
 
if __name__ == "__main__":
    # Create an admin for testing
    USERS["admin@a.com"] = {
        "id": "admin-1",
        "email": "admin@a.com",
        "hash": hash_password("adminpass"),
        "roles": ["user", "admin"],
    }
    print("Try:")
    print("  curl -X POST localhost:5000/signup -d 'email=a@a.com&password=hunter2'")
    print("  curl -X POST localhost:5000/login  -d 'email=a@a.com&password=hunter2' -c c.txt")
    print("  curl localhost:5000/me -b c.txt")
    print("  curl localhost:5000/admin -b c.txt   # 403")
    print("  curl -X POST localhost:5000/login -d 'email=admin@a.com&password=adminpass' -c ac.txt")
    print("  curl localhost:5000/admin -b ac.txt  # 200")
    app.run(host="127.0.0.1", port=5000)