From b4a84e70a288d3c5540facfc9d8701e9c85282b7 Mon Sep 17 00:00:00 2001 From: gabdevele Date: Sun, 12 Apr 2026 00:55:29 +0200 Subject: [PATCH] feat: implement proxy management endpoints and MongoDB integration --- app/api/proxy/router.py | 55 +++++++++++++++++---- app/core/mongodb.py | 12 +++++ app/main.py | 22 ++++++++- app/models/proxy.py | 36 ++++++++++++-- app/models/service.py | 12 +++-- app/utils/proxy/manager.py | 97 +++++++++++++++++++++++++++++++++++--- 6 files changed, 209 insertions(+), 25 deletions(-) diff --git a/app/api/proxy/router.py b/app/api/proxy/router.py index a55a94c..ef12b0d 100644 --- a/app/api/proxy/router.py +++ b/app/api/proxy/router.py @@ -1,4 +1,13 @@ -from fastapi import APIRouter +from fastapi import APIRouter, HTTPException, status + +from models.proxy import ProxyCreateInput, ProxyUpdateInput +from utils.proxy.manager import ( + create_proxy, + delete_proxy, + get_proxy_by_port, + list_proxies as list_proxy_documents, + update_proxy, +) router = APIRouter( prefix="/proxy", @@ -7,15 +16,45 @@ router = APIRouter( @router.get("/") -async def list_proxies(): - return +async def list_proxies_handler(): + return await list_proxy_documents() -@router.post("/") -async def create_proxy(): - return +@router.get("/{proxy_port}/") +async def get_proxy(proxy_port: int): + proxy = await get_proxy_by_port(proxy_port) + if proxy is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found") + return proxy + + +@router.post("/", status_code=status.HTTP_201_CREATED) +async def create_proxy_handler(payload: ProxyCreateInput): + existing_proxy = await get_proxy_by_port(payload.proxy_port) + if existing_proxy is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Proxy already exists on port {payload.proxy_port}", + ) + + return await create_proxy( + proxy_port=payload.proxy_port, + name=payload.name, + service={"address": payload.service.address, "port": payload.service.port}, + code=payload.code, + ) @router.put("/{proxy_port}/") -async def edit_proxy(): - return +async def edit_proxy(proxy_port: int, payload: ProxyUpdateInput): + updated_proxy = await update_proxy(proxy_port, payload) + if updated_proxy is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found") + return updated_proxy + + +@router.delete("/{proxy_port}/", status_code=status.HTTP_204_NO_CONTENT) +async def remove_proxy(proxy_port: int): + deleted = await delete_proxy(proxy_port) + if not deleted: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Proxy not found") diff --git a/app/core/mongodb.py b/app/core/mongodb.py index 7b0d65f..3f73dc1 100644 --- a/app/core/mongodb.py +++ b/app/core/mongodb.py @@ -13,6 +13,18 @@ class _MongoDB: mongodb = _MongoDB() + +async def init_mongodb() -> None: + mongodb.client = AsyncIOMotorClient(settings.MONGODB_URI) + await ensure_mongodb_available() + + +async def close_mongodb() -> None: + client = mongodb.client + mongodb.client = None + if client is not None: + client.close() + async def get_database() -> AsyncIOMotorDatabase: return mongodb.get_database() diff --git a/app/main.py b/app/main.py index 5fdb9da..42a662c 100644 --- a/app/main.py +++ b/app/main.py @@ -1,9 +1,27 @@ +from contextlib import asynccontextmanager + from fastapi import FastAPI -from api.proxy import router as proxy_router +from api.proxy.router import router as proxy_router +from core.mongodb import close_mongodb, init_mongodb +from core.redis import close_redis, init_redis +from utils.proxy.manager import ensure_proxy_indexes + + +@asynccontextmanager +async def lifespan(_: FastAPI): + await init_mongodb() + await init_redis() + await ensure_proxy_indexes() + try: + yield + finally: + await close_redis() + await close_mongodb() app = FastAPI( - title="thePlugAPI" + title="thePlugAPI", + lifespan=lifespan, ) app.include_router(proxy_router) diff --git a/app/models/proxy.py b/app/models/proxy.py index 4864a63..2fc78e9 100644 --- a/app/models/proxy.py +++ b/app/models/proxy.py @@ -1,8 +1,36 @@ -from pydantic import BaseModel, Field +from datetime import datetime -from models.service import ServiceDocument +from pydantic import BaseModel, ConfigDict, Field + +from models.service import ServiceEndpointDocument class ProxyDocument(BaseModel): - port: int = Field(..., description="Port number of the proxy") - service: ServiceDocument = Field(..., description="Associated service details") \ No newline at end of file + model_config = ConfigDict(extra="forbid") + + proxy_port: int = Field(..., description="Port exposed by the proxy") + name: str = Field(..., description="Human readable proxy name") + code: str = Field(..., description="Proxy source or identifier") + service: ServiceEndpointDocument = Field(..., description="Associated upstream service") + created_at: datetime | None = Field(default=None, description="Creation timestamp in UTC") + updated_at: datetime | None = Field(default=None, description="Last update timestamp in UTC") + + +class ProxyCreateInput(BaseModel): + model_config = ConfigDict(extra="forbid") + + proxy_port: int = Field(..., description="Port exposed by the proxy") + name: str = Field(..., description="Human readable proxy name") + code: str = Field(..., description="Proxy source or identifier") + service: ServiceEndpointDocument = Field(..., description="Associated upstream service") + + +class ProxyUpdateInput(BaseModel): + model_config = ConfigDict(extra="forbid") + + name: str | None = Field(default=None, description="Human readable proxy name") + code: str | None = Field(default=None, description="Proxy source or identifier") + service: ServiceEndpointDocument | None = Field( + default=None, + description="Associated upstream service", + ) \ No newline at end of file diff --git a/app/models/service.py b/app/models/service.py index 7caff45..5d67d28 100644 --- a/app/models/service.py +++ b/app/models/service.py @@ -1,6 +1,8 @@ -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field -class ServiceDocument(BaseModel): - name: str = Field(..., description="Name of the service") - port: int = Field(..., description="Port number of the service") - category: str = Field(..., description="Category of the service (e.g., web, pwn, etc.)") + +class ServiceEndpointDocument(BaseModel): + model_config = ConfigDict(extra="forbid") + + address: str = Field(..., description="Upstream service address") + port: int = Field(..., description="Upstream service port") diff --git a/app/utils/proxy/manager.py b/app/utils/proxy/manager.py index 6434ac6..c518448 100644 --- a/app/utils/proxy/manager.py +++ b/app/utils/proxy/manager.py @@ -1,10 +1,95 @@ +from datetime import datetime, timezone +from typing import Any, Mapping + +from motor.motor_asyncio import AsyncIOMotorCollection +from pymongo import ReturnDocument + +from core.mongodb import get_database +from models.proxy import ProxyDocument, ProxyUpdateInput +from models.service import ServiceEndpointDocument from utils.proxy.types import ServiceConfig +COLLECTION_NAME = "proxies" + + +def _utcnow() -> datetime: + return datetime.now(timezone.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( - port: int, - name: str, - service: ServiceConfig, - code: str -): - return + proxy_port: int, + name: str, + service: ServiceConfig, + code: str, +) -> ProxyDocument: + now = _utcnow() + proxy = ProxyDocument( + proxy_port=proxy_port, + name=name, + code=code, + service=ServiceEndpointDocument(address=service["address"], port=service["port"]), + 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 delete_proxy(proxy_port: int) -> bool: + collection = await _collection() + result = await collection.delete_one({"proxy_port": proxy_port}) + return result.deleted_count > 0