diff --git a/src/mcp/oauth/store.py b/src/mcp/oauth/store.py index 829be0f..2ec91dd 100644 --- a/src/mcp/oauth/store.py +++ b/src/mcp/oauth/store.py @@ -135,3 +135,73 @@ def delete_code(db_path, code: str) -> None: with _connect(db_path) as conn: conn.execute("DELETE FROM oauth_codes WHERE code_hash = ?", (_hash(code),)) conn.commit() + + +def save_token( + db_path, + *, + access_token: str, + refresh_token: str | None, + client_id: str, + user_id: int, + scope: str | None, + resource: str, + access_expires_at: str, + refresh_expires_at: str | None, +) -> int: + with _connect(db_path) as conn: + cur = conn.execute( + "INSERT INTO oauth_tokens (access_token_hash, refresh_token_hash, " + "client_id, user_id, scope, resource, issued_at, access_expires_at, " + "refresh_expires_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + _hash(access_token), + _hash(refresh_token) if refresh_token else None, + client_id, + user_id, + scope, + resource, + _utcnow_iso(), + access_expires_at, + refresh_expires_at, + ), + ) + conn.commit() + return cur.lastrowid + + +def get_token_by_access(db_path, access_token: str) -> dict | None: + with _connect(db_path) as conn: + row = conn.execute( + "SELECT * FROM oauth_tokens WHERE access_token_hash = ?", + (_hash(access_token),), + ).fetchone() + return dict(row) if row else None + + +def get_token_by_refresh(db_path, refresh_token: str) -> dict | None: + with _connect(db_path) as conn: + row = conn.execute( + "SELECT * FROM oauth_tokens WHERE refresh_token_hash = ?", + (_hash(refresh_token),), + ).fetchone() + return dict(row) if row else None + + +def revoke_token(db_path, token_id: int) -> None: + with _connect(db_path) as conn: + conn.execute( + "UPDATE oauth_tokens SET revoked_at = ? WHERE id = ?", + (_utcnow_iso(), token_id), + ) + conn.commit() + + +def increment_usage(db_path, token_id: int) -> None: + with _connect(db_path) as conn: + conn.execute( + "UPDATE oauth_tokens SET count_total = count_total + 1, last_used_at = ? " + "WHERE id = ?", + (_utcnow_iso(), token_id), + ) + conn.commit() diff --git a/tests/mcp/test_oauth_store.py b/tests/mcp/test_oauth_store.py index 28b482c..f089532 100644 --- a/tests/mcp/test_oauth_store.py +++ b/tests/mcp/test_oauth_store.py @@ -36,3 +36,25 @@ def test_save_get_delete_code(tmp_path): assert row["resource"] == "https://colibre.fr/_mcp" store.delete_code(db, "thecode") assert store.get_code(db, "thecode") is None + + +def test_save_get_revoke_token(tmp_path): + db = tmp_path / "u.sqlite" + store.init_schema(db) + tid = store.save_token( + db, + access_token="acc", + refresh_token="ref", + client_id="abc", + user_id=7, + scope="mcp", + resource="https://colibre.fr/_mcp", + access_expires_at="2999-01-01T00:00:00+00:00", + refresh_expires_at="2999-01-01T00:00:00+00:00", + ) + assert store.get_token_by_access(db, "acc")["id"] == tid + assert store.get_token_by_refresh(db, "ref")["user_id"] == 7 + store.increment_usage(db, tid) + assert store.get_token_by_access(db, "acc")["count_total"] == 1 + store.revoke_token(db, tid) + assert store.get_token_by_access(db, "acc")["revoked_at"] is not None