+8 -20

src/atproto/oauth.py

··· 1 1 - import sqlite3 2 2 - from typing import Any, NamedTuple 1 1 + from typing import Any, Callable, NamedTuple 3 2 import time 4 3 import json 5 4 from authlib.jose import JsonWebKey, Key ··· 20 19 refresh_token: str 21 20 scope: str 22 21 sub: str 22 22 + # only for parsing 23 23 + token_type: str | None 24 24 + expires_in: int | None 23 25 24 26 25 27 # Prepares and sends a pushed auth request (PAR) via HTTP POST to the Authorization Server. ··· 160 162 161 163 resp.raise_for_status() 162 164 token_body = resp.json() 163 163 - try: 164 164 - tokens = OAuthTokens(**token_body) 165 165 - except TypeError: 166 166 - raise Exception("invalid token body") 165 165 + tokens = OAuthTokens(**token_body) 167 166 168 167 return tokens, dpop_authserver_nonce 169 168 ··· 226 225 227 226 resp.raise_for_status() 228 227 token_body = resp.json() 229 229 - try: 230 230 - tokens = OAuthTokens(**token_body) 231 231 - except TypeError: 232 232 - raise Exception("invalid token body") 228 228 + tokens = OAuthTokens(**token_body) 233 229 234 230 return tokens, dpop_authserver_nonce 235 231 ··· 240 236 method: str, 241 237 url: str, 242 238 user: OAuthSession, 243 243 - db: sqlite3.Connection, 239 239 + update_dpop_pds_nonce: Callable[[str], None], 244 240 body: dict[str, Any] | None = None, 245 241 ) -> Response | None: 246 242 dpop_private_jwk = JsonWebKey.import_key(json.loads(user.dpop_private_jwk)) ··· 275 271 response.status_code in [400, 401] 276 272 and response.json()["error"] == "use_dpop_nonce" 277 273 ): 278 278 - # print(resp.headers) 279 274 dpop_pds_nonce = response.headers["DPoP-Nonce"] 280 275 print(f"retrying with new PDS DPoP nonce: {dpop_pds_nonce}") 281 281 - # update session database with new nonce 282 282 - cur = db.cursor() 283 283 - _ = cur.execute( 284 284 - "UPDATE oauth_session SET dpop_pds_nonce = ? WHERE did = ?;", 285 285 - [dpop_pds_nonce, user.did], 286 286 - ) 287 287 - db.commit() 288 288 - cur.close() 276 276 + update_dpop_pds_nonce(dpop_pds_nonce) 289 277 continue 290 278 break 291 279

+20 -29

src/main.py

··· 4 4 5 5 from .atproto import PdsUrl, get_record, resolve_did_from_handle, resolve_pds_from_did 6 6 from .atproto.oauth import pds_authed_req 7 7 - from .db import close_db_connection, get_db, init_db 8 8 - from .oauth import oauth 7 7 + from .db import close_db_connection, init_db 8 8 + from .oauth import get_auth_session, oauth, save_auth_session 9 9 from .types import OAuthSession 10 10 11 11 app = Flask(__name__) ··· 21 21 22 22 @app.before_request 23 23 def load_user_to_context(): 24 24 - user: OAuthSession | None = None 25 25 - did: str | None = session.get("user_did") 26 26 - if did is not None: 27 27 - db = get_db(app) 28 28 - row = db.execute( 29 29 - "select * from oauth_session where did = ?", 30 30 - (did,), 31 31 - ).fetchone() 32 32 - user = OAuthSession(**row) 33 33 - g.user = user 24 24 + g.user = get_auth_session(session) 34 25 35 26 36 27 def get_user() -> OAuthSession | None: ··· 84 75 if not username: 85 76 return redirect(url_for("page_login"), 303) 86 77 return redirect(url_for("oauth.oauth_start", username=username), 303) 78 78 + 79 79 + 80 80 + @app.route("/auth/logout") 81 81 + def auth_logout(): 82 82 + session.clear() 83 83 + return redirect("/", 303) 87 84 88 85 89 86 @app.get("/editor") ··· 174 171 return redirect("/editor", 303) 175 172 176 173 174 174 + @app.get("/terms") 175 175 + def page_terms(): 176 176 + return "come back soon" 177 177 + 178 178 + 177 179 def load_links(pds: str, did: str, reload: bool = False) -> list[dict[str, str]] | None: 178 180 if did in links and not reload: 179 181 app.logger.debug(f"returning cached links for {did}") ··· 228 230 "rkey": rkey, 229 231 "record": record, 230 232 } 233 233 + 234 234 + def update_dpop_pds_nonce(nonce: str): 235 235 + session_ = user._replace(dpop_pds_nonce=nonce) 236 236 + save_auth_session(session, session_) 237 237 + 231 238 response = pds_authed_req( 232 239 method="POST", 233 240 url=endpoint, 234 241 body=body, 235 242 user=user, 236 236 - db=get_db(app), 243 243 + update_dpop_pds_nonce=update_dpop_pds_nonce, 237 244 ) 238 245 if not response or not response.ok: 239 246 app.logger.warning("PDS HTTP ERROR") 240 240 - 241 241 - 242 242 - # AUTH 243 243 - 244 244 - 245 245 - @app.route("/auth/logout") 246 246 - def auth_logout(): 247 247 - user = get_user() 248 248 - if user is not None: 249 249 - db = get_db(app) 250 250 - cursor = db.cursor() 251 251 - _ = cursor.execute("delete from oauth_session where did = ?", (user.did,)) 252 252 - db.commit() 253 253 - cursor.close() 254 254 - session.clear() 255 255 - return redirect("/", 303)

+75 -52

src/oauth.py

··· 1 1 + from typing import NamedTuple 1 2 from authlib.jose import JsonWebKey, Key 2 3 from flask import Blueprint, current_app, jsonify, redirect, request, session, url_for 4 4 + from flask.sessions import SessionMixin 3 5 from urllib.parse import urlencode 4 6 5 7 import json ··· 14 16 ) 15 17 from .atproto.oauth import initial_token_request, send_par_auth_request 16 18 from .security import is_safe_url 17 17 - from .types import OAuthAuthRequest 18 18 - from .db import get_db 19 19 + from .types import OAuthAuthRequest, OAuthSession 19 20 20 21 oauth = Blueprint("oauth", __name__, url_prefix="/oauth") 21 22 ··· 87 88 par_request_uri: str = resp.json()["request_uri"] 88 89 current_app.logger.debug(f"saving oauth_auth_request to DB state={state}") 89 90 90 90 - db = get_db(current_app) 91 91 - cursor = db.cursor() 92 92 - _ = cursor.execute( 93 93 - "insert or replace into oauth_auth_requests values (?, ?, ?, ?, ?, ?, ?, ?, ?)", 94 94 - ( 95 95 - state, 96 96 - authserver_meta["issuer"], 97 97 - did, 98 98 - handle, 99 99 - pds_url, 100 100 - pkce_verifier, 101 101 - scope, 102 102 - dpop_authserver_nonce, 103 103 - dpop_private_jwk.as_json(is_private=True), 104 104 - ), 91 91 + oauth_request = OAuthAuthRequest( 92 92 + state, 93 93 + authserver_meta["issuer"], 94 94 + did, 95 95 + handle, 96 96 + pds_url, 97 97 + pkce_verifier, 98 98 + scope, 99 99 + dpop_authserver_nonce, 100 100 + dpop_private_jwk.as_json(is_private=True), 105 101 ) 106 106 - db.commit() 107 107 - cursor.close() 102 102 + save_auth_request(session, oauth_request) 108 103 109 104 auth_endpoint = authserver_meta["authorization_endpoint"] 110 105 assert is_safe_url(auth_endpoint) ··· 118 113 authserver_iss = request.args["iss"] 119 114 authorization_code = request.args["code"] 120 115 121 121 - db = get_db(current_app) 122 122 - cursor = db.cursor() 123 123 - 124 124 - row = cursor.execute( 125 125 - "select * from oauth_auth_requests where state = ?", (state,) 126 126 - ).fetchone() 127 127 - try: 128 128 - auth_request = OAuthAuthRequest(**row) 129 129 - except TypeError: 116 116 + auth_request = get_auth_request(session) 117 117 + if not auth_request: 130 118 return redirect(url_for("page_login"), 303) 131 119 132 120 current_app.logger.debug(f"Deleting auth request for state={state}") 133 133 - _ = cursor.execute("delete from oauth_auth_requests where state = ?", (state,)) 134 134 - db.commit() 121 121 + delete_auth_request(session) 135 122 136 123 assert auth_request.authserver_iss == authserver_iss 137 124 assert auth_request.state == state ··· 147 134 148 135 row = auth_request 149 136 150 150 - did = auth_request.did 151 137 if row.did: 152 138 # If we started with an account identifier, this is simple 153 139 did, handle, pds_url = row.did, row.handle, row.pds_url ··· 166 152 assert authserver_url == authserver_iss 167 153 168 154 assert row.scope == tokens.scope 155 155 + assert pds_url is not None 169 156 170 170 - current_app.logger.debug("storing user did and handle") 171 171 - db = get_db(current_app) 172 172 - cursor = db.cursor() 173 173 - _ = cursor.execute( 174 174 - "insert or replace into oauth_session values (?, ?, ?, ?, ?, ?, ?, ?, ?)", 175 175 - ( 176 176 - did, 177 177 - handle, 178 178 - pds_url, 179 179 - authserver_iss, 180 180 - tokens.access_token, 181 181 - tokens.refresh_token, 182 182 - dpop_authserver_nonce, 183 183 - None, 184 184 - auth_request.dpop_private_jwk, 185 185 - ), 157 157 + current_app.logger.debug("storing user oauth session") 158 158 + oauth_session = OAuthSession( 159 159 + did, 160 160 + handle, 161 161 + pds_url, 162 162 + authserver_iss, 163 163 + tokens.access_token, 164 164 + tokens.refresh_token, 165 165 + dpop_authserver_nonce, 166 166 + None, 167 167 + auth_request.dpop_private_jwk, 186 168 ) 187 187 - db.commit() 188 188 - cursor.close() 189 189 - 190 190 - session["user_did"] = did 191 191 - session["user_handle"] = auth_request.handle 169 169 + save_auth_session(session, oauth_session) 192 170 193 171 return redirect(url_for("page_login")) 194 172 ··· 221 199 CLIENT_SECRET_JWK = JsonWebKey.import_key(current_app.config["CLIENT_SECRET_JWK"]) 222 200 CLIENT_PUB_JWK = json.loads(CLIENT_SECRET_JWK.as_json(is_private=False)) 223 201 return jsonify({"keys": [CLIENT_PUB_JWK]}) 202 202 + 203 203 + 204 204 + # Session storage 205 205 + 206 206 + 207 207 + def save_auth_request(session: SessionMixin, request: OAuthAuthRequest): 208 208 + return _set_into_session(session, "oauth_auth_request", request) 209 209 + 210 210 + 211 211 + def save_auth_session(session: SessionMixin, auth_session: OAuthSession): 212 212 + return _set_into_session(session, "oauth_auth_session", auth_session) 213 213 + 214 214 + 215 215 + def delete_auth_request(session: SessionMixin): 216 216 + return _delete_from_session(session, "oauth_auth_request") 217 217 + 218 218 + 219 219 + def delete_auth_session(session: SessionMixin): 220 220 + return _delete_from_session(session, "oauth_auth_session") 221 221 + 222 222 + 223 223 + def get_auth_request(session: SessionMixin) -> OAuthAuthRequest | None: 224 224 + try: 225 225 + return OAuthAuthRequest(**session["oauth_auth_request"]) 226 226 + except TypeError as exception: 227 227 + current_app.logger.debug("unable to load oauth_auth_request") 228 228 + current_app.logger.debug(exception) 229 229 + return None 230 230 + 231 231 + 232 232 + def get_auth_session(session: SessionMixin) -> OAuthSession | None: 233 233 + try: 234 234 + return OAuthSession(**session["oauth_auth_session"]) 235 235 + except TypeError as exception: 236 236 + current_app.logger.debug("unable to load oauth_auth_session") 237 237 + current_app.logger.debug(exception) 238 238 + return None 239 239 + 240 240 + 241 241 + def _set_into_session(session: SessionMixin, key: str, value: NamedTuple): 242 242 + session[key] = value._asdict() 243 243 + 244 244 + 245 245 + def _delete_from_session(session: SessionMixin, key: str): 246 246 + del session[key]

+1 -23

src/schema.sql

··· 1 1 - create table if not exists oauth_auth_requests ( 2 2 - state text not null primary key, 3 3 - authserver_iss text not null, 4 4 - did text, 5 5 - handle text, 6 6 - pds_url text, 7 7 - pkce_verifier text not null, 8 8 - scope text not null, 9 9 - dpop_authserver_nonce text not null, 10 10 - dpop_private_jwk text not null 11 11 - ) strict, without rowid; 12 12 - 13 13 - create table if not exists oauth_sessions ( 14 14 - did text not null primary key, 15 15 - handle text, 16 16 - pds_url text not null, 17 17 - authserver_iss text not null, 18 18 - access_token text, 19 19 - refresh_token text, 20 20 - dpop_authserver_nonce text not null, 21 21 - dpop_pds_nonce text, 22 22 - dpop_private_jwk text not null 23 23 - ) strict, without rowid; 1 1 + -- empty for now