BIG BIG BIG BIG BIG UPDATES
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
gabdevele
co-authored by
Copilot
Unverified
Copilot
parent
46f7cd6a5f
commit
93ba8e9ac2
@@ -0,0 +1,102 @@
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from core.mongodb import get_database
|
||||
from models.auth import AuthUserDocument
|
||||
from motor.motor_asyncio import AsyncIOMotorCollection
|
||||
|
||||
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
|
||||
@@ -0,0 +1,70 @@
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from core.mongodb import get_database
|
||||
from models.example import AddonExampleCreateInput, AddonExampleDocument
|
||||
from motor.motor_asyncio import AsyncIOMotorCollection
|
||||
|
||||
COLLECTION_NAME = "addon_examples"
|
||||
|
||||
|
||||
def utcnow() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
async def collection() -> AsyncIOMotorCollection[dict[str, Any]]:
|
||||
database = await get_database()
|
||||
return database[COLLECTION_NAME]
|
||||
|
||||
|
||||
def deserialize(raw: Mapping[str, Any] | None) -> AddonExampleDocument | None:
|
||||
if raw is None:
|
||||
return None
|
||||
document = dict(raw)
|
||||
document.pop("_id", None)
|
||||
return AddonExampleDocument.model_validate(document)
|
||||
|
||||
|
||||
async def ensure_example_indexes() -> None:
|
||||
examples = await collection()
|
||||
await examples.create_index("id", unique=True, name="uq_addon_example_id")
|
||||
|
||||
|
||||
async def list_custom_examples() -> list[AddonExampleDocument]:
|
||||
examples = await collection()
|
||||
cursor = examples.find({}).sort("name", 1)
|
||||
|
||||
result: list[AddonExampleDocument] = []
|
||||
async for raw in cursor:
|
||||
example = deserialize(raw)
|
||||
if example is not None:
|
||||
result.append(example)
|
||||
return result
|
||||
|
||||
|
||||
async def create_custom_example(
|
||||
payload: AddonExampleCreateInput,
|
||||
) -> AddonExampleDocument:
|
||||
now = utcnow()
|
||||
example = AddonExampleDocument(
|
||||
id=payload.id,
|
||||
name=payload.name,
|
||||
description=payload.description,
|
||||
category=payload.category,
|
||||
code=payload.code,
|
||||
source_url=payload.source_url,
|
||||
built_in=False,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
|
||||
examples = await collection()
|
||||
await examples.insert_one(example.model_dump(exclude_none=True))
|
||||
return example
|
||||
|
||||
|
||||
async def delete_custom_example(example_id: str) -> bool:
|
||||
examples = await collection()
|
||||
result = await examples.delete_one({"id": example_id})
|
||||
return result.deleted_count > 0
|
||||
@@ -0,0 +1,113 @@
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from core.mongodb import get_database
|
||||
from models.proxy import ProxyCreateInput, ProxyDocument, ProxyUpdateInput
|
||||
from motor.motor_asyncio import AsyncIOMotorCollection
|
||||
from pymongo import ReturnDocument
|
||||
|
||||
COLLECTION_NAME = "proxies"
|
||||
|
||||
|
||||
def _utcnow() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
async def _collection() -> AsyncIOMotorCollection[dict[str, Any]]:
|
||||
database = await get_database()
|
||||
return database[COLLECTION_NAME]
|
||||
|
||||
|
||||
def _deserialize_proxy(raw: Mapping[str, Any] | None) -> ProxyDocument | None:
|
||||
if raw is None:
|
||||
return None
|
||||
document = dict(raw)
|
||||
document.pop("_id", None)
|
||||
return ProxyDocument.model_validate(document)
|
||||
|
||||
|
||||
async def ensure_proxy_indexes() -> None:
|
||||
collection = await _collection()
|
||||
await collection.create_index("proxy_port", unique=True, name="uq_proxy_port")
|
||||
|
||||
|
||||
async def create_proxy(payload: ProxyCreateInput) -> ProxyDocument:
|
||||
now = _utcnow()
|
||||
proxy = ProxyDocument(
|
||||
proxy_port=payload.proxy_port,
|
||||
name=payload.name,
|
||||
code=payload.code,
|
||||
mode=payload.mode,
|
||||
mode_spec=payload.mode_spec,
|
||||
protocol=payload.protocol,
|
||||
listen_host=payload.listen_host,
|
||||
mitm_options=payload.mitm_options,
|
||||
healthcheck=payload.healthcheck,
|
||||
service=payload.service,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
|
||||
collection = await _collection()
|
||||
await collection.insert_one(proxy.model_dump(exclude_none=True))
|
||||
return proxy
|
||||
|
||||
|
||||
async def get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
collection = await _collection()
|
||||
raw = await collection.find_one({"proxy_port": proxy_port})
|
||||
return _deserialize_proxy(raw)
|
||||
|
||||
|
||||
async def list_proxies() -> list[ProxyDocument]:
|
||||
collection = await _collection()
|
||||
cursor = collection.find({}).sort("proxy_port", 1)
|
||||
|
||||
proxies: list[ProxyDocument] = []
|
||||
async for raw in cursor:
|
||||
proxy = _deserialize_proxy(raw)
|
||||
if proxy is not None:
|
||||
proxies.append(proxy)
|
||||
return proxies
|
||||
|
||||
|
||||
async def update_proxy(
|
||||
proxy_port: int,
|
||||
updates: ProxyUpdateInput,
|
||||
) -> ProxyDocument | None:
|
||||
payload = updates.model_dump(exclude_none=True)
|
||||
if not payload:
|
||||
return await get_proxy_by_port(proxy_port)
|
||||
|
||||
payload["updated_at"] = _utcnow()
|
||||
|
||||
collection = await _collection()
|
||||
raw = await collection.find_one_and_update(
|
||||
{"proxy_port": proxy_port},
|
||||
{"$set": payload},
|
||||
return_document=ReturnDocument.AFTER,
|
||||
)
|
||||
return _deserialize_proxy(raw)
|
||||
|
||||
|
||||
async def update_proxy_runtime(
|
||||
proxy_port: int,
|
||||
values: dict[str, object],
|
||||
) -> ProxyDocument | None:
|
||||
payload = dict(values)
|
||||
payload["updated_at"] = _utcnow()
|
||||
|
||||
collection = await _collection()
|
||||
raw = await collection.find_one_and_update(
|
||||
{"proxy_port": proxy_port},
|
||||
{"$set": payload},
|
||||
return_document=ReturnDocument.AFTER,
|
||||
)
|
||||
return _deserialize_proxy(raw)
|
||||
|
||||
|
||||
async def delete_proxy(proxy_port: int) -> bool:
|
||||
collection = await _collection()
|
||||
result = await collection.delete_one({"proxy_port": proxy_port})
|
||||
return result.deleted_count > 0
|
||||
Reference in New Issue
Block a user