Merge pull request 'feature/crud-mongodb-proxy' (#3) from feature/crud-mongodb-proxy into master

Reviewed-on: #3

.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

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("/")
+@router.get("")
-async def list_proxies():
+async def list_proxies_handler():
-    return
+    return await list_proxy_documents()
 
 
-@router.post("/")
+@router.get("/{proxy_port}")
-async def create_proxy():
+async def get_proxy(proxy_port: int):
-    return
+    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}/")
+@router.post("", status_code=status.HTTP_201_CREATED)
-async def edit_proxy():
+async def create_proxy_handler(payload: ProxyCreateInput):
-    return
+    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")

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()
 

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

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)

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")
+    model_config = ConfigDict(extra="forbid")
-    service: ServiceDocument = Field(..., description="Associated service details")
+
+    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",
+    )

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")
+class ServiceEndpointDocument(BaseModel):
-    port: int = Field(..., description="Port number of the service")
+    model_config = ConfigDict(extra="forbid")
-    category: str = Field(..., description="Category of the service (e.g., web, pwn, etc.)")
+
+    address: str = Field(..., description="Upstream service address")
+    port: int = Field(..., description="Upstream service port")

app/utils/proxy/manager.py

@@ -7,4 +7,4 @@ async def create_proxy(
         service: ServiceConfig,
         code: str
 ):
-    return
+    return