104 lines
2.8 KiB
Python
104 lines
2.8 KiB
Python
from collections.abc import Mapping
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorCollection
|
|
|
|
from app.core.mongodb import get_database
|
|
from app.models.auth import AuthUserDocument
|
|
|
|
USERS_COLLECTION_NAME = "auth_users"
|
|
REVOKED_COLLECTION_NAME = "revoked_tokens"
|
|
|
|
|
|
def utcnow() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
async def users_collection() -> AsyncIOMotorCollection[dict[str, Any]]:
|
|
database = await get_database()
|
|
return database[USERS_COLLECTION_NAME]
|
|
|
|
|
|
async def revoked_tokens_collection() -> AsyncIOMotorCollection[dict[str, Any]]:
|
|
database = await get_database()
|
|
return database[REVOKED_COLLECTION_NAME]
|
|
|
|
|
|
def deserialize_user(raw: Mapping[str, Any] | None) -> AuthUserDocument | None:
|
|
if raw is None:
|
|
return None
|
|
document = dict(raw)
|
|
document.pop("_id", None)
|
|
return AuthUserDocument.model_validate(document)
|
|
|
|
|
|
async def ensure_auth_indexes() -> None:
|
|
users = await users_collection()
|
|
revoked = await revoked_tokens_collection()
|
|
|
|
await users.create_index("username", unique=True, name="uq_auth_username")
|
|
await revoked.create_index("jti", unique=True, name="uq_revoked_jti")
|
|
await revoked.create_index(
|
|
"expires_at",
|
|
expireAfterSeconds=0,
|
|
name="ttl_revoked_tokens",
|
|
)
|
|
|
|
|
|
async def ensure_single_account(*, username: str, password_hash: str) -> None:
|
|
users = await users_collection()
|
|
now = utcnow()
|
|
|
|
current = await users.find_one({}, sort=[("created_at", 1)])
|
|
if current is None:
|
|
await users.insert_one(
|
|
{
|
|
"username": username,
|
|
"password_hash": password_hash,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}
|
|
)
|
|
return
|
|
|
|
current_id = current["_id"]
|
|
await users.update_one(
|
|
{"_id": current_id},
|
|
{
|
|
"$set": {
|
|
"username": username,
|
|
"password_hash": password_hash,
|
|
"updated_at": now,
|
|
}
|
|
},
|
|
)
|
|
await users.delete_many({"_id": {"$ne": current_id}})
|
|
|
|
|
|
async def get_user_by_username(username: str) -> AuthUserDocument | None:
|
|
users = await users_collection()
|
|
raw = await users.find_one({"username": username})
|
|
return deserialize_user(raw)
|
|
|
|
|
|
async def revoke_token(*, jti: str, expires_at: datetime) -> None:
|
|
revoked = await revoked_tokens_collection()
|
|
await revoked.update_one(
|
|
{"jti": jti},
|
|
{
|
|
"$set": {
|
|
"jti": jti,
|
|
"expires_at": expires_at,
|
|
"revoked_at": utcnow(),
|
|
}
|
|
},
|
|
upsert=True,
|
|
)
|
|
|
|
|
|
async def is_token_revoked(jti: str) -> bool:
|
|
revoked = await revoked_tokens_collection()
|
|
raw = await revoked.find_one({"jti": jti}, projection={"_id": 1})
|
|
return raw is not None
|