Authorization & Roles

import secrets
from functools import wraps

from flask import Blueprint, Flask, abort, jsonify, request
from flask_login import LoginManager, UserMixin, current_user, login_required, login_user
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import Column, ForeignKey, Table, select
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from werkzeug.security import check_password_hash, generate_password_hash


class Base(DeclarativeBase):
    pass


db = SQLAlchemy(model_class=Base)
login_manager = LoginManager()

user_roles = Table(
    "user_roles",
    Base.metadata,
    Column("user_id", ForeignKey("users.id"), primary_key=True),
    Column("role_id", ForeignKey("roles.id"), primary_key=True),
)


class Role(db.Model):
    __tablename__ = "roles"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(unique=True)


class User(UserMixin, db.Model):
    __tablename__ = "users"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(unique=True)
    password_hash: Mapped[str]
    roles: Mapped[list[Role]] = relationship(secondary=user_roles)

    def has_role(self, *names: str) -> bool:
        return set(names) <= {role.name for role in self.roles}


class Book(db.Model):
    __tablename__ = "books"
    id: Mapped[int] = mapped_column(primary_key=True)
    title: Mapped[str]
    owner_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
    status: Mapped[str] = mapped_column(default="draft")


PASSWORD = secrets.token_urlsafe(12)            # demo password of every seeded user, different on each run
PASSWORD_HASH = generate_password_hash(PASSWORD)


@login_manager.user_loader
def load_user(user_id: str):
    return db.session.get(User, int(user_id))


@login_manager.unauthorized_handler
def unauthorized():
    abort(401)                                  # an API answers 401 instead of redirecting to a login page


def create_app(*register):
    app = Flask(__name__)
    app.config.update(SECRET_KEY=secrets.token_hex(), SQLALCHEMY_DATABASE_URI="sqlite://")
    db.init_app(app)
    login_manager.init_app(app)

    @app.post("/login")
    def login():
        data = request.get_json()
        user = db.session.scalar(select(User).where(User.name == data["name"]))
        if user is None or not check_password_hash(user.password_hash, data["password"]):
            abort(401)
        login_user(user)
        return jsonify(ok=True)

    with app.app_context():
        db.create_all()
        roles = {name: Role(name=name) for name in ("reader", "editor", "admin")}
        db.session.add_all(roles.values())
        users = {
            "alice": User(name="alice", password_hash=PASSWORD_HASH, roles=[roles["admin"]]),
            "bob": User(name="bob", password_hash=PASSWORD_HASH, roles=[roles["editor"]]),
            "carol": User(name="carol", password_hash=PASSWORD_HASH, roles=[roles["reader"]]),
        }
        db.session.add_all(users.values())
        db.session.flush()
        db.session.add_all([
            Book(title="Dune", owner_id=users["bob"].id),
            Book(title="Emma", owner_id=users["carol"].id),
        ])
        db.session.commit()
    for add_routes in register:
        add_routes(app)
    return app


def client_for(app, name=None):
    client = app.test_client()
    if name:
        assert client.post("/login", json={"name": name, "password": PASSWORD}).status_code == 200
    return client

def _guard(allowed):
    def decorator(view):
        @wraps(view)
        def wrapped(*args, **kwargs):
            if not current_user.is_authenticated:
                return login_manager.unauthorized()
            if not allowed(current_user):
                abort(403)
            return view(*args, **kwargs)
        return wrapped
    return decorator


def roles_required(*names):
    """The user must have all the roles."""
    return _guard(lambda user: user.has_role(*names))


def roles_accepted(*names):
    """The user must have at least one of the roles."""
    return _guard(lambda user: any(user.has_role(name) for name in names))


def role_routes(app):
    @app.post("/books")
    @roles_accepted("admin", "editor")
    def create_book():
        book = Book(title=request.get_json()["title"], owner_id=current_user.id)
        db.session.add(book)
        db.session.commit()
        return jsonify(id=book.id), 201

    @app.delete("/books/<int:book_id>")
    @roles_required("admin")
    def delete_book(book_id):
        return jsonify(deleted=book_id)


app = create_app(role_routes)
for who in (None, "carol", "bob", "alice"):
    client = client_for(app, who)
    print(who, client.post("/books", json={"title": "Persuasion"}).status_code, client.delete("/books/1").status_code)

ROLE_PERMISSIONS = {
    "reader": {"book:read"},
    "editor": {"book:read", "book:create", "book:edit"},
    "admin": {"book:read", "book:create", "book:edit", "book:delete", "user:manage"},
}


def permissions_of(user) -> frozenset[str]:
    return frozenset().union(*(ROLE_PERMISSIONS.get(role.name, set()) for role in user.roles))


def permission_required(permission: str):
    return _guard(lambda user: permission in permissions_of(user))


def permission_routes(app):
    @app.get("/books")
    @permission_required("book:read")
    def list_books():
        return jsonify(titles=[book.title for book in db.session.scalars(select(Book))])

    @app.post("/books")
    @permission_required("book:create")
    def create_book():
        return jsonify(created=True), 201

    @app.delete("/books/<int:book_id>")
    @permission_required("book:delete")
    def delete_book(book_id):
        return jsonify(deleted=book_id)


app = create_app(permission_routes)
carol = client_for(app, "carol")
print(carol.get("/books").get_json(), carol.post("/books", json={}).status_code)

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone


@dataclass(frozen=True)
class Subject:
    user_id: int
    roles: frozenset[str]


@dataclass(frozen=True)
class BookResource:
    owner_id: int
    status: str                       # "draft" or "published"
    locked_until: datetime | None = None


def can_edit(subject: Subject, book: BookResource, now: datetime) -> bool:
    if "admin" in subject.roles:
        return True
    if book.locked_until is not None and now < book.locked_until:
        return False                  # environment attribute: an editorial lock
    if book.status == "published":
        return False                  # resource attribute: published books are frozen
    return "editor" in subject.roles and book.owner_id == subject.user_id


now = datetime.now(timezone.utc)
bob = Subject(2, frozenset({"editor"}))
draft = BookResource(owner_id=2, status="draft")
print(can_edit(bob, draft, now))
print(can_edit(bob, BookResource(2, "published"), now))
print(can_edit(bob, BookResource(2, "draft", locked_until=now + timedelta(hours=1)), now))
print(can_edit(Subject(3, frozenset({"editor"})), draft, now))

from collections.abc import Callable


@dataclass(frozen=True)
class Policy:
    name: str
    effect: str                                           # "allow" or "deny"
    actions: frozenset[str]
    applies: Callable[[Subject, BookResource, datetime], bool]


POLICIES = [
    Policy("admins-manage-books", "allow", frozenset({"book:edit", "book:delete"}),
           lambda s, b, n: "admin" in s.roles),
    Policy("owners-edit-drafts", "allow", frozenset({"book:edit"}),
           lambda s, b, n: b.owner_id == s.user_id and b.status == "draft"),
    Policy("no-changes-while-locked", "deny", frozenset({"book:edit", "book:delete"}),
           lambda s, b, n: b.locked_until is not None and n < b.locked_until and "admin" not in s.roles),
]


def decide(policies, action, subject, book, now) -> bool:
    matching = [p for p in policies if action in p.actions and p.applies(subject, book, now)]
    if any(p.effect == "deny" for p in matching):
        return False                                      # deny overrides allow
    return any(p.effect == "allow" for p in matching)     # default deny


locked = BookResource(2, "draft", locked_until=now + timedelta(hours=1))
print(decide(POLICIES, "book:edit", bob, draft, now))
print(decide(POLICIES, "book:edit", bob, locked, now))
print(decide(POLICIES, "book:delete", bob, draft, now))
print(decide(POLICIES, "book:delete", Subject(1, frozenset({"admin"})), locked, now))

def visible_books(statement):
    """Restrict a select(Book) to what the current user may see: everything for admins, own books otherwise."""
    if current_user.has_role("admin"):
        return statement
    return statement.where(Book.owner_id == current_user.id)


def ownership_routes(app):
    @app.get("/books/<int:book_id>")
    @login_required
    def get_book(book_id):
        book = db.one_or_404(visible_books(select(Book).where(Book.id == book_id)))
        return jsonify(id=book.id, title=book.title)

    @app.patch("/books/<int:book_id>")
    @login_required
    def update_book(book_id):
        book = db.one_or_404(visible_books(select(Book).where(Book.id == book_id)))
        data = request.get_json()
        if "title" in data:                           # explicit allow-list: owner_id and status are not client fields
            book.title = data["title"]
        db.session.commit()
        return jsonify(id=book.id, title=book.title, owner_id=book.owner_id)

    @app.get("/my/books")
    @login_required
    def my_books():
        statement = visible_books(select(Book).order_by(Book.id))
        return jsonify(ids=[book.id for book in db.session.scalars(statement)])


app = create_app(ownership_routes)
bob = client_for(app, "bob")
print(bob.get("/books/1").status_code, bob.get("/books/2").status_code)
print(bob.patch("/books/1", json={"title": "Dune Messiah", "owner_id": 3}).get_json())
print(client_for(app, "alice").get("/my/books").get_json())

admin_bp = Blueprint("admin", __name__, url_prefix="/admin")


@admin_bp.before_request
def admin_only():
    if not current_user.is_authenticated:
        return login_manager.unauthorized()
    if not current_user.has_role("admin"):
        abort(403)


@admin_bp.get("/users")
def list_users():
    users = db.session.scalars(select(User).order_by(User.id))
    return jsonify([{"name": user.name, "roles": sorted(role.name for role in user.roles)} for user in users])


@admin_bp.post("/users/<int:user_id>/roles")
def grant_role(user_id):
    name = request.get_json().get("role")
    if name not in ROLE_PERMISSIONS:                      # only known roles can be granted
        abort(422)
    user = db.get_or_404(User, user_id)
    role = db.session.scalar(select(Role).where(Role.name == name))
    if role not in user.roles:
        user.roles.append(role)
        db.session.commit()
    return jsonify(name=user.name, roles=sorted(role.name for role in user.roles))


def admin_routes(app):
    app.register_blueprint(admin_bp)

    @app.get("/")
    def home():
        return jsonify(public=True)


app = create_app(admin_routes, permission_routes)
anonymous, bob, alice = client_for(app), client_for(app, "bob"), client_for(app, "alice")
print(anonymous.get("/").status_code, anonymous.get("/admin/users").status_code, bob.get("/admin/users").status_code)
print(alice.post("/admin/users/2/roles", json={"role": "admin"}).get_json())
print(bob.get("/admin/users").status_code, bob.delete("/books/1").status_code)

from flask_principal import (Identity, ItemNeed, Permission, Principal, RoleNeed, UserNeed, identity_loaded)


def principal_routes(app):
    principal = Principal(app, use_sessions=False)         # identity comes from Flask-Login, not from a session key
    admin_permission = Permission(RoleNeed("admin"))

    @principal.identity_loader
    def identity_from_flask_login():
        return Identity(current_user.id) if current_user.is_authenticated else None

    @identity_loaded.connect_via(app)
    def add_needs(sender, identity):
        identity.provides.add(UserNeed(identity.id))
        for role in current_user.roles:
            identity.provides.add(RoleNeed(role.name))
        for book_id in db.session.scalars(select(Book.id).where(Book.owner_id == identity.id)):
            identity.provides.add(ItemNeed("edit", book_id, "book"))

    @app.delete("/books/<int:book_id>")
    @admin_permission.require(http_exception=403)
    def delete_book(book_id):
        return jsonify(deleted=book_id)

    @app.patch("/books/<int:book_id>")
    def edit_book(book_id):
        Permission(ItemNeed("edit", book_id, "book"), RoleNeed("admin")).test(http_exception=403)
        return jsonify(edited=book_id)


app = create_app(principal_routes)
bob, alice, anonymous = client_for(app, "bob"), client_for(app, "alice"), client_for(app)
print(bob.patch("/books/1").status_code, bob.patch("/books/2").status_code, alice.patch("/books/2").status_code)
print(bob.delete("/books/1").status_code, alice.delete("/books/1").status_code, anonymous.delete("/books/1").status_code)

import flask_security as fsec
from flask_security.models import fsqla_v3 as fsqla


def create_security_app(**extra_config):
    app = Flask(__name__)
    app.config.update(
        SECRET_KEY=secrets.token_hex(),
        SECURITY_PASSWORD_SALT=secrets.token_hex(),
        SQLALCHEMY_DATABASE_URI="sqlite://",
        WTF_CSRF_ENABLED=False,                           # only for the test client; keep CSRF on in real applications
        **extra_config,
    )
    fs_db = SQLAlchemy()
    fsqla.FsModels.set_db_info(fs_db)

    class Role(fs_db.Model, fsqla.FsRoleMixin):           # defined here only so this page's other examples keep
        pass                                              # their own User and Role; at module level in a real project

    class User(fs_db.Model, fsqla.FsUserMixin):
        pass

    class WebAuthn(fs_db.Model, fsqla.FsWebAuthnMixin):   # the v3 mixins expect this model to exist
        pass

    fs_db.init_app(app)
    store = fsec.SQLAlchemyUserDatastore(fs_db, User, Role)
    fsec.Security(app, store)

    @app.get("/shelf")
    @fsec.auth_required()
    def shelf():
        return jsonify(user=fsec.current_user.email)

    @app.delete("/books/<int:book_id>")
    @fsec.roles_required("admin")
    def delete_book(book_id):
        return jsonify(deleted=book_id)

    @app.post("/books")
    @fsec.permissions_required("books-write")
    def create_book():
        return jsonify(created=True), 201

    with app.app_context():
        fs_db.create_all()
        store.find_or_create_role(name="admin", permissions={"books-write", "books-delete"})
        store.find_or_create_role(name="editor", permissions={"books-write"})
        store.create_user(email="alice@example.com", password=fsec.hash_password(PASSWORD), roles=["admin"])
        store.create_user(email="bob@example.com", password=fsec.hash_password(PASSWORD), roles=["editor"])
        fs_db.session.commit()
    return app


sec_app = create_security_app()


def sec_client(email):
    client = sec_app.test_client()
    assert client.post("/login", json={"email": email, "password": PASSWORD}).status_code == 200
    return client


alice, bob = sec_client("alice@example.com"), sec_client("bob@example.com")
anonymous = sec_app.test_client()
print(anonymous.get("/shelf").status_code, anonymous.get("/shelf", headers={"Accept": "application/json"}).status_code)
print(bob.delete("/books/1").status_code, alice.delete("/books/1").status_code, bob.post("/books").status_code)

from datetime import timedelta

from flask_jwt_extended import JWTManager, create_access_token, get_jwt, verify_jwt_in_request


def scope_required(scope: str):
    def decorator(view):
        @wraps(view)
        def wrapped(*args, **kwargs):
            verify_jwt_in_request()                           # 401 for a missing, expired or forged token
            if scope not in get_jwt().get("scope", "").split():
                abort(403)
            return view(*args, **kwargs)
        return wrapped
    return decorator


def jwt_routes(app):
    app.config["JWT_SECRET_KEY"] = secrets.token_hex()        # in production: from the environment
    app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(minutes=10)
    JWTManager(app)

    @app.get("/api/books")
    @scope_required("book:read")
    def api_books():
        return jsonify(roles=get_jwt()["roles"])

    @app.delete("/api/books/<int:book_id>")
    @scope_required("book:delete")
    def api_delete(book_id):
        return jsonify(deleted=book_id)


app = create_app(jwt_routes)
with app.app_context():
    reader_token = create_access_token(identity="3", additional_claims={"roles": ["reader"], "scope": "book:read"})
client = app.test_client()
headers = {"Authorization": f"Bearer {reader_token}"}
print(client.get("/api/books").status_code, client.get("/api/books", headers=headers).get_json())
print(client.delete("/api/books/1", headers=headers).status_code)

from flask import render_template_string


def template_routes(app):
    @app.context_processor
    def inject_can():
        return {"can": lambda permission: current_user.is_authenticated and permission in permissions_of(current_user)}

    @app.get("/shelf")
    @login_required
    def shelf():
        return render_template_string(
            "<p>{{ current_user.name }}</p>{% if can('book:delete') %}<button>Delete</button>{% endif %}"
        )


app = create_app(permission_routes, template_routes)
for who in ("alice", "bob"):
    client = client_for(app, who)
    print(who, "<button>" in client.get("/shelf").get_data(as_text=True), client.delete("/books/1").status_code)

app = create_app(permission_routes, ownership_routes, admin_routes, template_routes)

MATRIX = [
    # who,     method,   url,              expected status
    (None,    "GET",    "/books",          401),
    ("carol", "GET",    "/books",          200),
    ("carol", "POST",   "/books",          403),
    ("bob",   "POST",   "/books",          201),
    ("bob",   "GET",    "/books/1",        200),
    ("bob",   "GET",    "/books/2",        404),
    ("alice", "GET",    "/books/2",        200),
    ("bob",   "GET",    "/admin/users",    403),
    ("alice", "GET",    "/admin/users",    200),
    ("bob",   "DELETE", "/books/1",        403),
    ("alice", "DELETE", "/books/1",        200),
]

for who, method, url, expected in MATRIX:
    client = client_for(app, who)
    got = client.open(url, method=method, json={} if method == "POST" else None).status_code
    assert got == expected, (who, method, url, got, expected)
print(f"{len(MATRIX)} authorization cases passed")

public = {"static", "login", "home"}


def is_protected(endpoint: str) -> bool:
    blueprint = endpoint.rpartition(".")[0]
    has_decorator = hasattr(app.view_functions[endpoint], "__wrapped__")
    has_blueprint_hook = bool(blueprint) and bool(app.before_request_funcs.get(blueprint))
    return has_decorator or has_blueprint_hook


unprotected = [rule.endpoint for rule in app.url_map.iter_rules()
               if rule.endpoint not in public and not is_protected(rule.endpoint)]
print("unprotected endpoints:", unprotected)