114 lines
3.1 KiB
Python
114 lines
3.1 KiB
Python
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
|