Files
thePlug/backend/app/services/mitmproxy/templates.py
T

140 lines
4.6 KiB
Python

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
logger = logging.getLogger(__name__)
SUPPORTED_FILE_TEMPLATE_PROTOCOLS = frozenset({"http", "tcp"})
TEMPLATES_DIR = Path(__file__).resolve().parents[4] / "templates"
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"}
)
def _default_category(protocol: str, family: str) -> str:
return protocol if protocol in SUPPORTED_FILE_TEMPLATE_PROTOCOLS else family
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")
)
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)
),
"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 _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
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()}
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 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)
deduplicated: list[ProxyTemplateDocument] = []
duplicate_template_message = (
"Skipping template id %s from %s: duplicate file-backed template id"
)
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)}