feat: add proxy template origin field and implement socket utilities for log streaming
This commit is contained in:
gabdevele
+19
-128
@@ -5,15 +5,20 @@ from fastapi import APIRouter, Depends, HTTPException, WebSocket, status
|
||||
from pymongo.errors import DuplicateKeyError
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.security import (
|
||||
InvalidTokenError,
|
||||
authenticate_token,
|
||||
require_authenticated_user,
|
||||
from app.api.proxy.socket_utils import (
|
||||
authenticate_proxy_socket,
|
||||
proxy_log_path,
|
||||
read_log_snapshot,
|
||||
serialize_health_snapshot,
|
||||
snapshot_log_state,
|
||||
stream_proxy_logs,
|
||||
watch_socket_disconnect,
|
||||
)
|
||||
from app.core.config import settings
|
||||
from app.core.security import InvalidTokenError, require_authenticated_user
|
||||
from app.integrations.docker import get_all_services_in_dir
|
||||
from app.integrations.docker.discovery import service_scan_root
|
||||
from app.models.health import HealthHistoryResponse, HealthSnapshot
|
||||
from app.models.health import HealthHistoryResponse
|
||||
from app.models.proxy import ProxyCreateInput, ProxyUpdateInput
|
||||
from app.models.template import ProxyTemplateCreateInput
|
||||
from app.repositories.health import (
|
||||
@@ -51,7 +56,7 @@ from app.services.mitmproxy.runtime_validation import (
|
||||
backend_port,
|
||||
ensure_proxy_port_not_reserved,
|
||||
)
|
||||
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
|
||||
from app.services.mitmproxy.templates import file_template_ids, list_file_templates
|
||||
from app.services.mitmproxy.validator import (
|
||||
ProxyCodeValidationError,
|
||||
validate_proxy_code,
|
||||
@@ -152,125 +157,6 @@ async def get_proxy_or_404(proxy_port: int):
|
||||
return proxy
|
||||
|
||||
|
||||
async def authenticate_proxy_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)
|
||||
|
||||
|
||||
def proxy_log_path(proxy) -> Path:
|
||||
if proxy.log_path:
|
||||
return Path(proxy.log_path)
|
||||
return manager.log_path_for(proxy)
|
||||
|
||||
|
||||
def serialize_health_snapshot(
|
||||
snapshot: HealthSnapshot | None,
|
||||
) -> dict[str, object] | None:
|
||||
if snapshot is None:
|
||||
return None
|
||||
return snapshot.model_dump(mode="json")
|
||||
|
||||
|
||||
async def read_log_snapshot(path: Path, limit: int = 8000) -> str:
|
||||
snapshot, _offset, _marker = await snapshot_log_state(path, limit=limit)
|
||||
return snapshot
|
||||
|
||||
|
||||
async def snapshot_log_state(
|
||||
path: Path,
|
||||
limit: int = 8000,
|
||||
marker_width: int = 64,
|
||||
) -> tuple[str, int, bytes]:
|
||||
if not await asyncio.to_thread(path.is_file):
|
||||
return "", 0, b""
|
||||
|
||||
data = await asyncio.to_thread(path.read_bytes)
|
||||
size = len(data)
|
||||
marker = data[max(0, size - marker_width) : size]
|
||||
snapshot = data[-limit:].decode("utf-8", errors="replace")
|
||||
return snapshot, size, marker
|
||||
|
||||
|
||||
def _read_log_marker_bytes(path: Path, offset: int, width: int) -> bytes:
|
||||
if offset <= 0 or not path.is_file():
|
||||
return b""
|
||||
|
||||
start = max(0, offset - width)
|
||||
with path.open("rb") as handle:
|
||||
handle.seek(start)
|
||||
return handle.read(offset - start)
|
||||
|
||||
|
||||
async def read_log_marker(path: Path, offset: int, width: int = 64) -> bytes:
|
||||
return await asyncio.to_thread(_read_log_marker_bytes, path, offset, width)
|
||||
|
||||
|
||||
async def watch_socket_disconnect(websocket: WebSocket) -> None:
|
||||
while True:
|
||||
await websocket.receive_text()
|
||||
|
||||
|
||||
async def stream_proxy_logs(
|
||||
websocket: WebSocket,
|
||||
path: Path,
|
||||
*,
|
||||
offset: int,
|
||||
marker: bytes,
|
||||
) -> None:
|
||||
|
||||
while True:
|
||||
if not await asyncio.to_thread(path.exists):
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
try:
|
||||
size = await asyncio.to_thread(lambda: path.stat().st_size)
|
||||
except OSError:
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
if marker and size >= offset and await read_log_marker(path, offset) != marker:
|
||||
snapshot, offset, marker = await snapshot_log_state(path)
|
||||
await websocket.send_json({"type": "snapshot", "log": snapshot})
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
if size < offset:
|
||||
snapshot, offset, marker = await snapshot_log_state(path)
|
||||
await websocket.send_json({"type": "snapshot", "log": snapshot})
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
if size > offset:
|
||||
def read_log_chunk(
|
||||
chunk_offset: int = offset,
|
||||
chunk_size: int = size,
|
||||
) -> str:
|
||||
with path.open("rb") as handle:
|
||||
handle.seek(chunk_offset)
|
||||
return handle.read(chunk_size - chunk_offset).decode(
|
||||
"utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
|
||||
chunk = await asyncio.to_thread(read_log_chunk)
|
||||
offset = size
|
||||
marker = await read_log_marker(path, offset)
|
||||
if chunk:
|
||||
await websocket.send_json({"type": "append", "log": chunk})
|
||||
|
||||
await asyncio.sleep(0.4)
|
||||
|
||||
|
||||
@router.get("", dependencies=auth_dependencies)
|
||||
async def list_proxies_handler():
|
||||
return await list_proxy_documents()
|
||||
@@ -295,7 +181,7 @@ async def list_mitmproxy_options():
|
||||
|
||||
@router.get("/mitmproxy/templates", dependencies=auth_dependencies)
|
||||
async def list_mitmproxy_templates():
|
||||
return [*BUILTIN_TEMPLATES, *await list_custom_templates()]
|
||||
return [*list_file_templates(), *await list_custom_templates()]
|
||||
|
||||
|
||||
@router.post(
|
||||
@@ -305,6 +191,11 @@ async def list_mitmproxy_templates():
|
||||
)
|
||||
async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
|
||||
await _validate_code_or_raise(payload.code)
|
||||
if payload.id in file_template_ids():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=f"Template already exists: {payload.id}",
|
||||
)
|
||||
try:
|
||||
return await create_custom_template(payload)
|
||||
except DuplicateKeyError as error:
|
||||
@@ -320,7 +211,7 @@ async def create_mitmproxy_template(payload: ProxyTemplateCreateInput):
|
||||
dependencies=auth_dependencies,
|
||||
)
|
||||
async def delete_mitmproxy_template(template_id: str):
|
||||
if any(template.id == template_id for template in BUILTIN_TEMPLATES):
|
||||
if template_id in file_template_ids():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail="Built-in templates cannot be deleted",
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import WebSocket
|
||||
|
||||
from app.core.security import InvalidTokenError, authenticate_token
|
||||
from app.models.health import HealthSnapshot
|
||||
from app.services.mitmproxy.manager import manager
|
||||
|
||||
|
||||
async def authenticate_proxy_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)
|
||||
|
||||
|
||||
def proxy_log_path(proxy) -> Path:
|
||||
if proxy.log_path:
|
||||
return Path(proxy.log_path)
|
||||
return manager.log_path_for(proxy)
|
||||
|
||||
|
||||
def serialize_health_snapshot(
|
||||
snapshot: HealthSnapshot | None,
|
||||
) -> dict[str, object] | None:
|
||||
if snapshot is None:
|
||||
return None
|
||||
return snapshot.model_dump(mode="json")
|
||||
|
||||
|
||||
async def read_log_snapshot(path: Path, limit: int = 8000) -> str:
|
||||
snapshot, _offset, _marker = await snapshot_log_state(path, limit=limit)
|
||||
return snapshot
|
||||
|
||||
|
||||
async def snapshot_log_state(
|
||||
path: Path,
|
||||
limit: int = 8000,
|
||||
marker_width: int = 64,
|
||||
) -> tuple[str, int, bytes]:
|
||||
if not await asyncio.to_thread(path.is_file):
|
||||
return "", 0, b""
|
||||
|
||||
data = await asyncio.to_thread(path.read_bytes)
|
||||
size = len(data)
|
||||
marker = data[max(0, size - marker_width) : size]
|
||||
snapshot = data[-limit:].decode("utf-8", errors="replace")
|
||||
return snapshot, size, marker
|
||||
|
||||
|
||||
def _read_log_marker_bytes(path: Path, offset: int, width: int) -> bytes:
|
||||
if offset <= 0 or not path.is_file():
|
||||
return b""
|
||||
|
||||
start = max(0, offset - width)
|
||||
with path.open("rb") as handle:
|
||||
handle.seek(start)
|
||||
return handle.read(offset - start)
|
||||
|
||||
|
||||
async def read_log_marker(path: Path, offset: int, width: int = 64) -> bytes:
|
||||
return await asyncio.to_thread(_read_log_marker_bytes, path, offset, width)
|
||||
|
||||
|
||||
async def watch_socket_disconnect(websocket: WebSocket) -> None:
|
||||
while True:
|
||||
await websocket.receive_text()
|
||||
|
||||
|
||||
async def stream_proxy_logs(
|
||||
websocket: WebSocket,
|
||||
path: Path,
|
||||
*,
|
||||
offset: int,
|
||||
marker: bytes,
|
||||
) -> None:
|
||||
while True:
|
||||
if not await asyncio.to_thread(path.exists):
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
try:
|
||||
size = await asyncio.to_thread(lambda: path.stat().st_size)
|
||||
except OSError:
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
if marker and size >= offset and await read_log_marker(path, offset) != marker:
|
||||
snapshot, offset, marker = await snapshot_log_state(path)
|
||||
await websocket.send_json({"type": "snapshot", "log": snapshot})
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
if size < offset:
|
||||
snapshot, offset, marker = await snapshot_log_state(path)
|
||||
await websocket.send_json({"type": "snapshot", "log": snapshot})
|
||||
await asyncio.sleep(0.4)
|
||||
continue
|
||||
|
||||
if size > offset:
|
||||
|
||||
def read_log_chunk(
|
||||
chunk_offset: int = offset,
|
||||
chunk_size: int = size,
|
||||
) -> str:
|
||||
with path.open("rb") as handle:
|
||||
handle.seek(chunk_offset)
|
||||
return handle.read(chunk_size - chunk_offset).decode(
|
||||
"utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
|
||||
chunk = await asyncio.to_thread(read_log_chunk)
|
||||
offset = size
|
||||
marker = await read_log_marker(path, offset)
|
||||
if chunk:
|
||||
await websocket.send_json({"type": "append", "log": chunk})
|
||||
|
||||
await asyncio.sleep(0.4)
|
||||
@@ -19,8 +19,14 @@ class ProxyTemplateDocument(BaseModel):
|
||||
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")
|
||||
origin: str = Field(default="custom", description="Template source")
|
||||
created_at: datetime | None = None
|
||||
updated_at: datetime | None = None
|
||||
source_path: str | None = Field(
|
||||
default=None,
|
||||
description="Filesystem source path",
|
||||
exclude=True,
|
||||
)
|
||||
|
||||
|
||||
class ProxyTemplateCreateInput(BaseModel):
|
||||
|
||||
@@ -62,6 +62,7 @@ async def create_custom_template(
|
||||
code=payload.code,
|
||||
source_url=payload.source_url,
|
||||
built_in=False,
|
||||
origin="custom",
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
|
||||
@@ -4,12 +4,23 @@ import shutil
|
||||
import signal
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from app.integrations.docker import change_port
|
||||
from app.integrations.docker.ports import get_random_port
|
||||
from app.models.proxy import ProxyDocument
|
||||
from app.services.mitmproxy.options import MitmOptionError, validate_options
|
||||
from app.services.mitmproxy.runtime_files import (
|
||||
conf_dir_for,
|
||||
config_path_for,
|
||||
ensure_log_path,
|
||||
log_path_for,
|
||||
read_cmdline,
|
||||
read_log_tail,
|
||||
remove_dir,
|
||||
remove_file,
|
||||
script_path_for,
|
||||
write_config,
|
||||
write_script,
|
||||
)
|
||||
from app.services.mitmproxy.runtime_validation import (
|
||||
ReservedProxyPortError,
|
||||
ensure_proxy_port_not_reserved,
|
||||
@@ -300,23 +311,13 @@ class MitmdumpManager:
|
||||
return candidate
|
||||
|
||||
async def write_script(self, proxy: ProxyDocument) -> Path:
|
||||
path = self.script_path_for(proxy)
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
await asyncio.to_thread(path.write_text, proxy.code, encoding="utf-8")
|
||||
return path
|
||||
return await write_script(self.work_dir, proxy)
|
||||
|
||||
async def write_config(self, proxy: ProxyDocument) -> Path:
|
||||
path = self.config_path_for(proxy)
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
options = {"block_global": False, **proxy.mitm_options}
|
||||
config_text = yaml.safe_dump(options, sort_keys=True)
|
||||
await asyncio.to_thread(path.write_text, config_text, encoding="utf-8")
|
||||
return path
|
||||
return await write_config(self.work_dir, proxy)
|
||||
|
||||
async def log_path(self, proxy: ProxyDocument) -> Path:
|
||||
path = self.log_path_for(proxy)
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
return path
|
||||
return await ensure_log_path(self.work_dir, proxy)
|
||||
|
||||
async def spawn(self, command: list[str], log_path: Path):
|
||||
with log_path.open("ab") as log_file:
|
||||
@@ -408,49 +409,28 @@ class MitmdumpManager:
|
||||
return expected_script in cmdline and expected_confdir in cmdline
|
||||
|
||||
def read_cmdline(self, pid: int) -> list[str]:
|
||||
path = Path(f"/proc/{pid}/cmdline")
|
||||
try:
|
||||
raw = path.read_bytes()
|
||||
except OSError:
|
||||
return []
|
||||
decoded = raw.decode("utf-8", errors="replace")
|
||||
return [part for part in decoded.split("\x00") if part]
|
||||
return read_cmdline(pid)
|
||||
|
||||
def script_path_for(self, proxy: ProxyDocument) -> Path:
|
||||
return self.work_dir / "addons" / f"{proxy.proxy_port}.py"
|
||||
return script_path_for(self.work_dir, proxy)
|
||||
|
||||
def conf_dir_for(self, proxy: ProxyDocument) -> Path:
|
||||
return self.work_dir / "conf" / str(proxy.proxy_port)
|
||||
return conf_dir_for(self.work_dir, proxy)
|
||||
|
||||
def config_path_for(self, proxy: ProxyDocument) -> Path:
|
||||
return self.conf_dir_for(proxy) / "config.yaml"
|
||||
return config_path_for(self.work_dir, proxy)
|
||||
|
||||
def log_path_for(self, proxy: ProxyDocument) -> Path:
|
||||
return self.work_dir / "logs" / f"{proxy.proxy_port}.log"
|
||||
return log_path_for(self.work_dir, proxy)
|
||||
|
||||
async def read_log_tail(self, path: Path, limit: int = 4000) -> str:
|
||||
if not await asyncio.to_thread(path.exists):
|
||||
return ""
|
||||
content = await asyncio.to_thread(
|
||||
path.read_text,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
return content[-limit:].strip()
|
||||
return await read_log_tail(path, limit=limit)
|
||||
|
||||
def remove_file(self, path: Path) -> None:
|
||||
try:
|
||||
path.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
return
|
||||
remove_file(path)
|
||||
|
||||
def remove_dir(self, path: Path) -> None:
|
||||
try:
|
||||
shutil.rmtree(path)
|
||||
except FileNotFoundError:
|
||||
return
|
||||
except OSError:
|
||||
return
|
||||
remove_dir(path)
|
||||
|
||||
def lock_for(self, proxy_port: int) -> asyncio.Lock:
|
||||
lock = self.locks.get(proxy_port)
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
import asyncio
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from app.models.proxy import ProxyDocument
|
||||
|
||||
|
||||
def script_path_for(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
return work_dir / "addons" / f"{proxy.proxy_port}.py"
|
||||
|
||||
|
||||
def conf_dir_for(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
return work_dir / "conf" / str(proxy.proxy_port)
|
||||
|
||||
|
||||
def config_path_for(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
return conf_dir_for(work_dir, proxy) / "config.yaml"
|
||||
|
||||
|
||||
def log_path_for(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
return work_dir / "logs" / f"{proxy.proxy_port}.log"
|
||||
|
||||
|
||||
async def write_script(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
path = script_path_for(work_dir, proxy)
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
await asyncio.to_thread(path.write_text, proxy.code, encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
async def write_config(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
path = config_path_for(work_dir, proxy)
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
options = {"block_global": False, **proxy.mitm_options}
|
||||
config_text = yaml.safe_dump(options, sort_keys=True)
|
||||
await asyncio.to_thread(path.write_text, config_text, encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
async def ensure_log_path(work_dir: Path, proxy: ProxyDocument) -> Path:
|
||||
path = log_path_for(work_dir, proxy)
|
||||
await asyncio.to_thread(path.parent.mkdir, parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
async def read_log_tail(path: Path, limit: int = 4000) -> str:
|
||||
if not await asyncio.to_thread(path.exists):
|
||||
return ""
|
||||
content = await asyncio.to_thread(
|
||||
path.read_text,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
return content[-limit:].strip()
|
||||
|
||||
|
||||
def remove_file(path: Path) -> None:
|
||||
try:
|
||||
path.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
return
|
||||
|
||||
|
||||
def remove_dir(path: Path) -> None:
|
||||
try:
|
||||
shutil.rmtree(path)
|
||||
except FileNotFoundError:
|
||||
return
|
||||
except OSError:
|
||||
return
|
||||
|
||||
|
||||
def read_cmdline(pid: int) -> list[str]:
|
||||
path = Path(f"/proc/{pid}/cmdline")
|
||||
try:
|
||||
raw = path.read_bytes()
|
||||
except OSError:
|
||||
return []
|
||||
decoded = raw.decode("utf-8", errors="replace")
|
||||
return [part for part in decoded.split("\x00") if part]
|
||||
@@ -1,524 +1,139 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from pydantic import ValidationError
|
||||
|
||||
from app.models.template import ProxyTemplateDocument
|
||||
|
||||
DOCS_TEMPLATES_URL = "https://docs.mitmproxy.org/stable/addons/examples/"
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SUPPORTED_FILE_TEMPLATE_PROTOCOLS = frozenset({"http", "tcp"})
|
||||
TEMPLATES_DIR = Path(__file__).resolve().parents[4] / "templates"
|
||||
|
||||
|
||||
def builtin_template(
|
||||
*,
|
||||
id: str,
|
||||
name: str,
|
||||
protocol: str,
|
||||
family: str,
|
||||
tags: list[str],
|
||||
action: str,
|
||||
summary: str,
|
||||
behavior: str,
|
||||
code: str,
|
||||
) -> ProxyTemplateDocument:
|
||||
category = protocol if protocol in {"http", "tcp"} else family
|
||||
return ProxyTemplateDocument(
|
||||
id=id,
|
||||
name=name,
|
||||
description=summary,
|
||||
category=category,
|
||||
protocol=protocol,
|
||||
family=family,
|
||||
tags=tags,
|
||||
action=action,
|
||||
summary=summary,
|
||||
behavior=behavior,
|
||||
source_url=DOCS_TEMPLATES_URL,
|
||||
built_in=True,
|
||||
code=code,
|
||||
def template_directory() -> Path:
|
||||
return TEMPLATES_DIR
|
||||
|
||||
|
||||
def template_files(directory: Path | None = None) -> list[Path]:
|
||||
root = directory or template_directory()
|
||||
if not root.exists():
|
||||
return []
|
||||
return sorted(
|
||||
path
|
||||
for path in root.iterdir()
|
||||
if path.is_file() and path.suffix.lower() in {".yaml", ".yml"}
|
||||
)
|
||||
|
||||
|
||||
BUILTIN_TEMPLATES = [
|
||||
builtin_template(
|
||||
id="http-block-path-traversal",
|
||||
name="Block path traversal",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["path-traversal", "web", "block"],
|
||||
action="block",
|
||||
summary="Stops common path traversal probes in path and query.",
|
||||
behavior=(
|
||||
"Logs the matched traversal token and immediately drops the client "
|
||||
"connection before the request reaches the upstream service."
|
||||
),
|
||||
code="""import logging
|
||||
from urllib.parse import unquote
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
TRAVERSAL_MARKERS = (
|
||||
"../",
|
||||
"..\\\\",
|
||||
"%2e%2e",
|
||||
"%252e%252e",
|
||||
"..%2f",
|
||||
"..%5c",
|
||||
)
|
||||
def _default_category(protocol: str, family: str) -> str:
|
||||
return protocol if protocol in SUPPORTED_FILE_TEMPLATE_PROTOCOLS else family
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
raw_target = (
|
||||
f"{flow.request.path}?{flow.request.query}"
|
||||
if flow.request.query
|
||||
else flow.request.path
|
||||
def _string_value(value: object, default: str = "") -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if value is None:
|
||||
return default
|
||||
return str(value)
|
||||
|
||||
|
||||
def _normalize_template_payload(
|
||||
raw: Mapping[str, object],
|
||||
source: Path,
|
||||
) -> tuple[dict[str, object], str]:
|
||||
protocol = _string_value(raw.get("protocol")).strip().lower()
|
||||
if protocol not in SUPPORTED_FILE_TEMPLATE_PROTOCOLS:
|
||||
expected = sorted(SUPPORTED_FILE_TEMPLATE_PROTOCOLS)
|
||||
raise ValueError(
|
||||
f"unsupported protocol {protocol!r}; expected one of {expected}"
|
||||
)
|
||||
|
||||
if "id" not in raw or "name" not in raw or "code" not in raw:
|
||||
raise ValueError("missing one of required fields: id, name, code")
|
||||
|
||||
description = _string_value(
|
||||
raw.get("description") or raw.get("summary") or raw.get("behavior")
|
||||
)
|
||||
normalized = unquote(unquote(raw_target)).lower()
|
||||
for marker in TRAVERSAL_MARKERS:
|
||||
if marker in normalized:
|
||||
logging.warning(
|
||||
"[http-block-path-traversal] blocked %s %s from %s",
|
||||
flow.request.method,
|
||||
flow.request.pretty_url,
|
||||
flow.client_conn.peername,
|
||||
)
|
||||
flow.kill()
|
||||
return
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-block-blacklisted-substrings",
|
||||
name="Block blacklisted substrings",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["blacklist-string", "web", "block"],
|
||||
action="block",
|
||||
summary="Blocks requests containing forbidden strings in path, query, or body.",
|
||||
behavior=(
|
||||
"Scans URL and request body for blacklisted markers, then serves a "
|
||||
"blank HTML page with 200 OK and logs the first hit."
|
||||
family = _string_value(raw.get("family") or protocol)
|
||||
source_url = raw.get("source_url")
|
||||
tags = raw.get("tags")
|
||||
normalized_tags = [str(tag) for tag in tags] if isinstance(tags, list) else []
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"id": _string_value(raw["id"]).strip(),
|
||||
"name": _string_value(raw["name"]).strip(),
|
||||
"description": description,
|
||||
"category": _string_value(
|
||||
raw.get("category") or _default_category(protocol, family)
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
BLOCKED_STRINGS = (
|
||||
"flag{",
|
||||
"give_me_shell",
|
||||
"debug=1",
|
||||
"/etc/passwd",
|
||||
)
|
||||
"protocol": protocol,
|
||||
"family": family,
|
||||
"tags": normalized_tags,
|
||||
"action": _string_value(raw.get("action") or "custom"),
|
||||
"summary": _string_value(raw.get("summary") or description),
|
||||
"behavior": _string_value(raw.get("behavior") or description),
|
||||
"code": _string_value(raw["code"]),
|
||||
"source_url": _string_value(source_url) if source_url is not None else None,
|
||||
"built_in": True,
|
||||
"origin": "file",
|
||||
}
|
||||
return payload, str(source)
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
body = flow.request.get_text(strict=False)
|
||||
haystack = "\\n".join(
|
||||
[flow.request.path, str(flow.request.query), body]
|
||||
).lower()
|
||||
for marker in BLOCKED_STRINGS:
|
||||
if marker.lower() in haystack:
|
||||
logging.warning(
|
||||
"[http-block-blacklisted-substrings] blocked marker=%r url=%s",
|
||||
marker,
|
||||
flow.request.pretty_url,
|
||||
)
|
||||
flow.response = http.Response.make(
|
||||
200,
|
||||
b"",
|
||||
{"content-type": "text/html; charset=utf-8"},
|
||||
)
|
||||
return
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-block-large-body",
|
||||
name="Block large request bodies",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["length-limit", "body-size", "web", "block"],
|
||||
action="block",
|
||||
summary="Rejects oversized HTTP request bodies before they hit the service.",
|
||||
behavior=(
|
||||
"Checks the buffered request body length, returns 413 Payload Too "
|
||||
"Large, and logs the observed size."
|
||||
),
|
||||
code="""import logging
|
||||
def _load_template_from_file(path: Path) -> ProxyTemplateDocument | None:
|
||||
try:
|
||||
raw = yaml.safe_load(path.read_text(encoding="utf-8"))
|
||||
except (OSError, yaml.YAMLError) as error:
|
||||
logger.warning("Skipping template file %s: %s", path, error)
|
||||
return None
|
||||
|
||||
from mitmproxy import http
|
||||
if not isinstance(raw, dict):
|
||||
logger.warning("Skipping template file %s: expected mapping root", path)
|
||||
return None
|
||||
normalized_raw = {str(key): value for key, value in raw.items()}
|
||||
|
||||
MAX_BODY_BYTES = 8192
|
||||
try:
|
||||
payload, source_path = _normalize_template_payload(normalized_raw, path)
|
||||
document = ProxyTemplateDocument.model_validate(payload)
|
||||
except (ValidationError, ValueError) as error:
|
||||
logger.warning("Skipping template file %s: %s", path, error)
|
||||
return None
|
||||
|
||||
document.source_path = source_path
|
||||
return document
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
size = len(flow.request.raw_content or b"")
|
||||
if size > MAX_BODY_BYTES:
|
||||
logging.warning(
|
||||
"[http-block-large-body] blocked %s bytes on %s",
|
||||
size,
|
||||
flow.request.pretty_url,
|
||||
)
|
||||
flow.response = http.Response.make(
|
||||
413,
|
||||
b"",
|
||||
{"content-type": "text/plain"},
|
||||
)
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-block-suspicious-methods",
|
||||
name="Block suspicious HTTP methods",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["method-filter", "web", "block"],
|
||||
action="block",
|
||||
summary="Denies uncommon or dangerous HTTP methods.",
|
||||
behavior=(
|
||||
"Returns 405 Method Not Allowed for methods outside the allowlist "
|
||||
"and logs the blocked method and target URL."
|
||||
),
|
||||
code="""import logging
|
||||
def list_file_templates(directory: Path | None = None) -> list[ProxyTemplateDocument]:
|
||||
by_id: dict[str, list[ProxyTemplateDocument]] = defaultdict(list)
|
||||
for path in template_files(directory):
|
||||
template = _load_template_from_file(path)
|
||||
if template is None:
|
||||
continue
|
||||
by_id[template.id].append(template)
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
ALLOWED_METHODS = {"GET", "POST", "HEAD"}
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
if flow.request.method not in ALLOWED_METHODS:
|
||||
logging.warning(
|
||||
"[http-block-suspicious-methods] blocked method=%s url=%s",
|
||||
flow.request.method,
|
||||
flow.request.pretty_url,
|
||||
)
|
||||
flow.response = http.Response.make(
|
||||
405,
|
||||
b"",
|
||||
{"allow": "GET, POST, HEAD"},
|
||||
)
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-block-admin-paths",
|
||||
name="Block admin and debug paths",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["admin-path", "debug-path", "web", "block"],
|
||||
action="block",
|
||||
summary="Protects common admin, console, and debug endpoints.",
|
||||
behavior=(
|
||||
"Matches the request path against a small denylist and returns a "
|
||||
"blank 404 page before the request is forwarded."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
BLOCKED_PREFIXES = (
|
||||
"/admin",
|
||||
"/debug",
|
||||
"/console",
|
||||
"/actuator",
|
||||
)
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
if flow.request.path.startswith(BLOCKED_PREFIXES):
|
||||
logging.warning(
|
||||
"[http-block-admin-paths] blocked path=%s",
|
||||
flow.request.path,
|
||||
)
|
||||
flow.response = http.Response.make(
|
||||
404,
|
||||
b"",
|
||||
{"content-type": "text/html; charset=utf-8"},
|
||||
)
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-sanitize-headers-basic",
|
||||
name="Sanitize risky forwarding headers",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["rewrite", "headers", "web"],
|
||||
action="rewrite",
|
||||
summary=(
|
||||
"Removes a few risky forwarding headers before the request is "
|
||||
"sent upstream."
|
||||
),
|
||||
behavior=(
|
||||
"Deletes selected override or forwarding headers and logs what was "
|
||||
"rewritten without blocking the request."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
STRIP_HEADERS = (
|
||||
"x-forwarded-host",
|
||||
"x-http-method-override",
|
||||
"x-original-url",
|
||||
"x-rewrite-url",
|
||||
)
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
removed = []
|
||||
for header in STRIP_HEADERS:
|
||||
if header in flow.request.headers:
|
||||
removed.append(header)
|
||||
del flow.request.headers[header]
|
||||
if removed:
|
||||
logging.info(
|
||||
"[http-sanitize-headers-basic] stripped headers=%s url=%s",
|
||||
",".join(removed),
|
||||
flow.request.pretty_url,
|
||||
)
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-block-command-injection-markers",
|
||||
name="Block command injection markers",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["command-injection", "web", "block"],
|
||||
action="block",
|
||||
summary=(
|
||||
"Rejects obvious shell metacharacters and command-substitution "
|
||||
"markers."
|
||||
),
|
||||
behavior=(
|
||||
"Scans URL and body for shell-style injection markers, logs the "
|
||||
"first hit, and drops the client connection."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
BLOCKED_MARKERS = (";", "&&", "||", "$(", "`", "|sh", "|bash")
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
body = flow.request.get_text(strict=False)
|
||||
haystack = "\\n".join(
|
||||
[flow.request.pretty_url, body]
|
||||
).lower()
|
||||
for marker in BLOCKED_MARKERS:
|
||||
if marker in haystack:
|
||||
logging.warning(
|
||||
"[http-block-command-injection-markers] blocked marker=%r url=%s",
|
||||
marker,
|
||||
flow.request.pretty_url,
|
||||
)
|
||||
flow.kill()
|
||||
return
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="http-block-sqli-markers",
|
||||
name="Block obvious SQLi markers",
|
||||
protocol="http",
|
||||
family="web-hardening",
|
||||
tags=["sqli", "web", "block"],
|
||||
action="block",
|
||||
summary="Stops common SQL injection probes before they reach the app.",
|
||||
behavior=(
|
||||
"Matches a short list of high-signal SQLi markers in path, query, "
|
||||
"and body, then serves a blank HTML page and logs the hit."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import http
|
||||
|
||||
BLOCKED_MARKERS = (
|
||||
"' or 1=1",
|
||||
'" or 1=1',
|
||||
"union select",
|
||||
"sleep(",
|
||||
"information_schema",
|
||||
"--",
|
||||
)
|
||||
|
||||
|
||||
def request(flow: http.HTTPFlow) -> None:
|
||||
body = flow.request.get_text(strict=False)
|
||||
haystack = "\\n".join(
|
||||
[flow.request.path, str(flow.request.query), body]
|
||||
).lower()
|
||||
for marker in BLOCKED_MARKERS:
|
||||
if marker in haystack:
|
||||
logging.warning(
|
||||
"[http-block-sqli-markers] blocked marker=%r url=%s",
|
||||
marker,
|
||||
flow.request.pretty_url,
|
||||
)
|
||||
flow.response = http.Response.make(
|
||||
200,
|
||||
b"",
|
||||
{"content-type": "text/html; charset=utf-8"},
|
||||
)
|
||||
return
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="tcp-block-repeated-byte-threshold",
|
||||
name="Block repeated-byte flood",
|
||||
protocol="tcp",
|
||||
family="tcp-hardening",
|
||||
tags=["bof", "repeated-byte", "tcp", "block"],
|
||||
action="block",
|
||||
summary="Closes connections that send suspiciously long runs of the same byte.",
|
||||
behavior=(
|
||||
"Measures the longest repeated-byte run in each client message, logs "
|
||||
"the byte and run length, and closes the TCP flow on threshold hit."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import tcp
|
||||
|
||||
MAX_REPEAT = 128
|
||||
|
||||
|
||||
def longest_run(data: bytes) -> tuple[int, int]:
|
||||
best_byte = -1
|
||||
best_run = 0
|
||||
current_byte = -1
|
||||
current_run = 0
|
||||
for byte in data:
|
||||
if byte == current_byte:
|
||||
current_run += 1
|
||||
else:
|
||||
current_byte = byte
|
||||
current_run = 1
|
||||
if current_run > best_run:
|
||||
best_byte = byte
|
||||
best_run = current_run
|
||||
return best_byte, best_run
|
||||
|
||||
|
||||
def tcp_message(flow: tcp.TCPFlow) -> None:
|
||||
message = flow.messages[-1]
|
||||
if not message.from_client:
|
||||
return
|
||||
repeated_byte, run_length = longest_run(message.content)
|
||||
if run_length > MAX_REPEAT:
|
||||
logging.warning(
|
||||
"[tcp-block-repeated-byte-threshold] blocked byte=0x%02x run=%d",
|
||||
repeated_byte,
|
||||
run_length,
|
||||
)
|
||||
message.drop()
|
||||
flow.kill()
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="tcp-block-message-length-threshold",
|
||||
name="Block oversized TCP messages",
|
||||
protocol="tcp",
|
||||
family="tcp-hardening",
|
||||
tags=["length-limit", "tcp", "block"],
|
||||
action="block",
|
||||
summary="Cuts off client messages that exceed a fixed size threshold.",
|
||||
behavior=(
|
||||
"Checks every client TCP message size, logs the offending length, "
|
||||
"drops the message, and kills the flow when the threshold is hit."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import tcp
|
||||
|
||||
MAX_MESSAGE_BYTES = 4096
|
||||
|
||||
|
||||
def tcp_message(flow: tcp.TCPFlow) -> None:
|
||||
message = flow.messages[-1]
|
||||
if not message.from_client:
|
||||
return
|
||||
size = len(message.content)
|
||||
if size > MAX_MESSAGE_BYTES:
|
||||
logging.warning(
|
||||
"[tcp-block-message-length-threshold] blocked message size=%d",
|
||||
size,
|
||||
)
|
||||
message.drop()
|
||||
flow.kill()
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="tcp-block-blacklisted-byte-sequences",
|
||||
name="Block blacklisted TCP byte sequences",
|
||||
protocol="tcp",
|
||||
family="tcp-hardening",
|
||||
tags=["blacklist-bytes", "tcp", "block"],
|
||||
action="block",
|
||||
summary=(
|
||||
"Blocks payloads containing forbidden byte sequences or marker "
|
||||
"strings."
|
||||
),
|
||||
behavior=(
|
||||
"Looks for dangerous byte patterns in client messages, then drops "
|
||||
"the message and closes the flow."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import tcp
|
||||
|
||||
BLOCKED_SEQUENCES = (
|
||||
b"/bin/sh",
|
||||
b"cat flag",
|
||||
b"AAAA%p%p",
|
||||
b"\\x90\\x90\\x90\\x90",
|
||||
)
|
||||
|
||||
|
||||
def tcp_message(flow: tcp.TCPFlow) -> None:
|
||||
message = flow.messages[-1]
|
||||
if not message.from_client:
|
||||
return
|
||||
for sequence in BLOCKED_SEQUENCES:
|
||||
if sequence in message.content:
|
||||
logging.warning(
|
||||
"[tcp-block-blacklisted-byte-sequences] blocked sequence=%r",
|
||||
sequence,
|
||||
)
|
||||
message.drop()
|
||||
flow.kill()
|
||||
return
|
||||
""",
|
||||
),
|
||||
builtin_template(
|
||||
id="tcp-block-non-printable-threshold",
|
||||
name="Block binary-noise threshold",
|
||||
protocol="tcp",
|
||||
family="tcp-hardening",
|
||||
tags=["non-printable", "binary-noise", "tcp", "block"],
|
||||
action="block",
|
||||
summary=(
|
||||
"Blocks client messages with an excessive amount of non-printable "
|
||||
"bytes."
|
||||
),
|
||||
behavior=(
|
||||
"Treats unusually binary-heavy client payloads as suspicious, logs "
|
||||
"the ratio, and closes the flow when the threshold is exceeded."
|
||||
),
|
||||
code="""import logging
|
||||
|
||||
from mitmproxy import tcp
|
||||
|
||||
NON_PRINTABLE_RATIO = 0.35
|
||||
|
||||
|
||||
def tcp_message(flow: tcp.TCPFlow) -> None:
|
||||
message = flow.messages[-1]
|
||||
if not message.from_client or not message.content:
|
||||
return
|
||||
non_printable = sum(
|
||||
1
|
||||
for byte in message.content
|
||||
if byte < 0x20 and byte not in {0x09, 0x0a, 0x0d}
|
||||
deduplicated: list[ProxyTemplateDocument] = []
|
||||
duplicate_template_message = (
|
||||
"Skipping template id %s from %s: duplicate file-backed template id"
|
||||
)
|
||||
ratio = non_printable / len(message.content)
|
||||
if ratio > NON_PRINTABLE_RATIO:
|
||||
logging.warning(
|
||||
"[tcp-block-non-printable-threshold] blocked ratio=%.2f size=%d",
|
||||
ratio,
|
||||
len(message.content),
|
||||
)
|
||||
message.drop()
|
||||
flow.kill()
|
||||
""",
|
||||
),
|
||||
]
|
||||
for template_id, templates in by_id.items():
|
||||
if len(templates) > 1:
|
||||
for template in templates:
|
||||
logger.warning(
|
||||
duplicate_template_message,
|
||||
template_id,
|
||||
template.source_path or "<unknown>",
|
||||
)
|
||||
continue
|
||||
deduplicated.append(templates[0])
|
||||
|
||||
return sorted(deduplicated, key=lambda template: template.name.lower())
|
||||
|
||||
|
||||
def file_template_ids(directory: Path | None = None) -> set[str]:
|
||||
return {template.id for template in list_file_templates(directory)}
|
||||
|
||||
@@ -9,7 +9,7 @@ from app.core.security import get_auth_context, require_authenticated_user
|
||||
from app.models.health import HealthSnapshot
|
||||
from app.models.proxy import ProxyCreateInput, ProxyDocument
|
||||
from app.models.service import ServiceEndpointDocument
|
||||
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
|
||||
from app.services.mitmproxy.templates import list_file_templates
|
||||
from app.services.mitmproxy.validator import ProxyCodeValidationError
|
||||
|
||||
proxy_router_module = importlib.import_module("app.api.proxy.router")
|
||||
@@ -864,6 +864,8 @@ def test_stop_proxy_updates_runtime_state(client, monkeypatch: object):
|
||||
|
||||
|
||||
def test_list_mitmproxy_templates_returns_builtins_and_custom(client, monkeypatch):
|
||||
file_templates = list_file_templates()
|
||||
|
||||
async def fake_list_custom_templates():
|
||||
return []
|
||||
|
||||
@@ -872,13 +874,19 @@ def test_list_mitmproxy_templates_returns_builtins_and_custom(client, monkeypatc
|
||||
"list_custom_templates",
|
||||
fake_list_custom_templates,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"list_file_templates",
|
||||
lambda: file_templates,
|
||||
)
|
||||
|
||||
response = client.get("/proxy/mitmproxy/templates")
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert len(body) == len(BUILTIN_TEMPLATES)
|
||||
assert len(body) == len(file_templates)
|
||||
assert body[0]["built_in"] is True
|
||||
assert body[0]["origin"] == "file"
|
||||
assert body[0]["protocol"] in {"http", "tcp"}
|
||||
assert isinstance(body[0]["tags"], list)
|
||||
assert body[0]["action"] in {"block", "rewrite"}
|
||||
@@ -906,6 +914,7 @@ def test_create_mitmproxy_template_validates_and_persists(client, monkeypatch):
|
||||
"code": payload.code,
|
||||
"source_url": payload.source_url,
|
||||
"built_in": False,
|
||||
"origin": "custom",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
@@ -939,4 +948,33 @@ def test_create_mitmproxy_template_validates_and_persists(client, monkeypatch):
|
||||
assert response.status_code == 201
|
||||
assert response.json()["id"] == "custom-header"
|
||||
assert response.json()["protocol"] == "http"
|
||||
assert response.json()["origin"] == "custom"
|
||||
assert captured["validated"] == "def response(flow):\n pass\n"
|
||||
|
||||
|
||||
def test_create_mitmproxy_template_rejects_file_backed_id(client, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
proxy_router_module,
|
||||
"file_template_ids",
|
||||
lambda: {"http-block-path-traversal"},
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/proxy/mitmproxy/templates",
|
||||
json={
|
||||
"id": "http-block-path-traversal",
|
||||
"name": "Custom header",
|
||||
"description": "Adds a custom response header",
|
||||
"category": "http",
|
||||
"protocol": "http",
|
||||
"family": "custom",
|
||||
"tags": [],
|
||||
"action": "custom",
|
||||
"summary": "Saved custom template.",
|
||||
"behavior": "Applies the current script as-is.",
|
||||
"code": "def response(flow):\n pass\n",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 409
|
||||
assert "Template already exists" in response.json()["detail"]
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
from app.services.mitmproxy.templates import BUILTIN_TEMPLATES
|
||||
import logging
|
||||
|
||||
import yaml
|
||||
|
||||
from app.services.mitmproxy.templates import list_file_templates
|
||||
|
||||
|
||||
def test_builtin_templates_use_four_space_indentation():
|
||||
for template in BUILTIN_TEMPLATES:
|
||||
def test_file_templates_use_four_space_indentation():
|
||||
for template in list_file_templates():
|
||||
for line in template.code.splitlines():
|
||||
assert "\t" not in line, f"{template.id} contains a tab: {line!r}"
|
||||
if not line.strip():
|
||||
@@ -11,3 +15,42 @@ def test_builtin_templates_use_four_space_indentation():
|
||||
assert indent % 4 == 0, (
|
||||
f"{template.id} has non-4-space indentation: {line!r}"
|
||||
)
|
||||
|
||||
|
||||
def test_list_file_templates_loads_yaml_backed_templates():
|
||||
templates = list_file_templates()
|
||||
|
||||
assert templates
|
||||
assert all(template.origin == "file" for template in templates)
|
||||
assert all(template.built_in is True for template in templates)
|
||||
assert all(template.protocol in {"http", "tcp"} for template in templates)
|
||||
|
||||
|
||||
def test_list_file_templates_skips_invalid_files(tmp_path, caplog):
|
||||
path = tmp_path / "invalid.yaml"
|
||||
path.write_text("[]\n", encoding="utf-8")
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
templates = list_file_templates(tmp_path)
|
||||
|
||||
assert templates == []
|
||||
assert "Skipping template file" in caplog.text
|
||||
|
||||
|
||||
def test_list_file_templates_skips_duplicate_ids(tmp_path, caplog):
|
||||
payload = {
|
||||
"id": "duplicate",
|
||||
"name": "Duplicate",
|
||||
"protocol": "http",
|
||||
"description": "Duplicate template",
|
||||
"code": "from mitmproxy import http\n",
|
||||
}
|
||||
document = yaml.safe_dump(payload, sort_keys=False)
|
||||
(tmp_path / "one.yaml").write_text(document, encoding="utf-8")
|
||||
(tmp_path / "two.yaml").write_text(document, encoding="utf-8")
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
templates = list_file_templates(tmp_path)
|
||||
|
||||
assert templates == []
|
||||
assert "duplicate file-backed template id" in caplog.text
|
||||
|
||||
Reference in New Issue
Block a user