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,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,
|
||||
)
|
||||
Reference in New Issue
Block a user