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,11 @@
|
||||
API_PORT=3671
|
||||
FRONTEND_PORT=5173
|
||||
MONGO_INITDB_ROOT_PASSWORD=pass# esegui 'openssl rand -hex 5' in un terminale e copia qui
|
||||
MONGODB_URI=mongodb://theplug:pass@localhost:27017/theplug?authSource=admin
|
||||
MONGODB_DB_NAME=theplug
|
||||
REDIS_URL=redis://localhost:6379
|
||||
JWT_SECRET_KEY= # esegui 'openssl rand -hex 20' in un terminale e copia qui
|
||||
JWT_ALGORITHM=HS256
|
||||
JWT_EXPIRE_MINUTES=30
|
||||
AUTH_USERNAME=admin
|
||||
AUTH_PASSWORD= # esegui 'openssl rand -hex 20' in un terminale e copia qui
|
||||
@@ -0,0 +1,27 @@
|
||||
FROM python:3.14-slim
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PYTHONPATH=/app/src
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY pyproject.toml ./
|
||||
RUN pip install --no-cache-dir --upgrade pip && \
|
||||
pip install --no-cache-dir \
|
||||
"authlib>=1.6.0" \
|
||||
"basedpyright>=1.39.0" \
|
||||
"fastapi[standard]>=0.135.3" \
|
||||
"mitmproxy>=12.2.2" \
|
||||
"motor>=3.7.1" \
|
||||
"pwdlib[argon2,bcrypt]>=0.3.0" \
|
||||
"pyyaml>=6.0.3" \
|
||||
"pydantic-settings>=2.13.1" \
|
||||
"redis>=7.4.0" \
|
||||
"ruff>=0.15.10"
|
||||
|
||||
COPY app ./src
|
||||
|
||||
EXPOSE 3671
|
||||
|
||||
CMD ["fastapi", "run", "src/main.py", "--host", "0.0.0.0", "--port", "3671"]
|
||||
@@ -0,0 +1,3 @@
|
||||
from api.auth.router import router
|
||||
|
||||
__all__ = ["router"]
|
||||
@@ -0,0 +1,41 @@
|
||||
from datetime import UTC, datetime
|
||||
from typing import Annotated
|
||||
|
||||
from core.security import (
|
||||
AuthContext,
|
||||
create_access_token,
|
||||
get_auth_context,
|
||||
verify_password,
|
||||
)
|
||||
from fastapi import APIRouter, Depends, HTTPException, Response, status
|
||||
from models.auth import LoginInput, TokenResponse
|
||||
from repositories.auth import get_user_by_username, revoke_token
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/auth",
|
||||
tags=["Auth"],
|
||||
)
|
||||
|
||||
|
||||
@router.post("/login", response_model=TokenResponse)
|
||||
async def login_handler(payload: LoginInput):
|
||||
user = await get_user_by_username(payload.username)
|
||||
if user is None or not verify_password(payload.password, user.password_hash):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid username or password",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
token = create_access_token(subject=user.username)
|
||||
return TokenResponse(access_token=token)
|
||||
|
||||
|
||||
@router.post("/logout", status_code=status.HTTP_204_NO_CONTENT)
|
||||
async def logout_handler(
|
||||
auth_context: Annotated[AuthContext, Depends(get_auth_context)],
|
||||
) -> Response:
|
||||
expires_at = datetime.fromtimestamp(auth_context.payload.exp, tz=UTC)
|
||||
if expires_at > datetime.now(UTC):
|
||||
await revoke_token(jti=auth_context.payload.jti, expires_at=expires_at)
|
||||
return Response(status_code=status.HTTP_204_NO_CONTENT)
|
||||
@@ -0,0 +1,62 @@
|
||||
from core.security import (
|
||||
InvalidTokenError,
|
||||
authenticate_token,
|
||||
require_authenticated_user,
|
||||
)
|
||||
from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
|
||||
from services.editor import python_lsp
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
router = APIRouter(prefix="/proxy/editor/python", tags=["Editor"])
|
||||
|
||||
|
||||
@router.get("/session", dependencies=[Depends(require_authenticated_user)])
|
||||
async def python_editor_session():
|
||||
if python_lsp.lsp_binary() is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="basedpyright-langserver is not installed",
|
||||
)
|
||||
|
||||
session = python_lsp.get_session()
|
||||
return {
|
||||
"root_uri": session.root_uri,
|
||||
"file_uri": session.file_uri,
|
||||
}
|
||||
|
||||
|
||||
async def authenticate_editor_socket(websocket: WebSocket) -> None:
|
||||
token = websocket.query_params.get("token")
|
||||
if token is None:
|
||||
authorization = websocket.headers.get("authorization", "")
|
||||
scheme, _, value = authorization.partition(" ")
|
||||
if scheme.lower() == "bearer" and value:
|
||||
token = value
|
||||
|
||||
if not token:
|
||||
raise InvalidTokenError
|
||||
|
||||
await authenticate_token(token)
|
||||
|
||||
|
||||
@router.websocket("/ws")
|
||||
async def python_editor_socket(websocket: WebSocket):
|
||||
try:
|
||||
await authenticate_editor_socket(websocket)
|
||||
except InvalidTokenError:
|
||||
await websocket.close(code=4401)
|
||||
return
|
||||
|
||||
if python_lsp.lsp_binary() is None:
|
||||
await websocket.close(code=1011, reason="basedpyright unavailable")
|
||||
return
|
||||
|
||||
session = python_lsp.get_session()
|
||||
await websocket.accept()
|
||||
|
||||
try:
|
||||
await python_lsp.bridge(websocket, session)
|
||||
except WebSocketDisconnect:
|
||||
return
|
||||
except RuntimeError:
|
||||
await websocket.close(code=1011, reason="language server failed")
|
||||
@@ -0,0 +1,300 @@
|
||||
from pathlib import Path
|
||||
|
||||
from core.security import require_authenticated_user
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from models.example import AddonExampleCreateInput
|
||||
from models.proxy import ProxyCreateInput, ProxyUpdateInput
|
||||
from pymongo.errors import DuplicateKeyError
|
||||
from repositories.examples import (
|
||||
create_custom_example,
|
||||
delete_custom_example,
|
||||
list_custom_examples,
|
||||
)
|
||||
from repositories.proxy import (
|
||||
create_proxy,
|
||||
delete_proxy,
|
||||
get_proxy_by_port,
|
||||
update_proxy,
|
||||
update_proxy_runtime,
|
||||
)
|
||||
from repositories.proxy import (
|
||||
list_proxies as list_proxy_documents,
|
||||
)
|
||||
from services.mitmproxy.examples import BUILTIN_EXAMPLES
|
||||
from services.mitmproxy.health import check_proxy
|
||||
from services.mitmproxy.manager import MitmdumpError, manager
|
||||
from services.mitmproxy.options import (
|
||||
MitmOptionError,
|
||||
load_options,
|
||||
public_options,
|
||||
validate_options,
|
||||
)
|
||||
from services.mitmproxy.validator import ProxyCodeValidationError, validate_proxy_code
|
||||
|
||||
SUPPORTED_MODES = {
|
||||
"regular",
|
||||
"local",
|
||||
"wireguard",
|
||||
"reverse",
|
||||
"transparent",
|
||||
"tun",
|
||||
"upstream",
|
||||
"socks5",
|
||||
"dns",
|
||||
}
|
||||
SUPPORTED_PROTOCOLS = {
|
||||
"http",
|
||||
"https",
|
||||
"tcp",
|
||||
"tls",
|
||||
"udp",
|
||||
"dtls",
|
||||
"dns",
|
||||
"http3",
|
||||
"quic",
|
||||
}
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/proxy",
|
||||
tags=["Proxy"],
|
||||
dependencies=[Depends(require_authenticated_user)],
|
||||
)
|
||||
|
||||
|
||||
async def _validate_code_or_raise(code: str) -> None:
|
||||
try:
|
||||
await validate_proxy_code(code)
|
||||
except ProxyCodeValidationError as error:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail=error.to_http_detail(),
|
||||
) from error
|
||||
|
||||
|
||||
def validate_mode(mode: str) -> None:
|
||||
if mode not in SUPPORTED_MODES:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail=f"Unsupported proxy mode: {mode}",
|
||||
)
|
||||
|
||||
|
||||
def validate_mode_spec(mode_spec: str | None) -> None:
|
||||
if mode_spec is None:
|
||||
return
|
||||
mode = mode_spec.split(":", 1)[0].split("@", 1)[0]
|
||||
validate_mode(mode)
|
||||
|
||||
|
||||
def validate_protocol(protocol: str) -> None:
|
||||
if protocol not in SUPPORTED_PROTOCOLS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail=f"Unsupported reverse proxy protocol: {protocol}",
|
||||
)
|
||||
|
||||
|
||||
async def validate_mitm_options(options: dict[str, object]) -> None:
|
||||
try:
|
||||
await validate_options(options)
|
||||
except MitmOptionError as error:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail=str(error),
|
||||
) from error
|
||||
|
||||
|
||||
async def get_proxy_or_404(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.get("")
|
||||
async def list_proxies_handler():
|
||||
return await list_proxy_documents()
|
||||
|
||||
|
||||
@router.get("/services")
|
||||
async def list_services(request: Request):
|
||||
return getattr(request.app.state, "services", [])
|
||||
|
||||
|
||||
@router.get("/mitmproxy/options")
|
||||
async def list_mitmproxy_options():
|
||||
return public_options(await load_options())
|
||||
|
||||
|
||||
@router.get("/mitmproxy/examples")
|
||||
async def list_mitmproxy_examples():
|
||||
return [*BUILTIN_EXAMPLES, *await list_custom_examples()]
|
||||
|
||||
|
||||
@router.post("/mitmproxy/examples", status_code=status.HTTP_201_CREATED)
|
||||
async def create_mitmproxy_example(payload: AddonExampleCreateInput):
|
||||
await _validate_code_or_raise(payload.code)
|
||||
try:
|
||||
return await create_custom_example(payload)
|
||||
except DuplicateKeyError as error:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"Example already exists: {payload.id}",
|
||||
) from error
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/mitmproxy/examples/{example_id}",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def delete_mitmproxy_example(example_id: str):
|
||||
if any(example.id == example_id for example in BUILTIN_EXAMPLES):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Built-in examples cannot be deleted",
|
||||
)
|
||||
|
||||
deleted = await delete_custom_example(example_id)
|
||||
if not deleted:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Example not found",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{proxy_port}")
|
||||
async def get_proxy(proxy_port: int):
|
||||
return await get_proxy_or_404(proxy_port)
|
||||
|
||||
|
||||
@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}",
|
||||
)
|
||||
|
||||
await _validate_code_or_raise(payload.code)
|
||||
validate_mode(payload.mode)
|
||||
validate_mode_spec(payload.mode_spec)
|
||||
validate_protocol(payload.protocol)
|
||||
await validate_mitm_options(payload.mitm_options)
|
||||
|
||||
return await create_proxy(payload)
|
||||
|
||||
|
||||
@router.put("/{proxy_port}")
|
||||
async def edit_proxy(proxy_port: int, payload: ProxyUpdateInput):
|
||||
if payload.code is not None:
|
||||
await _validate_code_or_raise(payload.code)
|
||||
if payload.mode is not None:
|
||||
validate_mode(payload.mode)
|
||||
if "mode_spec" in payload.model_fields_set:
|
||||
validate_mode_spec(payload.mode_spec)
|
||||
if payload.protocol is not None:
|
||||
validate_protocol(payload.protocol)
|
||||
if payload.mitm_options is not None:
|
||||
await validate_mitm_options(payload.mitm_options)
|
||||
|
||||
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.post("/{proxy_port}/start")
|
||||
async def start_proxy(proxy_port: int):
|
||||
proxy = await get_proxy_or_404(proxy_port)
|
||||
|
||||
try:
|
||||
runtime = await manager.start(proxy)
|
||||
except MitmdumpError as error:
|
||||
updated = await update_proxy_runtime(
|
||||
proxy_port,
|
||||
{"status": "error", "last_error": str(error)},
|
||||
)
|
||||
return updated or proxy
|
||||
|
||||
return await update_proxy_runtime(proxy_port, runtime)
|
||||
|
||||
|
||||
@router.post("/{proxy_port}/stop")
|
||||
async def stop_proxy(proxy_port: int):
|
||||
proxy = await get_proxy_or_404(proxy_port)
|
||||
runtime = await manager.stop(proxy)
|
||||
return await update_proxy_runtime(proxy_port, runtime)
|
||||
|
||||
|
||||
@router.post("/{proxy_port}/restart")
|
||||
async def restart_proxy(proxy_port: int):
|
||||
proxy = await get_proxy_or_404(proxy_port)
|
||||
stopped = await manager.stop(proxy)
|
||||
proxy = await update_proxy_runtime(proxy_port, stopped) or proxy
|
||||
try:
|
||||
runtime = await manager.start(proxy)
|
||||
except MitmdumpError as error:
|
||||
updated = await update_proxy_runtime(
|
||||
proxy_port,
|
||||
{"status": "error", "last_error": str(error)},
|
||||
)
|
||||
return updated or proxy
|
||||
return await update_proxy_runtime(proxy_port, runtime)
|
||||
|
||||
|
||||
@router.get("/{proxy_port}/status")
|
||||
async def proxy_status(proxy_port: int):
|
||||
proxy = await get_proxy_or_404(proxy_port)
|
||||
runtime_status = await manager.status(proxy)
|
||||
if runtime_status != proxy.status:
|
||||
values: dict[str, object] = {"status": runtime_status}
|
||||
if runtime_status != "running":
|
||||
values["pid"] = None
|
||||
proxy = await update_proxy_runtime(proxy_port, values) or proxy
|
||||
return {"proxy_port": proxy.proxy_port, "status": runtime_status, "pid": proxy.pid}
|
||||
|
||||
|
||||
@router.get("/{proxy_port}/health")
|
||||
async def proxy_health(proxy_port: int):
|
||||
proxy = await get_proxy_or_404(proxy_port)
|
||||
return await check_proxy(proxy)
|
||||
|
||||
|
||||
@router.get("/{proxy_port}/logs")
|
||||
async def proxy_logs(proxy_port: int):
|
||||
proxy = await get_proxy_or_404(proxy_port)
|
||||
if proxy.log_path is None:
|
||||
return {"proxy_port": proxy.proxy_port, "log": ""}
|
||||
|
||||
path = Path(proxy.log_path)
|
||||
if not path.is_file():
|
||||
return {"proxy_port": proxy.proxy_port, "log": ""}
|
||||
|
||||
return {
|
||||
"proxy_port": proxy.proxy_port,
|
||||
"log": path.read_text(encoding="utf-8", errors="replace")[-8000:],
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/{proxy_port}", status_code=status.HTTP_204_NO_CONTENT)
|
||||
async def remove_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",
|
||||
)
|
||||
await manager.stop(proxy)
|
||||
deleted = await delete_proxy(proxy_port)
|
||||
if not deleted:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Proxy not found",
|
||||
)
|
||||
@@ -0,0 +1,23 @@
|
||||
from pydantic import Field
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=(".env", "backend/.env"),
|
||||
env_file_encoding="utf-8",
|
||||
case_sensitive=False,
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
MONGODB_URI: str
|
||||
MONGODB_DB_NAME: str
|
||||
REDIS_URL: str
|
||||
JWT_SECRET_KEY: str
|
||||
JWT_ALGORITHM: str = "HS256"
|
||||
JWT_EXPIRE_MINUTES: int = Field(default=30, gt=0)
|
||||
AUTH_USERNAME: str
|
||||
AUTH_PASSWORD: str
|
||||
|
||||
|
||||
settings = Settings()
|
||||
@@ -0,0 +1,37 @@
|
||||
from core.config import settings
|
||||
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
||||
|
||||
|
||||
class _MongoDB:
|
||||
client: AsyncIOMotorClient | None = None
|
||||
|
||||
def get_database(self) -> AsyncIOMotorDatabase:
|
||||
if self.client is None:
|
||||
raise RuntimeError("MongoDB client has not been initialized")
|
||||
return self.client[settings.MONGODB_DB_NAME]
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
async def ensure_mongodb_available() -> None:
|
||||
client = mongodb.client
|
||||
if client is None:
|
||||
raise RuntimeError("MongoDB client has not been initialized")
|
||||
await client.admin.command("ping")
|
||||
@@ -0,0 +1,22 @@
|
||||
import redis.asyncio as aioredis
|
||||
from core.config import settings
|
||||
|
||||
_redis_pool: aioredis.Redis | None = None
|
||||
|
||||
|
||||
async def init_redis() -> None:
|
||||
global _redis_pool
|
||||
_redis_pool = aioredis.from_url(
|
||||
settings.REDIS_URL, encoding="utf-8", decode_responses=True
|
||||
)
|
||||
|
||||
|
||||
async def close_redis() -> None:
|
||||
if _redis_pool:
|
||||
await _redis_pool.aclose()
|
||||
|
||||
|
||||
async def get_redis() -> aioredis.Redis:
|
||||
if _redis_pool is None:
|
||||
raise RuntimeError("Redis pool has not been initialized")
|
||||
return _redis_pool
|
||||
@@ -0,0 +1,116 @@
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Annotated
|
||||
from uuid import uuid4
|
||||
|
||||
from authlib.jose import jwt # type: ignore
|
||||
from authlib.jose.errors import JoseError # type: ignore
|
||||
from core.config import settings
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
from models.auth import AccessTokenPayload, AuthUserDocument
|
||||
from pwdlib import PasswordHash
|
||||
from pwdlib.hashers.argon2 import Argon2Hasher
|
||||
from repositories.auth import get_user_by_username, is_token_revoked
|
||||
|
||||
pwd_context = PasswordHash((Argon2Hasher(),))
|
||||
bearer_scheme = HTTPBearer(auto_error=False)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AuthContext:
|
||||
token: str
|
||||
user: AuthUserDocument
|
||||
payload: AccessTokenPayload
|
||||
|
||||
|
||||
class InvalidTokenError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def token_auth_error() -> HTTPException:
|
||||
return HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Could not validate credentials",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
|
||||
def hash_password(password: str) -> str:
|
||||
return pwd_context.hash(password)
|
||||
|
||||
|
||||
def verify_password(plain_password: str, password_hash: str) -> bool:
|
||||
if hasattr(pwd_context, "verify"):
|
||||
return pwd_context.verify(plain_password, password_hash)
|
||||
return pwd_context.check(plain_password, password_hash)
|
||||
|
||||
|
||||
def create_access_token(*, subject: str) -> str:
|
||||
now = datetime.now(UTC)
|
||||
expires_delta = timedelta(minutes=settings.JWT_EXPIRE_MINUTES)
|
||||
expires_at = int((now + expires_delta).timestamp())
|
||||
|
||||
header = {
|
||||
"alg": settings.JWT_ALGORITHM,
|
||||
"typ": "JWT",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"sub": subject,
|
||||
"jti": str(uuid4()),
|
||||
"exp": expires_at,
|
||||
}
|
||||
|
||||
token = jwt.encode(header, payload, settings.JWT_SECRET_KEY)
|
||||
|
||||
if isinstance(token, bytes):
|
||||
return token.decode("utf-8")
|
||||
|
||||
return token
|
||||
|
||||
|
||||
def decode_access_token(token: str) -> AccessTokenPayload:
|
||||
try:
|
||||
claims = jwt.decode(token, settings.JWT_SECRET_KEY)
|
||||
claims.validate()
|
||||
return AccessTokenPayload.model_validate(dict(claims))
|
||||
except JoseError, ValueError, TypeError:
|
||||
raise InvalidTokenError from None
|
||||
|
||||
|
||||
async def authenticate_token(token: str) -> AuthContext:
|
||||
try:
|
||||
payload = decode_access_token(token)
|
||||
except InvalidTokenError:
|
||||
raise InvalidTokenError from None
|
||||
|
||||
if await is_token_revoked(payload.jti):
|
||||
raise InvalidTokenError
|
||||
|
||||
user = await get_user_by_username(payload.sub)
|
||||
if user is None:
|
||||
raise InvalidTokenError
|
||||
|
||||
return AuthContext(token=token, user=user, payload=payload)
|
||||
|
||||
|
||||
async def get_auth_context(
|
||||
credentials: Annotated[
|
||||
HTTPAuthorizationCredentials | None,
|
||||
Depends(bearer_scheme),
|
||||
],
|
||||
) -> AuthContext:
|
||||
if credentials is None or credentials.scheme.lower() != "bearer":
|
||||
raise token_auth_error()
|
||||
|
||||
try:
|
||||
return await authenticate_token(credentials.credentials)
|
||||
except InvalidTokenError:
|
||||
raise token_auth_error() from None
|
||||
|
||||
|
||||
async def require_authenticated_user(
|
||||
auth_context: Annotated[AuthContext, Depends(get_auth_context)],
|
||||
) -> AuthUserDocument:
|
||||
return auth_context.user
|
||||
@@ -0,0 +1,15 @@
|
||||
from integrations.docker.compose import (
|
||||
find_docker_compose,
|
||||
get_all_services_in_dir,
|
||||
get_services,
|
||||
)
|
||||
from integrations.docker.editor import change_port
|
||||
from integrations.docker.models import DockerContainer
|
||||
|
||||
__all__ = [
|
||||
"DockerContainer",
|
||||
"find_docker_compose",
|
||||
"get_services",
|
||||
"get_all_services_in_dir",
|
||||
"change_port",
|
||||
]
|
||||
@@ -0,0 +1,98 @@
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from integrations.docker.models import DockerContainer
|
||||
from integrations.docker.ports import _get_service_port
|
||||
|
||||
COMPOSE_NAMES = (
|
||||
"compose.yml",
|
||||
"compose.yaml",
|
||||
"docker-compose.yml",
|
||||
"docker-compose.yaml",
|
||||
)
|
||||
|
||||
|
||||
def _get_visible_folders(directory: Path) -> list[Path]:
|
||||
return [
|
||||
p
|
||||
for p in Path(directory).iterdir()
|
||||
if p.is_dir() and not p.name.startswith(".")
|
||||
]
|
||||
|
||||
|
||||
def find_docker_compose(path: Path) -> tuple[dict, Path] | None:
|
||||
if not path or not path.is_dir():
|
||||
return None
|
||||
|
||||
for name in COMPOSE_NAMES:
|
||||
file = path / name
|
||||
if file.is_file():
|
||||
content = yaml.safe_load(file.read_text()) or {}
|
||||
return content, file
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_services(folder: Path) -> list[DockerContainer] | None:
|
||||
services: list[DockerContainer] = []
|
||||
compose_data = find_docker_compose(folder)
|
||||
if compose_data is None:
|
||||
return None
|
||||
|
||||
docker_compose, _ = compose_data
|
||||
compose_services = docker_compose.get("services") or {}
|
||||
if not compose_services:
|
||||
return None
|
||||
|
||||
for service_name, config in compose_services.items():
|
||||
mappings = config.get("ports") if isinstance(config, dict) else None
|
||||
if not isinstance(mappings, list) or not mappings:
|
||||
continue
|
||||
|
||||
ports = [
|
||||
p
|
||||
for p in (_get_service_port(mapping) for mapping in mappings)
|
||||
if p is not None
|
||||
]
|
||||
if not ports:
|
||||
continue
|
||||
|
||||
services.append({"name": service_name, "ports": ports, "folder": folder})
|
||||
|
||||
return services
|
||||
|
||||
|
||||
def get_all_services_in_dir(directory: str | Path) -> list[DockerContainer]:
|
||||
path = Path(directory)
|
||||
if not path.is_dir():
|
||||
return []
|
||||
|
||||
services: list[DockerContainer] = []
|
||||
for sub_folder in _get_visible_folders(path):
|
||||
sub_services = get_services(sub_folder)
|
||||
if sub_services:
|
||||
services.extend(sub_services)
|
||||
|
||||
return services
|
||||
|
||||
|
||||
async def run_docker_command(file: Path, command: list[str]) -> dict | None:
|
||||
if not command:
|
||||
return None
|
||||
|
||||
full_command = ["docker", "compose", "-f", f"{file}", *command]
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*full_command,
|
||||
cwd=str(file.parent),
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
stdout, stderr = await process.communicate()
|
||||
|
||||
if process.returncode == 0:
|
||||
return {"status": "success", "output": stdout.decode()}
|
||||
else:
|
||||
return {"status": "error", "reason": stderr.decode()}
|
||||
@@ -0,0 +1,60 @@
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from integrations.docker.compose import find_docker_compose, run_docker_command
|
||||
from integrations.docker.ports import _get_port_index, _replace_port_in_mapping
|
||||
|
||||
|
||||
def _backup_path_for(compose_path: Path) -> Path:
|
||||
if compose_path.suffix in {".yml", ".yaml"}:
|
||||
return compose_path.with_name(f"{compose_path.stem}.bak{compose_path.suffix}")
|
||||
return compose_path.with_name(f"{compose_path.name}.bak")
|
||||
|
||||
|
||||
async def change_port(
|
||||
service_name: str,
|
||||
folder: Path,
|
||||
port_from: int,
|
||||
port_to: int,
|
||||
) -> bool:
|
||||
compose_data = find_docker_compose(folder)
|
||||
if compose_data is None:
|
||||
return False
|
||||
|
||||
docker_compose, compose_path = compose_data
|
||||
compose_services = docker_compose.get("services") or {}
|
||||
if not compose_services:
|
||||
return False
|
||||
|
||||
service_config = compose_services.get(service_name)
|
||||
if not isinstance(service_config, dict):
|
||||
return False
|
||||
|
||||
mappings = service_config.get("ports")
|
||||
if not isinstance(mappings, list):
|
||||
return False
|
||||
|
||||
idx = _get_port_index(mappings, port_from)
|
||||
if idx is None:
|
||||
return False
|
||||
|
||||
mappings[idx] = _replace_port_in_mapping(mappings[idx], port_from, port_to)
|
||||
|
||||
backup_path = _backup_path_for(compose_path)
|
||||
original_text = compose_path.read_text()
|
||||
|
||||
backup_path.write_text(compose_path.read_text())
|
||||
compose_path.write_text(
|
||||
yaml.dump(docker_compose, default_flow_style=False, sort_keys=False)
|
||||
)
|
||||
|
||||
result = await run_docker_command(
|
||||
compose_path,
|
||||
["up", "-d", service_name],
|
||||
)
|
||||
|
||||
if result is None or result.get("status") != "success":
|
||||
compose_path.write_text(original_text)
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,8 @@
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
||||
|
||||
class DockerContainer(TypedDict):
|
||||
name: str
|
||||
ports: list[int]
|
||||
folder: Path
|
||||
@@ -0,0 +1,80 @@
|
||||
import socket
|
||||
from collections.abc import Mapping
|
||||
|
||||
|
||||
def get_random_port() -> int:
|
||||
# Binding to port 0 asks the OS for an available ephemeral port.
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
port = s.getsockname()[1]
|
||||
return port
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Could not find an available port: {e}") from e
|
||||
|
||||
|
||||
def _get_service_port(mapping: str | int | Mapping[str, object]) -> int | None:
|
||||
if isinstance(mapping, int):
|
||||
return mapping
|
||||
|
||||
if isinstance(mapping, Mapping):
|
||||
published = mapping.get("published")
|
||||
if published is None:
|
||||
return None
|
||||
try:
|
||||
return int(published)
|
||||
except TypeError, ValueError:
|
||||
return None
|
||||
|
||||
parts = str(mapping).split(":")
|
||||
match parts:
|
||||
case [""]:
|
||||
return None
|
||||
case [port] | [port, _]:
|
||||
try:
|
||||
return int(port)
|
||||
except ValueError:
|
||||
return None
|
||||
case ["localhost" | "127.0.0.1", _, _]:
|
||||
return None
|
||||
case [_, port, _]:
|
||||
try:
|
||||
return int(port)
|
||||
except ValueError:
|
||||
return None
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _replace_port_in_mapping(
|
||||
mapping: str | int | Mapping[str, object],
|
||||
port_from: int,
|
||||
port_to: int,
|
||||
) -> str | dict[str, object]:
|
||||
if isinstance(mapping, Mapping):
|
||||
updated = dict(mapping)
|
||||
if _get_service_port(updated) == port_from:
|
||||
updated["published"] = str(port_to)
|
||||
return updated
|
||||
|
||||
parts = str(mapping).split(":")
|
||||
match parts:
|
||||
case [_]:
|
||||
return str(port_to)
|
||||
case [_, container_port]:
|
||||
return f"{port_to}:{container_port}"
|
||||
case [ip, _, container_port]:
|
||||
return f"{ip}:{port_to}:{container_port}"
|
||||
case _:
|
||||
return str(port_to)
|
||||
|
||||
|
||||
def _get_port_index(
|
||||
mappings: list[str | int | Mapping[str, object]],
|
||||
port: int,
|
||||
) -> int | None:
|
||||
for i, mapping in enumerate(mappings):
|
||||
if _get_service_port(mapping) == port:
|
||||
return i
|
||||
return None
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
import sys
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
from api.auth.router import router as auth_router
|
||||
from api.editor.router import router as editor_router
|
||||
from api.proxy.router import router as proxy_router
|
||||
from core.config import settings
|
||||
from core.mongodb import close_mongodb, init_mongodb
|
||||
from core.redis import close_redis, init_redis
|
||||
from core.security import hash_password
|
||||
from fastapi import FastAPI
|
||||
from integrations.docker import get_all_services_in_dir
|
||||
from repositories.auth import ensure_auth_indexes, ensure_single_account
|
||||
from repositories.examples import ensure_example_indexes
|
||||
from repositories.proxy import ensure_proxy_indexes
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI):
|
||||
await init_mongodb()
|
||||
await init_redis()
|
||||
await ensure_auth_indexes()
|
||||
await ensure_single_account(
|
||||
username=settings.AUTH_USERNAME,
|
||||
password_hash=hash_password(settings.AUTH_PASSWORD),
|
||||
)
|
||||
await ensure_proxy_indexes()
|
||||
await ensure_example_indexes()
|
||||
|
||||
home_dir = str(Path.home())
|
||||
app.state.services = get_all_services_in_dir(home_dir) or []
|
||||
if not app.state.services:
|
||||
warning = (
|
||||
"\n" + "!" * 60 + "\n"
|
||||
" WARNING: No services found in home directory!\n"
|
||||
f" Scanned: {home_dir}\n" + "!" * 60 + "\n"
|
||||
)
|
||||
print(warning, file=sys.stderr)
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await close_redis()
|
||||
await close_mongodb()
|
||||
|
||||
|
||||
DEBUG = os.getenv("DEBUG", "false").lower() == "true"
|
||||
|
||||
app = FastAPI(
|
||||
title="thePlugAPI",
|
||||
lifespan=lifespan,
|
||||
docs_url="/docs" if DEBUG else None,
|
||||
redoc_url="/redoc" if DEBUG else None,
|
||||
openapi_url="/openapi.json" if DEBUG else None,
|
||||
)
|
||||
|
||||
app.include_router(auth_router)
|
||||
app.include_router(editor_router)
|
||||
app.include_router(proxy_router)
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"Hello": "World!"}
|
||||
@@ -0,0 +1,43 @@
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class AuthUserDocument(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
username: str = Field(
|
||||
...,
|
||||
description="Unique username for the single account",
|
||||
)
|
||||
password_hash: str = Field(..., description="Hashed password")
|
||||
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 LoginInput(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
username: str = Field(..., description="Username")
|
||||
password: str = Field(..., description="Plain text password")
|
||||
|
||||
|
||||
class TokenResponse(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
access_token: str = Field(..., description="JWT access token")
|
||||
token_type: str = Field(default="bearer", description="Token type")
|
||||
|
||||
|
||||
class AccessTokenPayload(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
sub: str = Field(..., description="Subject username")
|
||||
jti: str = Field(..., description="Unique token identifier")
|
||||
exp: int = Field(..., description="Token expiration as unix timestamp")
|
||||
@@ -0,0 +1,28 @@
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class AddonExampleDocument(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str = Field(..., description="Stable example identifier")
|
||||
name: str = Field(..., description="Display name")
|
||||
description: str = Field(default="", description="Short explanation")
|
||||
category: str = Field(default="custom", description="Example category")
|
||||
code: str = Field(..., description="Addon source code")
|
||||
source_url: str | None = Field(default=None, description="Reference URL")
|
||||
built_in: bool = Field(default=False, description="Provided by thePlug")
|
||||
created_at: datetime | None = None
|
||||
updated_at: datetime | None = None
|
||||
|
||||
|
||||
class AddonExampleCreateInput(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
id: str = Field(..., description="Stable example identifier")
|
||||
name: str = Field(..., description="Display name")
|
||||
description: str = Field(default="", description="Short explanation")
|
||||
category: str = Field(default="custom", description="Example category")
|
||||
code: str = Field(..., description="Addon source code")
|
||||
source_url: str | None = Field(default=None, description="Reference URL")
|
||||
@@ -0,0 +1,115 @@
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from models.service import ServiceEndpointDocument
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class HealthCheckDocument(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
enabled: bool = Field(default=True, description="Enable active health checks")
|
||||
protocol: str | None = Field(
|
||||
default=None,
|
||||
description="Override protocol used by the health check",
|
||||
)
|
||||
path: str = Field(default="/", description="HTTP path to check")
|
||||
method: str = Field(default="GET", description="HTTP method to use")
|
||||
expected_status: int | None = Field(
|
||||
default=None,
|
||||
description="Expected HTTP status code",
|
||||
)
|
||||
udp_payload: str = Field(default="", description="UDP payload to send")
|
||||
expect_response: bool = Field(
|
||||
default=True,
|
||||
description="Wait for a response when the protocol supports it",
|
||||
)
|
||||
verify_tls: bool = Field(
|
||||
default=False,
|
||||
description="Verify TLS certificates during health checks",
|
||||
)
|
||||
timeout_seconds: float = Field(
|
||||
default=2.0,
|
||||
gt=0,
|
||||
le=30,
|
||||
description="Per-target health check timeout",
|
||||
)
|
||||
|
||||
|
||||
class ProxyDocument(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")
|
||||
mode: str = Field(default="reverse", description="mitmproxy proxy mode")
|
||||
mode_spec: str | None = Field(
|
||||
default=None,
|
||||
description="Raw mitmproxy mode specification. Overrides mode/protocol target.",
|
||||
)
|
||||
protocol: str = Field(default="http", description="mitmproxy reverse protocol")
|
||||
listen_host: str = Field(default="0.0.0.0", description="Proxy listen host")
|
||||
mitm_options: dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="mitmdump options passed through --set",
|
||||
)
|
||||
healthcheck: HealthCheckDocument = Field(default_factory=HealthCheckDocument)
|
||||
status: str = Field(default="stopped", description="Runtime proxy status")
|
||||
pid: int | None = Field(default=None, description="mitmdump process id")
|
||||
script_path: str | None = Field(default=None, description="Generated addon path")
|
||||
config_path: str | None = Field(default=None, description="Generated config path")
|
||||
log_path: str | None = Field(default=None, description="mitmdump log path")
|
||||
last_error: str | None = Field(default=None, description="Last runtime error")
|
||||
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")
|
||||
mode: str = Field(default="reverse", description="mitmproxy proxy mode")
|
||||
mode_spec: str | None = Field(
|
||||
default=None,
|
||||
description="Raw mitmproxy mode specification. Overrides mode/protocol target.",
|
||||
)
|
||||
protocol: str = Field(default="http", description="mitmproxy reverse protocol")
|
||||
listen_host: str = Field(default="0.0.0.0", description="Proxy listen host")
|
||||
mitm_options: dict[str, Any] = Field(default_factory=dict)
|
||||
healthcheck: HealthCheckDocument = Field(default_factory=HealthCheckDocument)
|
||||
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")
|
||||
mode: str | None = Field(default=None, description="mitmproxy proxy mode")
|
||||
mode_spec: str | None = Field(
|
||||
default=None,
|
||||
description="Raw mitmproxy mode specification. Overrides mode/protocol target.",
|
||||
)
|
||||
protocol: str | None = Field(default=None, description="mitmproxy reverse protocol")
|
||||
listen_host: str | None = Field(default=None, description="Proxy listen host")
|
||||
mitm_options: dict[str, Any] | None = Field(default=None)
|
||||
healthcheck: HealthCheckDocument | None = Field(default=None)
|
||||
service: ServiceEndpointDocument | None = Field(
|
||||
default=None,
|
||||
description="Associated upstream service",
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class ServiceEndpointDocument(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
host: str = Field(default="127.0.0.1", description="Upstream service host")
|
||||
port: int = Field(..., description="Original upstream service port")
|
||||
name: str | None = Field(default=None, description="Docker Compose service name")
|
||||
folder: str | None = Field(
|
||||
default=None,
|
||||
description="Docker Compose project folder",
|
||||
)
|
||||
upstream_port: int | None = Field(
|
||||
default=None,
|
||||
description="Current port used by mitmdump as upstream",
|
||||
)
|
||||
@@ -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
|
||||
@@ -0,0 +1,323 @@
|
||||
import asyncio
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
WORKSPACE_ROOT = Path(tempfile.gettempdir()) / "theplug-lsp"
|
||||
CONFIG_FILE = "pyrightconfig.json"
|
||||
ENTRY_FILE = "proxy_addon.py"
|
||||
|
||||
STUB_FILES = {
|
||||
"typings/mitmproxy/__init__.pyi": """from . import (
|
||||
command,
|
||||
ctx,
|
||||
dns,
|
||||
http,
|
||||
tcp,
|
||||
tls,
|
||||
udp,
|
||||
websocket,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"command",
|
||||
"ctx",
|
||||
"dns",
|
||||
"http",
|
||||
"tcp",
|
||||
"tls",
|
||||
"udp",
|
||||
"websocket",
|
||||
]
|
||||
""",
|
||||
"typings/mitmproxy/ctx.pyi": """class _Logger:
|
||||
def debug(self, message: object) -> None: ...
|
||||
def info(self, message: object) -> None: ...
|
||||
def warn(self, message: object) -> None: ...
|
||||
def error(self, message: object) -> None: ...
|
||||
def alert(self, message: object) -> None: ...
|
||||
|
||||
|
||||
class _Options:
|
||||
def __getattr__(self, name: str) -> object: ...
|
||||
|
||||
|
||||
log: _Logger
|
||||
options: _Options
|
||||
""",
|
||||
"typings/mitmproxy/http.pyi": """from typing import Any
|
||||
|
||||
|
||||
class Headers(dict[str, str]):
|
||||
...
|
||||
|
||||
|
||||
class Request:
|
||||
host: str
|
||||
path: str
|
||||
headers: Headers
|
||||
query: dict[str, str]
|
||||
content: bytes
|
||||
|
||||
|
||||
class Response:
|
||||
headers: Headers
|
||||
status_code: int
|
||||
content: bytes
|
||||
|
||||
@classmethod
|
||||
def make(
|
||||
cls,
|
||||
status_code: int,
|
||||
content: bytes | str = b"",
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> "Response": ...
|
||||
|
||||
|
||||
class WebSocketMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
|
||||
|
||||
class WebSocketData:
|
||||
messages: list[WebSocketMessage]
|
||||
|
||||
|
||||
class HTTPFlow:
|
||||
request: Request
|
||||
response: Response | None
|
||||
websocket: WebSocketData | None
|
||||
metadata: dict[str, Any]
|
||||
""",
|
||||
"typings/mitmproxy/tcp.pyi": """class TCPMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
|
||||
|
||||
class TCPFlow:
|
||||
messages: list[TCPMessage]
|
||||
""",
|
||||
"typings/mitmproxy/dns.pyi": """class Question:
|
||||
name: str
|
||||
type: int
|
||||
|
||||
|
||||
class ResourceRecord:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
type: int,
|
||||
class_: int,
|
||||
ttl: int,
|
||||
data: bytes,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class Message:
|
||||
question: Question | None
|
||||
|
||||
def succeed(self, records: list[ResourceRecord]) -> "Message": ...
|
||||
|
||||
|
||||
class _Types:
|
||||
A: int
|
||||
|
||||
|
||||
class _Classes:
|
||||
IN: int
|
||||
|
||||
|
||||
types: _Types
|
||||
classes: _Classes
|
||||
|
||||
|
||||
class DNSFlow:
|
||||
request: Message
|
||||
response: Message | None
|
||||
""",
|
||||
"typings/mitmproxy/command.pyi": """from collections.abc import Callable
|
||||
from typing import TypeVar
|
||||
|
||||
F = TypeVar("F", bound=Callable[..., object])
|
||||
|
||||
|
||||
def command(name: str) -> Callable[[F], F]: ...
|
||||
""",
|
||||
"typings/mitmproxy/websocket.pyi": """class WebSocketMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
""",
|
||||
"typings/mitmproxy/tls.pyi": """class ClientHelloData: ...
|
||||
""",
|
||||
"typings/mitmproxy/udp.pyi": """class UDPMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
""",
|
||||
}
|
||||
|
||||
|
||||
class LspFrameBuffer:
|
||||
def __init__(self) -> None:
|
||||
self.buffer = bytearray()
|
||||
|
||||
def feed(self, chunk: bytes) -> list[str]:
|
||||
self.buffer.extend(chunk)
|
||||
messages: list[str] = []
|
||||
while True:
|
||||
header_end = self.buffer.find(b"\r\n\r\n")
|
||||
if header_end < 0:
|
||||
return messages
|
||||
|
||||
header = self.buffer[:header_end].decode("ascii", errors="ignore")
|
||||
size = content_length(header)
|
||||
if size is None:
|
||||
del self.buffer[: header_end + 4]
|
||||
continue
|
||||
|
||||
frame_end = header_end + 4 + size
|
||||
if len(self.buffer) < frame_end:
|
||||
return messages
|
||||
|
||||
body = bytes(self.buffer[header_end + 4 : frame_end])
|
||||
del self.buffer[:frame_end]
|
||||
messages.append(body.decode("utf-8"))
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PythonEditorSession:
|
||||
root_path: Path
|
||||
file_path: Path
|
||||
|
||||
@property
|
||||
def root_uri(self) -> str:
|
||||
return self.root_path.as_uri()
|
||||
|
||||
@property
|
||||
def file_uri(self) -> str:
|
||||
return self.file_path.as_uri()
|
||||
|
||||
|
||||
def content_length(header: str) -> int | None:
|
||||
for line in header.split("\r\n"):
|
||||
if ":" not in line:
|
||||
continue
|
||||
name, value = line.split(":", 1)
|
||||
if name.lower() == "content-length":
|
||||
return int(value.strip())
|
||||
return None
|
||||
|
||||
|
||||
def encode_message(message: str) -> bytes:
|
||||
payload = message.encode("utf-8")
|
||||
header = f"Content-Length: {len(payload)}\r\n\r\n".encode("ascii")
|
||||
return header + payload
|
||||
|
||||
|
||||
def lsp_binary() -> str | None:
|
||||
return shutil.which("basedpyright-langserver")
|
||||
|
||||
|
||||
def get_session() -> PythonEditorSession:
|
||||
root_path = WORKSPACE_ROOT
|
||||
file_path = root_path / ENTRY_FILE
|
||||
|
||||
root_path.mkdir(parents=True, exist_ok=True)
|
||||
file_path.touch(exist_ok=True)
|
||||
|
||||
config = {
|
||||
"include": [ENTRY_FILE],
|
||||
"stubPath": "typings",
|
||||
"typeCheckingMode": "basic",
|
||||
"pythonVersion": "3.14",
|
||||
"reportMissingTypeStubs": False,
|
||||
"reportMissingImports": "warning",
|
||||
"useLibraryCodeForTypes": True,
|
||||
}
|
||||
(root_path / CONFIG_FILE).write_text(
|
||||
json.dumps(config, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
for relative_path, content in STUB_FILES.items():
|
||||
path = root_path / relative_path
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
|
||||
return PythonEditorSession(root_path=root_path, file_path=file_path)
|
||||
|
||||
|
||||
async def open_process(session: PythonEditorSession) -> asyncio.subprocess.Process:
|
||||
binary = lsp_binary()
|
||||
if binary is None:
|
||||
raise RuntimeError("basedpyright-langserver is not installed")
|
||||
|
||||
return await asyncio.create_subprocess_exec(
|
||||
binary,
|
||||
"--stdio",
|
||||
cwd=session.root_path,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
|
||||
async def forward_client(websocket, writer: asyncio.StreamWriter) -> None:
|
||||
while True:
|
||||
message = await websocket.receive_text()
|
||||
writer.write(encode_message(message))
|
||||
await writer.drain()
|
||||
|
||||
|
||||
async def forward_server(reader: asyncio.StreamReader, websocket) -> None:
|
||||
frames = LspFrameBuffer()
|
||||
while True:
|
||||
chunk = await reader.read(65536)
|
||||
if not chunk:
|
||||
return
|
||||
for message in frames.feed(chunk):
|
||||
await websocket.send_text(message)
|
||||
|
||||
|
||||
async def drain_errors(reader: asyncio.StreamReader) -> None:
|
||||
while True:
|
||||
chunk = await reader.read(4096)
|
||||
if not chunk:
|
||||
return
|
||||
|
||||
|
||||
async def bridge(websocket, session: PythonEditorSession) -> None:
|
||||
process = await open_process(session)
|
||||
stdin = process.stdin
|
||||
stdout = process.stdout
|
||||
stderr = process.stderr
|
||||
|
||||
if stdin is None or stdout is None or stderr is None:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
raise RuntimeError("basedpyright-langserver did not expose stdio pipes")
|
||||
|
||||
send_task = asyncio.create_task(forward_client(websocket, stdin))
|
||||
receive_task = asyncio.create_task(forward_server(stdout, websocket))
|
||||
error_task = asyncio.create_task(drain_errors(stderr))
|
||||
|
||||
try:
|
||||
done, pending = await asyncio.wait(
|
||||
{send_task, receive_task, error_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
for task in done:
|
||||
task.result()
|
||||
finally:
|
||||
if process.returncode is None:
|
||||
stdin.close()
|
||||
await stdin.wait_closed()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=1)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
@@ -0,0 +1,134 @@
|
||||
from models.example import AddonExampleDocument
|
||||
|
||||
DOCS_EXAMPLES_URL = "https://docs.mitmproxy.org/stable/addons/examples/"
|
||||
|
||||
BUILTIN_EXAMPLES = [
|
||||
AddonExampleDocument(
|
||||
id="http-add-response-header",
|
||||
name="HTTP add response header",
|
||||
description=(
|
||||
"Injects a fixed response header on every HTTP response so you can "
|
||||
"verify that traffic is passing through the proxy response path."
|
||||
),
|
||||
category="http",
|
||||
source_url=DOCS_EXAMPLES_URL,
|
||||
built_in=True,
|
||||
code="""from mitmproxy import http
|
||||
|
||||
|
||||
def response(flow: http.HTTPFlow) -> None:
|
||||
flow.response.headers["x-theplug"] = "enabled"
|
||||
""",
|
||||
),
|
||||
AddonExampleDocument(
|
||||
id="http-modify-query",
|
||||
name="HTTP modify query",
|
||||
description=(
|
||||
"Adds or overwrites a query parameter before each request is "
|
||||
"forwarded upstream. Useful for forced debug or feature flags."
|
||||
),
|
||||
category="http",
|
||||
source_url=DOCS_EXAMPLES_URL,
|
||||
built_in=True,
|
||||
code="""from mitmproxy import http
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
flow.request.query["debug"] = "1"
|
||||
""",
|
||||
),
|
||||
AddonExampleDocument(
|
||||
id="http-reply-from-proxy",
|
||||
name="HTTP reply from proxy",
|
||||
description=(
|
||||
"Serves a local HTTP response for a selected path without touching "
|
||||
"the upstream service. Useful for probes, maintenance, or stubs."
|
||||
),
|
||||
category="http",
|
||||
source_url=DOCS_EXAMPLES_URL,
|
||||
built_in=True,
|
||||
code="""from mitmproxy import http
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
if flow.request.path == "/healthz":
|
||||
flow.response = http.Response.make(
|
||||
200,
|
||||
b"ok",
|
||||
{"content-type": "text/plain"},
|
||||
)
|
||||
""",
|
||||
),
|
||||
AddonExampleDocument(
|
||||
id="websocket-rewrite",
|
||||
name="WebSocket rewrite",
|
||||
description=(
|
||||
"Rewrites matching bytes inside live WebSocket frames to alter "
|
||||
"runtime flags or payload markers in transit."
|
||||
),
|
||||
category="websocket",
|
||||
source_url=DOCS_EXAMPLES_URL,
|
||||
built_in=True,
|
||||
code="""from mitmproxy import http
|
||||
|
||||
|
||||
def websocket_message(flow: http.HTTPFlow) -> None:
|
||||
if flow.websocket is None:
|
||||
return
|
||||
message = flow.websocket.messages[-1]
|
||||
message.content = message.content.replace(b"debug=false", b"debug=true")
|
||||
""",
|
||||
),
|
||||
AddonExampleDocument(
|
||||
id="tcp-log-messages",
|
||||
name="TCP log messages",
|
||||
description=(
|
||||
"Logs raw TCP message direction and size for each frame so you can "
|
||||
"observe binary traffic without altering it."
|
||||
),
|
||||
category="tcp",
|
||||
source_url=DOCS_EXAMPLES_URL,
|
||||
built_in=True,
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import tcp
|
||||
|
||||
|
||||
def tcp_message(flow: tcp.TCPFlow) -> None:
|
||||
message = flow.messages[-1]
|
||||
direction = "client" if message.from_client else "server"
|
||||
logging.info("%s sent %d bytes", direction, len(message.content))
|
||||
""",
|
||||
),
|
||||
AddonExampleDocument(
|
||||
id="dns-spoof-a-record",
|
||||
name="DNS spoof A record",
|
||||
description=(
|
||||
"Overrides one DNS A record while forwarding every other query "
|
||||
"unchanged. Useful for selective service redirection."
|
||||
),
|
||||
category="dns",
|
||||
source_url=DOCS_EXAMPLES_URL,
|
||||
built_in=True,
|
||||
code="""import ipaddress
|
||||
|
||||
from mitmproxy import dns
|
||||
|
||||
|
||||
def dns_request(flow: dns.DNSFlow) -> None:
|
||||
question = flow.request.question
|
||||
if question and question.name == "target.local" and question.type == dns.types.A:
|
||||
flow.response = flow.request.succeed(
|
||||
[
|
||||
dns.ResourceRecord(
|
||||
name="target.local",
|
||||
type=dns.types.A,
|
||||
class_=dns.classes.IN,
|
||||
ttl=60,
|
||||
data=ipaddress.ip_address("127.0.0.1").packed,
|
||||
)
|
||||
]
|
||||
)
|
||||
""",
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,185 @@
|
||||
import asyncio
|
||||
import ssl
|
||||
from dataclasses import dataclass
|
||||
|
||||
from models.proxy import HealthCheckDocument, ProxyDocument
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Target:
|
||||
name: str
|
||||
host: str
|
||||
port: int
|
||||
protocol: str
|
||||
|
||||
|
||||
async def check_proxy(proxy: ProxyDocument) -> dict[str, object]:
|
||||
if not proxy.healthcheck.enabled:
|
||||
return {"enabled": False, "overall": "disabled", "checks": []}
|
||||
|
||||
upstream_port = proxy.service.upstream_port or proxy.service.port
|
||||
checks = await asyncio.gather(
|
||||
check_target(
|
||||
Target(
|
||||
name="upstream",
|
||||
host=proxy.service.host,
|
||||
port=upstream_port,
|
||||
protocol=health_protocol(proxy),
|
||||
),
|
||||
proxy.healthcheck,
|
||||
),
|
||||
check_target(
|
||||
Target(
|
||||
name="proxy",
|
||||
host=listen_host_for_check(proxy.listen_host),
|
||||
port=proxy.proxy_port,
|
||||
protocol=health_protocol(proxy),
|
||||
),
|
||||
proxy.healthcheck,
|
||||
),
|
||||
)
|
||||
overall = "healthy" if all(check["healthy"] for check in checks) else "unhealthy"
|
||||
return {"enabled": True, "overall": overall, "checks": checks}
|
||||
|
||||
|
||||
def health_protocol(proxy: ProxyDocument) -> str:
|
||||
return proxy.healthcheck.protocol or proxy.protocol
|
||||
|
||||
|
||||
def listen_host_for_check(host: str) -> str:
|
||||
if host in {"", "0.0.0.0", "::"}:
|
||||
return "127.0.0.1"
|
||||
return host
|
||||
|
||||
|
||||
async def check_target(
|
||||
target: Target,
|
||||
config: HealthCheckDocument,
|
||||
) -> dict[str, object]:
|
||||
try:
|
||||
match target.protocol:
|
||||
case "http":
|
||||
return await check_http(target, config, use_tls=False)
|
||||
case "https":
|
||||
return await check_http(target, config, use_tls=True)
|
||||
case "tcp":
|
||||
return await check_tcp(target, config, use_tls=False)
|
||||
case "tls":
|
||||
return await check_tcp(target, config, use_tls=True)
|
||||
case "udp" | "dtls" | "quic" | "dns" | "http3":
|
||||
return await check_udp(target, config)
|
||||
case _:
|
||||
return result(target, False, f"Unsupported protocol: {target.protocol}")
|
||||
except TimeoutError:
|
||||
return result(target, False, "Timed out")
|
||||
except OSError as error:
|
||||
return result(target, False, str(error))
|
||||
|
||||
|
||||
async def check_http(
|
||||
target: Target,
|
||||
config: HealthCheckDocument,
|
||||
*,
|
||||
use_tls: bool,
|
||||
) -> dict[str, object]:
|
||||
reader, writer = await open_connection(target, config, use_tls=use_tls)
|
||||
try:
|
||||
path = config.path if config.path.startswith("/") else f"/{config.path}"
|
||||
request = (
|
||||
f"{config.method} {path} HTTP/1.1\r\n"
|
||||
f"Host: {target.host}\r\n"
|
||||
"Connection: close\r\n\r\n"
|
||||
)
|
||||
writer.write(request.encode("ascii"))
|
||||
await writer.drain()
|
||||
line = await asyncio.wait_for(reader.readline(), config.timeout_seconds)
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
|
||||
parts = line.decode("iso-8859-1", errors="replace").split()
|
||||
status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else None
|
||||
expected = config.expected_status
|
||||
healthy = status == expected if expected is not None else status is not None
|
||||
return result(target, healthy, None, {"status": status})
|
||||
|
||||
|
||||
async def check_tcp(
|
||||
target: Target,
|
||||
config: HealthCheckDocument,
|
||||
*,
|
||||
use_tls: bool,
|
||||
) -> dict[str, object]:
|
||||
_reader, writer = await open_connection(target, config, use_tls=use_tls)
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
return result(target, True)
|
||||
|
||||
|
||||
async def check_udp(target: Target, config: HealthCheckDocument) -> dict[str, object]:
|
||||
loop = asyncio.get_running_loop()
|
||||
transport, protocol = await loop.create_datagram_endpoint(
|
||||
lambda: DatagramProbe(config.expect_response),
|
||||
remote_addr=(target.host, target.port),
|
||||
)
|
||||
try:
|
||||
payload = config.udp_payload.encode()
|
||||
transport.sendto(payload)
|
||||
if not config.expect_response:
|
||||
return result(target, True, None, {"mode": "send-only"})
|
||||
data = await asyncio.wait_for(protocol.response, config.timeout_seconds)
|
||||
return result(target, True, None, {"bytes": len(data)})
|
||||
finally:
|
||||
transport.close()
|
||||
|
||||
|
||||
async def open_connection(
|
||||
target: Target,
|
||||
config: HealthCheckDocument,
|
||||
*,
|
||||
use_tls: bool,
|
||||
):
|
||||
if use_tls and config.verify_tls:
|
||||
context = ssl.create_default_context()
|
||||
elif use_tls:
|
||||
context = ssl.create_default_context()
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
else:
|
||||
context = None
|
||||
return await asyncio.wait_for(
|
||||
asyncio.open_connection(target.host, target.port, ssl=context),
|
||||
config.timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
class DatagramProbe(asyncio.DatagramProtocol):
|
||||
def __init__(self, expect_response: bool):
|
||||
self.response = asyncio.get_running_loop().create_future()
|
||||
if not expect_response:
|
||||
self.response.set_result(b"")
|
||||
|
||||
def datagram_received(self, data, _addr):
|
||||
if not self.response.done():
|
||||
self.response.set_result(data)
|
||||
|
||||
def error_received(self, exc):
|
||||
if not self.response.done():
|
||||
self.response.set_exception(exc)
|
||||
|
||||
|
||||
def result(
|
||||
target: Target,
|
||||
healthy: bool,
|
||||
error: str | None = None,
|
||||
details: dict[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"target": target.name,
|
||||
"host": target.host,
|
||||
"port": target.port,
|
||||
"protocol": target.protocol,
|
||||
"healthy": healthy,
|
||||
"error": error,
|
||||
"details": details or {},
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import signal
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from integrations import docker
|
||||
from integrations.docker.ports import get_random_port
|
||||
from models.proxy import ProxyDocument
|
||||
from services.mitmproxy.options import MitmOptionError, validate_options
|
||||
|
||||
|
||||
class MitmdumpError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class MitmdumpManager:
|
||||
def __init__(self, work_dir: str = "/tmp/theplug", command: str = "mitmdump"):
|
||||
self.work_dir = Path(work_dir)
|
||||
self.command = command
|
||||
self.processes: dict[int, asyncio.subprocess.Process] = {}
|
||||
self.locks: dict[int, asyncio.Lock] = {}
|
||||
|
||||
async def start(self, proxy: ProxyDocument) -> dict[str, object]:
|
||||
async with self.lock_for(proxy.proxy_port):
|
||||
return await self.start_unlocked(proxy)
|
||||
|
||||
async def stop(self, proxy: ProxyDocument) -> dict[str, object]:
|
||||
async with self.lock_for(proxy.proxy_port):
|
||||
return await self.stop_unlocked(proxy)
|
||||
|
||||
async def start_unlocked(self, proxy: ProxyDocument) -> dict[str, object]:
|
||||
if await self.status(proxy) == "running":
|
||||
raise MitmdumpError("Proxy is already running")
|
||||
|
||||
command_path = shutil.which(self.command)
|
||||
if command_path is None:
|
||||
raise MitmdumpError("mitmdump executable not found in PATH")
|
||||
|
||||
try:
|
||||
await validate_options(proxy.mitm_options)
|
||||
except MitmOptionError as error:
|
||||
raise MitmdumpError(str(error)) from error
|
||||
upstream_port = await self.prepare_upstream(proxy)
|
||||
script_path = self.write_script(proxy)
|
||||
config_path = self.write_config(proxy)
|
||||
log_path = self.log_path(proxy)
|
||||
command = self.build_command(
|
||||
command_path,
|
||||
proxy,
|
||||
upstream_port,
|
||||
script_path,
|
||||
config_path,
|
||||
)
|
||||
|
||||
try:
|
||||
process = await self.spawn(command, log_path)
|
||||
except Exception as error:
|
||||
await self.restore_upstream(proxy, upstream_port)
|
||||
raise MitmdumpError(str(error)) from error
|
||||
|
||||
failed = await self.exited_early(process)
|
||||
if failed:
|
||||
await self.restore_upstream(proxy, upstream_port)
|
||||
reason = self.read_log_tail(log_path) or "mitmdump exited during startup"
|
||||
raise MitmdumpError(reason)
|
||||
|
||||
self.processes[proxy.proxy_port] = process
|
||||
return {
|
||||
"status": "running",
|
||||
"pid": process.pid,
|
||||
"script_path": str(script_path),
|
||||
"config_path": str(config_path),
|
||||
"log_path": str(log_path),
|
||||
"service.upstream_port": upstream_port,
|
||||
"last_error": None,
|
||||
}
|
||||
|
||||
async def stop_unlocked(self, proxy: ProxyDocument) -> dict[str, object]:
|
||||
process = self.processes.pop(proxy.proxy_port, None)
|
||||
if process is not None and process.returncode is None:
|
||||
await self.stop_process(process)
|
||||
elif proxy.pid is not None:
|
||||
await self.stop_pid(proxy.pid)
|
||||
|
||||
await self.restore_upstream(proxy, proxy.service.upstream_port)
|
||||
return {
|
||||
"status": "stopped",
|
||||
"pid": None,
|
||||
"service.upstream_port": None,
|
||||
"last_error": None,
|
||||
}
|
||||
|
||||
async def status(self, proxy: ProxyDocument) -> str:
|
||||
process = self.processes.get(proxy.proxy_port)
|
||||
if process is not None:
|
||||
if process.returncode is None:
|
||||
return "running"
|
||||
self.processes.pop(proxy.proxy_port, None)
|
||||
return "stopped"
|
||||
|
||||
if proxy.pid is None:
|
||||
return "stopped" if proxy.status == "running" else proxy.status
|
||||
|
||||
try:
|
||||
os.kill(proxy.pid, 0)
|
||||
except ProcessLookupError:
|
||||
return "stopped"
|
||||
except PermissionError:
|
||||
return "running"
|
||||
return "running"
|
||||
|
||||
async def prepare_upstream(self, proxy: ProxyDocument) -> int | None:
|
||||
if proxy.mode_spec is not None or proxy.mode != "reverse":
|
||||
return None
|
||||
if proxy.service.upstream_port is not None:
|
||||
return proxy.service.upstream_port
|
||||
|
||||
if proxy.service.name is None or proxy.service.folder is None:
|
||||
return proxy.service.port
|
||||
|
||||
upstream_port = get_random_port()
|
||||
changed = await docker.change_port(
|
||||
proxy.service.name,
|
||||
Path(proxy.service.folder),
|
||||
proxy.service.port,
|
||||
upstream_port,
|
||||
)
|
||||
if not changed:
|
||||
raise MitmdumpError("Unable to move Docker service port")
|
||||
return upstream_port
|
||||
|
||||
async def restore_upstream(
|
||||
self,
|
||||
proxy: ProxyDocument,
|
||||
upstream_port: int | None,
|
||||
) -> None:
|
||||
if upstream_port is None:
|
||||
return
|
||||
if proxy.service.name is None or proxy.service.folder is None:
|
||||
return
|
||||
if upstream_port == proxy.service.port:
|
||||
return
|
||||
|
||||
await docker.change_port(
|
||||
proxy.service.name,
|
||||
Path(proxy.service.folder),
|
||||
upstream_port,
|
||||
proxy.service.port,
|
||||
)
|
||||
|
||||
def build_command(
|
||||
self,
|
||||
command_path: str,
|
||||
proxy: ProxyDocument,
|
||||
upstream_port: int | None,
|
||||
script_path: Path,
|
||||
config_path: Path,
|
||||
) -> list[str]:
|
||||
command = [
|
||||
command_path,
|
||||
"--mode",
|
||||
self.mode_for(proxy, upstream_port),
|
||||
"--listen-host",
|
||||
proxy.listen_host,
|
||||
"--set",
|
||||
f"confdir={config_path.parent}",
|
||||
"-s",
|
||||
str(script_path),
|
||||
]
|
||||
if self.needs_listen_port(proxy):
|
||||
command.extend(["--listen-port", str(proxy.proxy_port)])
|
||||
return command
|
||||
|
||||
def mode_for(self, proxy: ProxyDocument, upstream_port: int | None) -> str:
|
||||
if proxy.mode_spec:
|
||||
return proxy.mode_spec
|
||||
|
||||
match proxy.mode:
|
||||
case "reverse":
|
||||
port = upstream_port or proxy.service.port
|
||||
return (
|
||||
f"reverse:{proxy.protocol}://{proxy.service.host}:{port}"
|
||||
f"@{proxy.proxy_port}"
|
||||
)
|
||||
case "upstream":
|
||||
return (
|
||||
f"upstream:{proxy.protocol}://{proxy.service.host}:"
|
||||
f"{proxy.service.port}@{proxy.proxy_port}"
|
||||
)
|
||||
case "wireguard":
|
||||
return f"wireguard@{proxy.proxy_port}"
|
||||
case _:
|
||||
return proxy.mode
|
||||
|
||||
def needs_listen_port(self, proxy: ProxyDocument) -> bool:
|
||||
if proxy.mode_spec is not None:
|
||||
return False
|
||||
return proxy.mode in {"regular", "socks5", "dns", "transparent"}
|
||||
|
||||
def write_script(self, proxy: ProxyDocument) -> Path:
|
||||
path = self.work_dir / "addons" / f"{proxy.proxy_port}.py"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(proxy.code, encoding="utf-8")
|
||||
return path
|
||||
|
||||
def write_config(self, proxy: ProxyDocument) -> Path:
|
||||
path = self.work_dir / "conf" / str(proxy.proxy_port) / "config.yaml"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
options = {"block_global": False, **proxy.mitm_options}
|
||||
path.write_text(yaml.safe_dump(options, sort_keys=True), encoding="utf-8")
|
||||
return path
|
||||
|
||||
def log_path(self, proxy: ProxyDocument) -> Path:
|
||||
path = self.work_dir / "logs" / f"{proxy.proxy_port}.log"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
async def spawn(self, command: list[str], log_path: Path):
|
||||
with log_path.open("ab") as log_file:
|
||||
return await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
stdout=log_file,
|
||||
stderr=log_file,
|
||||
start_new_session=True,
|
||||
)
|
||||
|
||||
async def stop_process(self, process) -> None:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=5)
|
||||
except TimeoutError:
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
await process.wait()
|
||||
|
||||
async def exited_early(self, process) -> bool:
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=0.5)
|
||||
except TimeoutError:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def stop_pid(self, pid: int) -> None:
|
||||
try:
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0.2)
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
|
||||
try:
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
return
|
||||
|
||||
def read_log_tail(self, path: Path, limit: int = 4000) -> str:
|
||||
if not path.exists():
|
||||
return ""
|
||||
content = path.read_text(encoding="utf-8", errors="replace")
|
||||
return content[-limit:].strip()
|
||||
|
||||
def lock_for(self, proxy_port: int) -> asyncio.Lock:
|
||||
lock = self.locks.get(proxy_port)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
self.locks[proxy_port] = lock
|
||||
return lock
|
||||
|
||||
|
||||
manager = MitmdumpManager()
|
||||
@@ -0,0 +1,156 @@
|
||||
import asyncio
|
||||
import re
|
||||
import shutil
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
RESERVED_OPTIONS = {"mode", "listen_host", "listen_port", "scripts"}
|
||||
OPTIONS_DOCS_URL = "https://docs.mitmproxy.org/stable/concepts/options/"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class MitmOption:
|
||||
name: str
|
||||
value_type: str
|
||||
default: Any
|
||||
description: str
|
||||
choices: list[str]
|
||||
|
||||
|
||||
class MitmOptionError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
catalog_cache: dict[str, MitmOption] | None = None
|
||||
|
||||
|
||||
async def load_options() -> dict[str, MitmOption]:
|
||||
global catalog_cache
|
||||
if catalog_cache is None:
|
||||
catalog_cache = await read_options()
|
||||
return catalog_cache
|
||||
|
||||
|
||||
async def read_options(command: str = "mitmdump") -> dict[str, MitmOption]:
|
||||
command_path = shutil.which(command)
|
||||
if command_path is None:
|
||||
raise MitmOptionError("mitmdump executable not found in PATH")
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
command_path,
|
||||
"--options",
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
stdout, stderr = await process.communicate()
|
||||
if process.returncode != 0:
|
||||
message = stderr.decode("utf-8", errors="replace").strip()
|
||||
raise MitmOptionError(message or "Unable to load mitmdump options")
|
||||
|
||||
text = stdout.decode("utf-8", errors="replace")
|
||||
defaults = yaml.safe_load(text) or {}
|
||||
return parse_options(text, defaults)
|
||||
|
||||
|
||||
def parse_options(text: str, defaults: dict[str, Any]) -> dict[str, MitmOption]:
|
||||
comments: list[str] = []
|
||||
options: dict[str, MitmOption] = {}
|
||||
|
||||
for line in text.splitlines():
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
comments = []
|
||||
continue
|
||||
if stripped.startswith("#"):
|
||||
comments.append(stripped[1:].strip())
|
||||
continue
|
||||
if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*:", stripped):
|
||||
continue
|
||||
|
||||
name = stripped.split(":", 1)[0]
|
||||
description = " ".join(comments)
|
||||
options[name] = MitmOption(
|
||||
name=name,
|
||||
value_type=parse_type(description),
|
||||
default=defaults.get(name),
|
||||
description=description,
|
||||
choices=parse_choices(description),
|
||||
)
|
||||
comments = []
|
||||
|
||||
return options
|
||||
|
||||
|
||||
def parse_type(description: str) -> str:
|
||||
match = re.search(r"Type ([^.]+)\.", description)
|
||||
if match is None:
|
||||
return "unknown"
|
||||
return match.group(1).strip()
|
||||
|
||||
|
||||
def parse_choices(description: str) -> list[str]:
|
||||
match = re.search(r"Valid values are ([^.]+)\.", description)
|
||||
if match is None:
|
||||
return []
|
||||
return re.findall(r"'([^']+)'", match.group(1))
|
||||
|
||||
|
||||
async def validate_options(options: dict[str, Any]) -> None:
|
||||
catalog = await load_options()
|
||||
errors: list[str] = []
|
||||
|
||||
for name, value in options.items():
|
||||
option = catalog.get(name)
|
||||
if option is None:
|
||||
errors.append(f"Unknown mitmdump option: {name}")
|
||||
continue
|
||||
if name in RESERVED_OPTIONS:
|
||||
errors.append(f"Option is managed by thePlug and cannot be set: {name}")
|
||||
continue
|
||||
message = validate_value(option, value)
|
||||
if message:
|
||||
errors.append(f"{name}: {message}")
|
||||
|
||||
if errors:
|
||||
raise MitmOptionError("; ".join(errors))
|
||||
|
||||
|
||||
def validate_value(option: MitmOption, value: Any) -> str | None:
|
||||
value_type = option.value_type
|
||||
if option.choices and str(value) not in option.choices:
|
||||
return f"expected one of {', '.join(option.choices)}"
|
||||
if value is None:
|
||||
if value_type.startswith("optional "):
|
||||
return None
|
||||
return "cannot be null"
|
||||
if value_type == "bool" and not isinstance(value, bool):
|
||||
return "expected bool"
|
||||
if value_type == "int" and not isinstance(value, int):
|
||||
return "expected int"
|
||||
if value_type == "str" and not isinstance(value, str):
|
||||
return "expected str"
|
||||
if value_type == "sequence of str":
|
||||
if not isinstance(value, list) or not all(isinstance(v, str) for v in value):
|
||||
return "expected list[str]"
|
||||
if value_type == "optional str" and not isinstance(value, str):
|
||||
return "expected str or null"
|
||||
if value_type == "optional int" and not isinstance(value, int):
|
||||
return "expected int or null"
|
||||
return None
|
||||
|
||||
|
||||
def public_options(options: dict[str, MitmOption]) -> list[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
"name": option.name,
|
||||
"type": option.value_type,
|
||||
"default": option.default,
|
||||
"description": option.description,
|
||||
"choices": option.choices,
|
||||
"reserved": option.name in RESERVED_OPTIONS,
|
||||
"docs_url": f"{OPTIONS_DOCS_URL}#{option.name}",
|
||||
}
|
||||
for option in options.values()
|
||||
]
|
||||
@@ -0,0 +1,164 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
TARGET_RUFF_CODES = ("F821", "F822", "F823")
|
||||
|
||||
|
||||
class ProxyCodeValidationError(Exception):
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
*,
|
||||
errors: list[dict[str, Any]] | None = None,
|
||||
) -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.errors = errors or []
|
||||
|
||||
def to_http_detail(self) -> dict[str, Any]:
|
||||
return {
|
||||
"message": self.message,
|
||||
"errors": self.errors,
|
||||
}
|
||||
|
||||
|
||||
def _validate_python_syntax(code: str) -> None:
|
||||
try:
|
||||
compile(code, "<proxy_code>", "exec")
|
||||
except SyntaxError as error:
|
||||
line = error.lineno if error.lineno is not None else 0
|
||||
column = error.offset if error.offset is not None else 0
|
||||
detail = {
|
||||
"code": "COMPILE",
|
||||
"message": error.msg,
|
||||
"location": {"row": line, "column": column},
|
||||
}
|
||||
raise ProxyCodeValidationError(
|
||||
"Proxy code has invalid Python syntax",
|
||||
errors=[detail],
|
||||
) from error
|
||||
|
||||
|
||||
def _get_ruff_executable() -> list[str]:
|
||||
ruff_path = shutil.which("ruff")
|
||||
if ruff_path:
|
||||
return [ruff_path]
|
||||
|
||||
uv_path = shutil.which("uv")
|
||||
if uv_path:
|
||||
return [uv_path, "run", "ruff"]
|
||||
|
||||
raise ProxyCodeValidationError("Unable to run ruff: executable not found")
|
||||
|
||||
|
||||
def _parse_ruff_output(output: str) -> list[dict[str, Any]]:
|
||||
if not output.strip():
|
||||
return []
|
||||
|
||||
try:
|
||||
parsed = json.loads(output)
|
||||
except json.JSONDecodeError as error:
|
||||
raise ProxyCodeValidationError("Unable to parse ruff JSON output") from error
|
||||
|
||||
if not isinstance(parsed, list):
|
||||
raise ProxyCodeValidationError("Unexpected ruff output format")
|
||||
|
||||
errors: list[dict[str, Any]] = []
|
||||
for item in parsed:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
code = item.get("code")
|
||||
if code not in TARGET_RUFF_CODES:
|
||||
continue
|
||||
|
||||
message = item.get("message")
|
||||
raw_location = item.get("location")
|
||||
location = raw_location if isinstance(raw_location, dict) else {}
|
||||
row = location.get("row", 0)
|
||||
column = location.get("column", 0)
|
||||
|
||||
errors.append(
|
||||
{
|
||||
"code": code,
|
||||
"message": str(message) if message is not None else "",
|
||||
"location": {"row": row, "column": column},
|
||||
}
|
||||
)
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
async def _run_ruff_check(code: str) -> list[dict[str, Any]]:
|
||||
command = _get_ruff_executable()
|
||||
|
||||
file_path = ""
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
suffix=".py",
|
||||
encoding="utf-8",
|
||||
delete=False,
|
||||
) as handle:
|
||||
handle.write(code)
|
||||
file_path = handle.name
|
||||
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*command,
|
||||
"check",
|
||||
"--output-format=json",
|
||||
"--select",
|
||||
",".join(TARGET_RUFF_CODES),
|
||||
file_path,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
stdout, stderr = await process.communicate()
|
||||
finally:
|
||||
if file_path:
|
||||
try:
|
||||
os.unlink(file_path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
stdout_text = stdout.decode("utf-8", errors="replace")
|
||||
stderr_text = stderr.decode("utf-8", errors="replace")
|
||||
|
||||
if process.returncode not in (0, 1):
|
||||
raise ProxyCodeValidationError(
|
||||
"Ruff check failed to execute",
|
||||
errors=[
|
||||
{
|
||||
"code": "RUFF",
|
||||
"message": stderr_text.strip() or "Unknown ruff failure",
|
||||
"location": {"row": 0, "column": 0},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
if stderr_text.strip() and process.returncode == 0:
|
||||
raise ProxyCodeValidationError(
|
||||
"Ruff emitted unexpected stderr output",
|
||||
errors=[
|
||||
{
|
||||
"code": "RUFF",
|
||||
"message": stderr_text.strip(),
|
||||
"location": {"row": 0, "column": 0},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
return _parse_ruff_output(stdout_text)
|
||||
|
||||
|
||||
async def validate_proxy_code(code: str) -> None:
|
||||
_validate_python_syntax(code)
|
||||
errors = await _run_ruff_check(code)
|
||||
if errors:
|
||||
raise ProxyCodeValidationError(
|
||||
"Proxy code contains ruff violations",
|
||||
errors=errors,
|
||||
)
|
||||
@@ -0,0 +1,39 @@
|
||||
[project]
|
||||
name = "theplug"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
requires-python = ">=3.14"
|
||||
dependencies = [
|
||||
"authlib>=1.6.0",
|
||||
"basedpyright>=1.39.0",
|
||||
"fastapi[standard]>=0.135.3",
|
||||
"mitmproxy>=12.2.2",
|
||||
"motor>=3.7.1",
|
||||
"pwdlib[argon2,bcrypt]>=0.3.0",
|
||||
"pydantic>=2.12.5",
|
||||
"pydantic-settings>=2.13.1",
|
||||
"pyyaml>=6.0.3",
|
||||
"redis>=7.4.0",
|
||||
"ruff>=0.15.10",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py314"
|
||||
line-length = 88
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "B", "UP"]
|
||||
|
||||
[tool.ruff.format]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
pythonpath = ["app"]
|
||||
addopts = "-q"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=9.0.3",
|
||||
]
|
||||
@@ -0,0 +1,112 @@
|
||||
import importlib
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from core.security import AuthContext, get_auth_context
|
||||
from models.auth import AccessTokenPayload, AuthUserDocument
|
||||
|
||||
auth_router_module = importlib.import_module("api.auth.router")
|
||||
|
||||
|
||||
def test_login_success_returns_bearer_token(client, monkeypatch: object):
|
||||
user = AuthUserDocument(username="admin", password_hash="stored-hash")
|
||||
|
||||
async def fake_get_user_by_username(username: str) -> AuthUserDocument | None:
|
||||
assert username == "admin"
|
||||
return user
|
||||
|
||||
def fake_verify_password(plain_password: str, password_hash: str) -> bool:
|
||||
return plain_password == "secret" and password_hash == "stored-hash"
|
||||
|
||||
def fake_create_access_token(*, subject: str) -> str:
|
||||
assert subject == "admin"
|
||||
return "token-123"
|
||||
|
||||
monkeypatch.setattr(
|
||||
auth_router_module,
|
||||
"get_user_by_username",
|
||||
fake_get_user_by_username,
|
||||
)
|
||||
monkeypatch.setattr(auth_router_module, "verify_password", fake_verify_password)
|
||||
monkeypatch.setattr(
|
||||
auth_router_module,
|
||||
"create_access_token",
|
||||
fake_create_access_token,
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/auth/login",
|
||||
json={"username": "admin", "password": "secret"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"access_token": "token-123", "token_type": "bearer"}
|
||||
|
||||
|
||||
def test_login_invalid_credentials_returns_401(client, monkeypatch: object):
|
||||
async def fake_get_user_by_username(_: str) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(
|
||||
auth_router_module,
|
||||
"get_user_by_username",
|
||||
fake_get_user_by_username,
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/auth/login",
|
||||
json={"username": "admin", "password": "wrong"},
|
||||
)
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.json()["detail"] == "Invalid username or password"
|
||||
assert response.headers["www-authenticate"] == "Bearer"
|
||||
|
||||
|
||||
def test_logout_revokes_active_token(client, monkeypatch: object):
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_revoke_token(*, jti: str, expires_at: datetime) -> None:
|
||||
captured["jti"] = jti
|
||||
captured["expires_at"] = expires_at
|
||||
|
||||
monkeypatch.setattr(auth_router_module, "revoke_token", fake_revoke_token)
|
||||
|
||||
response = client.post("/auth/logout")
|
||||
|
||||
assert response.status_code == 204
|
||||
assert captured["jti"] == "test-jti"
|
||||
assert isinstance(captured["expires_at"], datetime)
|
||||
|
||||
|
||||
def test_logout_skips_revocation_when_token_is_expired(
|
||||
client,
|
||||
fake_user: AuthUserDocument,
|
||||
monkeypatch: object,
|
||||
):
|
||||
state = {"called": False}
|
||||
|
||||
async def fake_revoke_token(*, jti: str, expires_at: datetime) -> None:
|
||||
state["called"] = True
|
||||
|
||||
monkeypatch.setattr(auth_router_module, "revoke_token", fake_revoke_token)
|
||||
|
||||
expired_payload = AccessTokenPayload(
|
||||
sub=fake_user.username,
|
||||
jti="expired-jti",
|
||||
exp=int((datetime.now(UTC) - timedelta(minutes=1)).timestamp()),
|
||||
)
|
||||
expired_context = AuthContext(
|
||||
token="expired-token",
|
||||
user=fake_user,
|
||||
payload=expired_payload,
|
||||
)
|
||||
|
||||
async def override_auth_context() -> AuthContext:
|
||||
return expired_context
|
||||
|
||||
client.app.dependency_overrides[get_auth_context] = override_auth_context
|
||||
|
||||
response = client.post("/auth/logout")
|
||||
|
||||
assert response.status_code == 204
|
||||
assert state["called"] is False
|
||||
@@ -0,0 +1,52 @@
|
||||
import pytest
|
||||
from services.editor import python_lsp
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
|
||||
def test_python_editor_session(client, monkeypatch):
|
||||
session = python_lsp.PythonEditorSession(
|
||||
root_path=python_lsp.WORKSPACE_ROOT,
|
||||
file_path=python_lsp.WORKSPACE_ROOT / python_lsp.ENTRY_FILE,
|
||||
)
|
||||
monkeypatch.setattr("api.editor.router.python_lsp.lsp_binary", lambda: "/bin/lsp")
|
||||
monkeypatch.setattr("api.editor.router.python_lsp.get_session", lambda: session)
|
||||
|
||||
response = client.get("/proxy/editor/python/session")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"root_uri": session.root_uri,
|
||||
"file_uri": session.file_uri,
|
||||
}
|
||||
|
||||
|
||||
def test_python_editor_socket_rejects_missing_token(client):
|
||||
with pytest.raises(WebSocketDisconnect) as error:
|
||||
with client.websocket_connect("/proxy/editor/python/ws"):
|
||||
pass
|
||||
|
||||
assert error.value.code == 4401
|
||||
|
||||
|
||||
def test_python_editor_socket_accepts_token(client, monkeypatch):
|
||||
accepted = {"value": False}
|
||||
|
||||
async def fake_authenticate(_: object) -> None:
|
||||
return None
|
||||
|
||||
async def fake_bridge(websocket, session) -> None:
|
||||
accepted["value"] = True
|
||||
await websocket.send_text(session.file_uri)
|
||||
await websocket.close()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"api.editor.router.authenticate_editor_socket",
|
||||
fake_authenticate,
|
||||
)
|
||||
monkeypatch.setattr("api.editor.router.python_lsp.lsp_binary", lambda: "/bin/lsp")
|
||||
monkeypatch.setattr("api.editor.router.python_lsp.bridge", fake_bridge)
|
||||
|
||||
with client.websocket_connect("/proxy/editor/python/ws?token=test-token") as socket:
|
||||
assert socket.receive_text().endswith("/proxy_addon.py")
|
||||
|
||||
assert accepted["value"] is True
|
||||
@@ -0,0 +1,460 @@
|
||||
import importlib
|
||||
|
||||
from core.security import get_auth_context, require_authenticated_user
|
||||
from models.proxy import ProxyCreateInput, ProxyDocument
|
||||
from models.service import ServiceEndpointDocument
|
||||
from services.mitmproxy.examples import BUILTIN_EXAMPLES
|
||||
from services.mitmproxy.validator import ProxyCodeValidationError
|
||||
|
||||
proxy_router_module = importlib.import_module("api.proxy.router")
|
||||
|
||||
|
||||
def make_proxy_document(
|
||||
*,
|
||||
proxy_port: int,
|
||||
name: str = "svc",
|
||||
code: str = "proxy-code",
|
||||
port: int = 9000,
|
||||
) -> ProxyDocument:
|
||||
return ProxyDocument(
|
||||
proxy_port=proxy_port,
|
||||
name=name,
|
||||
code=code,
|
||||
service=ServiceEndpointDocument(port=port),
|
||||
)
|
||||
|
||||
|
||||
def test_proxy_requires_auth_without_dependency_overrides(client):
|
||||
client.app.dependency_overrides.pop(require_authenticated_user, None)
|
||||
client.app.dependency_overrides.pop(get_auth_context, None)
|
||||
|
||||
response = client.get("/proxy")
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.json()["detail"] == "Could not validate credentials"
|
||||
|
||||
|
||||
def test_list_proxies_returns_documents(client, monkeypatch: object):
|
||||
docs = [
|
||||
make_proxy_document(proxy_port=8080, name="api"),
|
||||
make_proxy_document(proxy_port=8081, name="worker"),
|
||||
]
|
||||
|
||||
async def fake_list_proxies() -> list[ProxyDocument]:
|
||||
return docs
|
||||
|
||||
monkeypatch.setattr(proxy_router_module, "list_proxy_documents", fake_list_proxies)
|
||||
|
||||
response = client.get("/proxy")
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert [item["proxy_port"] for item in body] == [8080, 8081]
|
||||
|
||||
|
||||
def test_get_proxy_returns_404_when_not_found(client, monkeypatch: object):
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
|
||||
response = client.get("/proxy/9999")
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.json()["detail"] == "Proxy not found"
|
||||
|
||||
|
||||
def test_create_proxy_returns_409_on_port_conflict(client, monkeypatch: object):
|
||||
existing = make_proxy_document(proxy_port=3671)
|
||||
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
assert proxy_port == 3671
|
||||
return existing
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/proxy",
|
||||
json={
|
||||
"proxy_port": 3671,
|
||||
"name": "new",
|
||||
"code": "abc",
|
||||
"service": {"port": 8080},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 409
|
||||
assert response.json()["detail"] == "Proxy already exists on port 3671"
|
||||
|
||||
|
||||
def test_create_proxy_returns_created_document(client, monkeypatch: object):
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
assert proxy_port == 4040
|
||||
return None
|
||||
|
||||
async def fake_create_proxy(payload: ProxyCreateInput) -> ProxyDocument:
|
||||
captured["proxy_port"] = payload.proxy_port
|
||||
captured["name"] = payload.name
|
||||
captured["service"] = payload.service.model_dump(exclude_none=True)
|
||||
captured["code"] = payload.code
|
||||
captured["protocol"] = payload.protocol
|
||||
return make_proxy_document(
|
||||
proxy_port=payload.proxy_port,
|
||||
name=payload.name,
|
||||
code=payload.code,
|
||||
port=payload.service.port,
|
||||
)
|
||||
|
||||
async def fake_validate_proxy_code(code: str) -> None:
|
||||
assert code == "def proxy_handler():\n return 'ok'\n"
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"validate_proxy_code",
|
||||
fake_validate_proxy_code,
|
||||
)
|
||||
monkeypatch.setattr(proxy_router_module, "create_proxy", fake_create_proxy)
|
||||
|
||||
response = client.post(
|
||||
"/proxy",
|
||||
json={
|
||||
"proxy_port": 4040,
|
||||
"name": "edge",
|
||||
"code": "def proxy_handler():\n return 'ok'\n",
|
||||
"service": {"port": 8181},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
assert response.json()["proxy_port"] == 4040
|
||||
assert captured == {
|
||||
"proxy_port": 4040,
|
||||
"name": "edge",
|
||||
"service": {"host": "127.0.0.1", "port": 8181},
|
||||
"code": "def proxy_handler():\n return 'ok'\n",
|
||||
"protocol": "http",
|
||||
}
|
||||
|
||||
|
||||
def test_create_proxy_returns_422_when_code_validation_fails(
|
||||
client,
|
||||
monkeypatch: object,
|
||||
):
|
||||
create_called = False
|
||||
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
assert proxy_port == 4000
|
||||
return None
|
||||
|
||||
async def fake_validate_proxy_code(code: str) -> None:
|
||||
assert code == "invalid-code"
|
||||
raise ProxyCodeValidationError(
|
||||
"Proxy code contains ruff violations",
|
||||
errors=[
|
||||
{
|
||||
"code": "F821",
|
||||
"message": "Undefined name `missing_symbol`",
|
||||
"location": {"row": 1, "column": 1},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
async def fake_create_proxy(
|
||||
*,
|
||||
proxy_port: int,
|
||||
name: str,
|
||||
service: dict[str, object],
|
||||
code: str,
|
||||
) -> ProxyDocument:
|
||||
nonlocal create_called
|
||||
create_called = True
|
||||
return make_proxy_document(proxy_port=proxy_port, name=name, code=code)
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"validate_proxy_code",
|
||||
fake_validate_proxy_code,
|
||||
)
|
||||
monkeypatch.setattr(proxy_router_module, "create_proxy", fake_create_proxy)
|
||||
|
||||
response = client.post(
|
||||
"/proxy",
|
||||
json={
|
||||
"proxy_port": 4000,
|
||||
"name": "broken",
|
||||
"code": "invalid-code",
|
||||
"service": {"port": 9000},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
body = response.json()["detail"]
|
||||
assert body["message"] == "Proxy code contains ruff violations"
|
||||
assert body["errors"][0]["code"] == "F821"
|
||||
assert create_called is False
|
||||
|
||||
|
||||
def test_update_proxy_returns_422_when_code_validation_fails(
|
||||
client,
|
||||
monkeypatch: object,
|
||||
):
|
||||
update_called = False
|
||||
|
||||
async def fake_validate_proxy_code(code: str) -> None:
|
||||
assert code == "invalid-code"
|
||||
raise ProxyCodeValidationError(
|
||||
"Proxy code has invalid Python syntax",
|
||||
errors=[
|
||||
{
|
||||
"code": "COMPILE",
|
||||
"message": "invalid syntax",
|
||||
"location": {"row": 1, "column": 1},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
async def fake_update_proxy(
|
||||
proxy_port: int,
|
||||
payload: object,
|
||||
) -> ProxyDocument | None:
|
||||
nonlocal update_called
|
||||
update_called = True
|
||||
return make_proxy_document(proxy_port=proxy_port)
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"validate_proxy_code",
|
||||
fake_validate_proxy_code,
|
||||
)
|
||||
monkeypatch.setattr(proxy_router_module, "update_proxy", fake_update_proxy)
|
||||
|
||||
response = client.put(
|
||||
"/proxy/8123",
|
||||
json={"code": "invalid-code"},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
body = response.json()["detail"]
|
||||
assert body["message"] == "Proxy code has invalid Python syntax"
|
||||
assert body["errors"][0]["code"] == "COMPILE"
|
||||
assert update_called is False
|
||||
|
||||
|
||||
def test_start_proxy_updates_runtime_state(client, monkeypatch: object):
|
||||
proxy = make_proxy_document(proxy_port=4040)
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
assert proxy_port == 4040
|
||||
return proxy
|
||||
|
||||
class FakeManager:
|
||||
async def start(self, document: ProxyDocument) -> dict[str, object]:
|
||||
assert document.proxy_port == 4040
|
||||
return {
|
||||
"status": "running",
|
||||
"pid": 1234,
|
||||
"script_path": "/tmp/theplug/addons/4040.py",
|
||||
"log_path": "/tmp/theplug/logs/4040.log",
|
||||
"service.upstream_port": 9000,
|
||||
"last_error": None,
|
||||
}
|
||||
|
||||
async def fake_update_proxy_runtime(
|
||||
proxy_port: int,
|
||||
values: dict[str, object],
|
||||
) -> ProxyDocument:
|
||||
captured["proxy_port"] = proxy_port
|
||||
captured["values"] = values
|
||||
return proxy.model_copy(update={"status": "running", "pid": 1234})
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
monkeypatch.setattr(proxy_router_module, "manager", FakeManager())
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"update_proxy_runtime",
|
||||
fake_update_proxy_runtime,
|
||||
)
|
||||
|
||||
response = client.post("/proxy/4040/start")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "running"
|
||||
assert captured["proxy_port"] == 4040
|
||||
assert captured["values"]["pid"] == 1234
|
||||
|
||||
|
||||
def test_start_proxy_records_runtime_error(client, monkeypatch: object):
|
||||
proxy = make_proxy_document(proxy_port=4040)
|
||||
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
assert proxy_port == 4040
|
||||
return proxy
|
||||
|
||||
class FakeManager:
|
||||
async def start(self, document: ProxyDocument) -> dict[str, object]:
|
||||
assert document.proxy_port == 4040
|
||||
raise proxy_router_module.MitmdumpError("mitmdump executable not found")
|
||||
|
||||
async def fake_update_proxy_runtime(
|
||||
proxy_port: int,
|
||||
values: dict[str, object],
|
||||
) -> ProxyDocument:
|
||||
assert proxy_port == 4040
|
||||
assert values == {
|
||||
"status": "error",
|
||||
"last_error": "mitmdump executable not found",
|
||||
}
|
||||
return proxy.model_copy(
|
||||
update={"status": "error", "last_error": values["last_error"]}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
monkeypatch.setattr(proxy_router_module, "manager", FakeManager())
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"update_proxy_runtime",
|
||||
fake_update_proxy_runtime,
|
||||
)
|
||||
|
||||
response = client.post("/proxy/4040/start")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "error"
|
||||
assert response.json()["last_error"] == "mitmdump executable not found"
|
||||
|
||||
|
||||
def test_stop_proxy_updates_runtime_state(client, monkeypatch: object):
|
||||
proxy = make_proxy_document(proxy_port=4040).model_copy(
|
||||
update={"status": "running", "pid": 1234}
|
||||
)
|
||||
|
||||
async def fake_get_proxy_by_port(proxy_port: int) -> ProxyDocument | None:
|
||||
assert proxy_port == 4040
|
||||
return proxy
|
||||
|
||||
class FakeManager:
|
||||
async def stop(self, document: ProxyDocument) -> dict[str, object]:
|
||||
assert document.pid == 1234
|
||||
return {
|
||||
"status": "stopped",
|
||||
"pid": None,
|
||||
"service.upstream_port": None,
|
||||
"last_error": None,
|
||||
}
|
||||
|
||||
async def fake_update_proxy_runtime(
|
||||
proxy_port: int,
|
||||
values: dict[str, object],
|
||||
) -> ProxyDocument:
|
||||
assert proxy_port == 4040
|
||||
assert values["status"] == "stopped"
|
||||
return proxy.model_copy(update={"status": "stopped", "pid": None})
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"get_proxy_by_port",
|
||||
fake_get_proxy_by_port,
|
||||
)
|
||||
monkeypatch.setattr(proxy_router_module, "manager", FakeManager())
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"update_proxy_runtime",
|
||||
fake_update_proxy_runtime,
|
||||
)
|
||||
|
||||
response = client.post("/proxy/4040/stop")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "stopped"
|
||||
|
||||
|
||||
def test_list_mitmproxy_examples_returns_builtins_and_custom(client, monkeypatch):
|
||||
async def fake_list_custom_examples():
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"list_custom_examples",
|
||||
fake_list_custom_examples,
|
||||
)
|
||||
|
||||
response = client.get("/proxy/mitmproxy/examples")
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert len(body) == len(BUILTIN_EXAMPLES)
|
||||
assert body[0]["built_in"] is True
|
||||
|
||||
|
||||
def test_create_mitmproxy_example_validates_and_persists(client, monkeypatch):
|
||||
captured = {}
|
||||
|
||||
async def fake_validate_proxy_code(code: str) -> None:
|
||||
captured["validated"] = code
|
||||
|
||||
async def fake_create_custom_example(payload):
|
||||
captured["payload"] = payload
|
||||
return {
|
||||
"id": payload.id,
|
||||
"name": payload.name,
|
||||
"description": payload.description,
|
||||
"category": payload.category,
|
||||
"code": payload.code,
|
||||
"source_url": payload.source_url,
|
||||
"built_in": False,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"validate_proxy_code",
|
||||
fake_validate_proxy_code,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"create_custom_example",
|
||||
fake_create_custom_example,
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/proxy/mitmproxy/examples",
|
||||
json={
|
||||
"id": "custom-header",
|
||||
"name": "Custom header",
|
||||
"description": "Adds a custom response header",
|
||||
"category": "http",
|
||||
"code": "def response(flow):\n pass\n",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
assert response.json()["id"] == "custom-header"
|
||||
assert captured["validated"] == "def response(flow):\n pass\n"
|
||||
@@ -0,0 +1,70 @@
|
||||
import os
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
os.environ.setdefault(
|
||||
"MONGODB_URI",
|
||||
"mongodb://theplug:pass@localhost:27017/theplug?authSource=admin",
|
||||
)
|
||||
os.environ.setdefault("MONGODB_DB_NAME", "theplug")
|
||||
os.environ.setdefault("REDIS_URL", "redis://localhost:6379")
|
||||
os.environ.setdefault("JWT_SECRET_KEY", "test-secret-key")
|
||||
os.environ.setdefault("JWT_ALGORITHM", "HS256")
|
||||
os.environ.setdefault("JWT_EXPIRE_MINUTES", "30")
|
||||
os.environ.setdefault("AUTH_USERNAME", "admin")
|
||||
os.environ.setdefault("AUTH_PASSWORD", "admin")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_user():
|
||||
from models.auth import AuthUserDocument
|
||||
|
||||
return AuthUserDocument(
|
||||
username="admin",
|
||||
password_hash="hashed-password",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def auth_context(fake_user):
|
||||
from core.security import AuthContext
|
||||
from models.auth import AccessTokenPayload
|
||||
|
||||
exp = int((datetime.now(UTC) + timedelta(minutes=30)).timestamp())
|
||||
payload = AccessTokenPayload(sub=fake_user.username, jti="test-jti", exp=exp)
|
||||
return AuthContext(token="test-token", user=fake_user, payload=payload)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app(fake_user, auth_context) -> FastAPI:
|
||||
from api.auth.router import router as auth_router
|
||||
from api.editor.router import router as editor_router
|
||||
from api.proxy.router import router as proxy_router
|
||||
from core.security import AuthContext, get_auth_context, require_authenticated_user
|
||||
from models.auth import AuthUserDocument
|
||||
|
||||
app_instance = FastAPI()
|
||||
app_instance.include_router(auth_router)
|
||||
app_instance.include_router(editor_router)
|
||||
app_instance.include_router(proxy_router)
|
||||
|
||||
async def _override_current_user() -> AuthUserDocument:
|
||||
return fake_user
|
||||
|
||||
async def _override_auth_context() -> AuthContext:
|
||||
return auth_context
|
||||
|
||||
app_instance.dependency_overrides[require_authenticated_user] = (
|
||||
_override_current_user
|
||||
)
|
||||
app_instance.dependency_overrides[get_auth_context] = _override_auth_context
|
||||
return app_instance
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(app: FastAPI):
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
@@ -0,0 +1,43 @@
|
||||
import asyncio
|
||||
|
||||
from models.proxy import ProxyDocument
|
||||
from models.service import ServiceEndpointDocument
|
||||
from services.mitmproxy.health import check_proxy
|
||||
|
||||
|
||||
def test_check_proxy_reports_upstream_and_proxy_healthy():
|
||||
async def run():
|
||||
upstream = await asyncio.start_server(http_handler, "127.0.0.1", 0)
|
||||
proxy = await asyncio.start_server(http_handler, "127.0.0.1", 0)
|
||||
try:
|
||||
upstream_port = upstream.sockets[0].getsockname()[1]
|
||||
proxy_port = proxy.sockets[0].getsockname()[1]
|
||||
document = ProxyDocument(
|
||||
proxy_port=proxy_port,
|
||||
name="health",
|
||||
code="",
|
||||
service=ServiceEndpointDocument(port=upstream_port),
|
||||
)
|
||||
|
||||
health = await check_proxy(document)
|
||||
finally:
|
||||
upstream.close()
|
||||
proxy.close()
|
||||
await upstream.wait_closed()
|
||||
await proxy.wait_closed()
|
||||
|
||||
assert health["overall"] == "healthy"
|
||||
assert [check["target"] for check in health["checks"]] == [
|
||||
"upstream",
|
||||
"proxy",
|
||||
]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
async def http_handler(reader, writer):
|
||||
await reader.readline()
|
||||
writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")
|
||||
await writer.drain()
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
@@ -0,0 +1,75 @@
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
from models.proxy import ProxyDocument
|
||||
from models.service import ServiceEndpointDocument
|
||||
from services.mitmproxy.manager import MitmdumpManager
|
||||
|
||||
|
||||
def make_proxy(**updates):
|
||||
data = {
|
||||
"proxy_port": 8080,
|
||||
"name": "svc",
|
||||
"code": "",
|
||||
"service": ServiceEndpointDocument(port=9000),
|
||||
}
|
||||
data.update(updates)
|
||||
return ProxyDocument(**data)
|
||||
|
||||
|
||||
def test_reverse_mode_is_built_from_service():
|
||||
manager = MitmdumpManager()
|
||||
proxy = make_proxy(protocol="tcp")
|
||||
|
||||
command = manager.build_command(
|
||||
"mitmdump",
|
||||
proxy,
|
||||
9100,
|
||||
Path("/tmp/addon.py"),
|
||||
Path("/tmp/conf/config.yaml"),
|
||||
)
|
||||
|
||||
assert "reverse:tcp://127.0.0.1:9100@8080" in command
|
||||
assert "--listen-port" not in command
|
||||
|
||||
|
||||
def test_raw_mode_spec_overrides_generated_mode():
|
||||
manager = MitmdumpManager()
|
||||
proxy = make_proxy(mode="regular", mode_spec="wireguard:/tmp/wg.conf@51820")
|
||||
|
||||
command = manager.build_command(
|
||||
"mitmdump",
|
||||
proxy,
|
||||
None,
|
||||
Path("/tmp/addon.py"),
|
||||
Path("/tmp/conf/config.yaml"),
|
||||
)
|
||||
|
||||
assert "wireguard:/tmp/wg.conf@51820" in command
|
||||
assert "--listen-port" not in command
|
||||
|
||||
|
||||
def test_regular_mode_uses_proxy_port_as_listen_port():
|
||||
manager = MitmdumpManager()
|
||||
proxy = make_proxy(mode="regular")
|
||||
|
||||
command = manager.build_command(
|
||||
"mitmdump",
|
||||
proxy,
|
||||
None,
|
||||
Path("/tmp/addon.py"),
|
||||
Path("/tmp/conf/config.yaml"),
|
||||
)
|
||||
|
||||
assert "regular" in command
|
||||
assert command[-2:] == ["--listen-port", "8080"]
|
||||
|
||||
|
||||
def test_running_document_without_pid_is_treated_as_stopped():
|
||||
manager = MitmdumpManager()
|
||||
proxy = make_proxy(status="running", pid=None)
|
||||
|
||||
async def run():
|
||||
return await manager.status(proxy)
|
||||
|
||||
assert asyncio.run(run()) == "stopped"
|
||||
@@ -0,0 +1,51 @@
|
||||
import pytest
|
||||
from services.mitmproxy.options import (
|
||||
MitmOption,
|
||||
parse_options,
|
||||
public_options,
|
||||
validate_value,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_options_extracts_types_and_choices():
|
||||
text = """
|
||||
# Enable feature. Type bool.
|
||||
enabled: false
|
||||
|
||||
# Valid values are 'eager', 'lazy'.
|
||||
strategy: eager
|
||||
|
||||
# Hosts. Type sequence of str.
|
||||
hosts: []
|
||||
"""
|
||||
options = parse_options(
|
||||
text,
|
||||
{"enabled": False, "strategy": "eager", "hosts": []},
|
||||
)
|
||||
|
||||
assert options["enabled"].value_type == "bool"
|
||||
assert options["strategy"].choices == ["eager", "lazy"]
|
||||
assert options["hosts"].value_type == "sequence of str"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("option", "value", "message"),
|
||||
[
|
||||
(MitmOption("anticache", "bool", False, "", []), "true", "expected bool"),
|
||||
(MitmOption("flow_detail", "int", 1, "", []), "1", "expected int"),
|
||||
(MitmOption("allow_hosts", "sequence of str", [], "", []), "x", "list[str]"),
|
||||
(MitmOption("strategy", "unknown", "eager", "", ["eager"]), "lazy", "one of"),
|
||||
],
|
||||
)
|
||||
def test_validate_value_rejects_wrong_types(option, value, message):
|
||||
assert message in validate_value(option, value)
|
||||
|
||||
|
||||
def test_public_options_include_docs_links():
|
||||
options = {
|
||||
"flow_detail": MitmOption("flow_detail", "int", 1, "Flow detail.", []),
|
||||
}
|
||||
|
||||
public = public_options(options)
|
||||
|
||||
assert public[0]["docs_url"].endswith("#flow_detail")
|
||||
@@ -0,0 +1,30 @@
|
||||
from services.editor import python_lsp
|
||||
|
||||
|
||||
def test_encode_message():
|
||||
message = '{"jsonrpc":"2.0"}'
|
||||
|
||||
payload = python_lsp.encode_message(message)
|
||||
|
||||
assert payload.endswith(message.encode("utf-8"))
|
||||
assert payload.startswith(
|
||||
f"Content-Length: {len(message.encode('utf-8'))}\r\n\r\n".encode("ascii")
|
||||
)
|
||||
|
||||
|
||||
def test_frame_buffer_reads_split_messages():
|
||||
buffer = python_lsp.LspFrameBuffer()
|
||||
first = python_lsp.encode_message('{"id":1}')
|
||||
second = python_lsp.encode_message('{"id":2}')
|
||||
|
||||
parts = [
|
||||
first[:12],
|
||||
first[12:] + second[:10],
|
||||
second[10:],
|
||||
]
|
||||
|
||||
messages: list[str] = []
|
||||
for part in parts:
|
||||
messages.extend(buffer.feed(part))
|
||||
|
||||
assert messages == ['{"id":1}', '{"id":2}']
|
||||
Generated
+1567
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user