diff --git a/.env.example b/.env.example index f6f3450..72c629b 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,5 @@ API_PORT=3671 MONGO_INITDB_ROOT_PASSWORD=pass MONGODB_URL=mongodb://theplug:pass@mongodb:27017/theplug?authSource=admin +MONGODB_DB_NAME=theplug +REDIS_URL=redis://redis:6379 \ No newline at end of file diff --git a/app/api/proxy/router.py b/app/api/proxy/router.py index a55a94c..ccffb56 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 db.proxy import ( + create_proxy, + delete_proxy, + get_proxy_by_port, + list_proxies as list_proxy_documents, + update_proxy, +) +from models.proxy import ProxyCreateInput, ProxyUpdateInput router = APIRouter( prefix="/proxy", @@ -6,16 +15,46 @@ router = APIRouter( ) -@router.get("/") -async def list_proxies(): - return +@router.get("") +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.put("/{proxy_port}/") -async def edit_proxy(): - return +@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(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/__init__.py b/app/core/__init__.py new file mode 100644 index 0000000..e69de29 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/db/__init__.py b/app/db/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/db/proxy.py b/app/db/proxy.py new file mode 100644 index 0000000..5096b91 --- /dev/null +++ b/app/db/proxy.py @@ -0,0 +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( + 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 diff --git a/app/main.py b/app/main.py index 5fdb9da..febde27 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 db.proxy import ensure_proxy_indexes +from core.mongodb import close_mongodb, init_mongodb +from core.redis import close_redis, init_redis + + +@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/__init__.py b/app/models/__init__.py new file mode 100644 index 0000000..e69de29 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..58ad605 100644 --- a/app/utils/proxy/manager.py +++ b/app/utils/proxy/manager.py @@ -7,4 +7,4 @@ async def create_proxy( service: ServiceConfig, code: str ): - return + return \ No newline at end of file