feat(mcp): store OAuth — tokens access/refresh + usage (#114)
Implements OAuth token management functions for Task 3: - save_token: persists access/refresh tokens with hashing - get_token_by_access/refresh: retrieves tokens by their hashes - increment_usage: tracks token usage count and last_used_at - revoke_token: sets revoked_at timestamp All 3 tests passing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -135,3 +135,73 @@ def delete_code(db_path, code: str) -> None:
|
|||||||
with _connect(db_path) as conn:
|
with _connect(db_path) as conn:
|
||||||
conn.execute("DELETE FROM oauth_codes WHERE code_hash = ?", (_hash(code),))
|
conn.execute("DELETE FROM oauth_codes WHERE code_hash = ?", (_hash(code),))
|
||||||
conn.commit()
|
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()
|
||||||
|
|||||||
@@ -36,3 +36,25 @@ def test_save_get_delete_code(tmp_path):
|
|||||||
assert row["resource"] == "https://colibre.fr/_mcp"
|
assert row["resource"] == "https://colibre.fr/_mcp"
|
||||||
store.delete_code(db, "thecode")
|
store.delete_code(db, "thecode")
|
||||||
assert store.get_code(db, "thecode") is None
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user