diff --git a/src/api/tokens_db.py b/src/api/tokens_db.py index c1916c5..068e8aa 100644 --- a/src/api/tokens_db.py +++ b/src/api/tokens_db.py @@ -16,7 +16,8 @@ CREATE TABLE IF NOT EXISTS api_tokens ( created_at TEXT NOT NULL, last_used_at TEXT, count_total INTEGER NOT NULL DEFAULT 0, - revoked_at TEXT + revoked_at TEXT, + kind TEXT NOT NULL DEFAULT 'api' ); CREATE INDEX IF NOT EXISTS idx_api_tokens_hash ON api_tokens(token_hash); """ @@ -47,13 +48,15 @@ def init_schema(db_path) -> None: conn.commit() -def create_token(db_path, label: str, user_id: int | None = None) -> tuple[str, int]: +def create_token( + db_path, label: str, user_id: int | None = None, kind: str = "api" +) -> tuple[str, int]: token = TOKEN_PREFIX + secrets.token_hex(32) with _connect(db_path) as conn: cur = conn.execute( - "INSERT INTO api_tokens (token_hash, label, user_id, created_at) " - "VALUES (?, ?, ?, ?)", - (_hash(token), label, user_id, _utcnow_iso()), + "INSERT INTO api_tokens (token_hash, label, user_id, kind, created_at) " + "VALUES (?, ?, ?, ?, ?)", + (_hash(token), label, user_id, kind, _utcnow_iso()), ) conn.commit() return token, cur.lastrowid @@ -92,3 +95,24 @@ def list_tokens(db_path) -> list[dict]: with _connect(db_path) as conn: rows = conn.execute("SELECT * FROM api_tokens ORDER BY id").fetchall() return [dict(r) for r in rows] + + +def list_user_tokens(db_path, user_id: int, kind: str = "mcp") -> list[dict]: + with _connect(db_path) as conn: + rows = conn.execute( + "SELECT * FROM api_tokens WHERE user_id = ? AND kind = ? " + "ORDER BY created_at DESC, id DESC", + (user_id, kind), + ).fetchall() + return [dict(r) for r in rows] + + +def revoke_user_token(db_path, token_id: int, user_id: int) -> bool: + with _connect(db_path) as conn: + cur = conn.execute( + "UPDATE api_tokens SET revoked_at = ? " + "WHERE id = ? AND user_id = ? AND revoked_at IS NULL", + (_utcnow_iso(), token_id, user_id), + ) + conn.commit() + return cur.rowcount > 0 diff --git a/tests/api/test_tokens_db.py b/tests/api/test_tokens_db.py index 92a05ef..64b34bd 100644 --- a/tests/api/test_tokens_db.py +++ b/tests/api/test_tokens_db.py @@ -62,3 +62,42 @@ def test_list_tokens_returns_all(temp_db): tokens_db.create_token(temp_db, "b") rows = tokens_db.list_tokens(temp_db) assert [r["label"] for r in rows] == ["a", "b"] + + +def test_create_token_defaults_to_api_kind(temp_db): + token, token_id = tokens_db.create_token(temp_db, "x") + row = tokens_db.get_token_by_plaintext(temp_db, token) + assert row["kind"] == "api" + + +def test_create_token_with_mcp_kind(temp_db): + token, _ = tokens_db.create_token(temp_db, "x", user_id=7, kind="mcp") + row = tokens_db.get_token_by_plaintext(temp_db, token) + assert row["kind"] == "mcp" + assert row["user_id"] == 7 + + +def test_list_user_tokens_filters_by_user_and_kind(temp_db): + tokens_db.create_token(temp_db, "mcp-u1", user_id=1, kind="mcp") + tokens_db.create_token(temp_db, "api-u1", user_id=1, kind="api") + tokens_db.create_token(temp_db, "mcp-u2", user_id=2, kind="mcp") + rows = tokens_db.list_user_tokens(temp_db, 1, "mcp") + assert [r["label"] for r in rows] == ["mcp-u1"] + + +def test_revoke_user_token_revokes_own(temp_db): + token, token_id = tokens_db.create_token(temp_db, "x", user_id=1, kind="mcp") + assert tokens_db.revoke_user_token(temp_db, token_id, 1) is True + assert tokens_db.get_token_by_plaintext(temp_db, token)["revoked_at"] is not None + + +def test_revoke_user_token_refuses_other_owner(temp_db): + token, token_id = tokens_db.create_token(temp_db, "x", user_id=1, kind="mcp") + assert tokens_db.revoke_user_token(temp_db, token_id, 999) is False + assert tokens_db.get_token_by_plaintext(temp_db, token)["revoked_at"] is None + + +def test_revoke_user_token_already_revoked_returns_false(temp_db): + _, token_id = tokens_db.create_token(temp_db, "x", user_id=1, kind="mcp") + assert tokens_db.revoke_user_token(temp_db, token_id, 1) is True + assert tokens_db.revoke_user_token(temp_db, token_id, 1) is False