code-muse 0.0.1__py3-none-any.whl

code_muse/tools/image_tools.py

@@ -0,0 +1,185 @@
+"""Image loading tool for visual analysis.
+
+Provides ``load_image_for_analysis`` — a tiny generic tool that reads an image
+from disk, validates it, downsizes it if it's billboard-sized, and returns it
+as a ``ToolReturn`` with ``BinaryContent`` so multimodal models can see it.
+
+Lives outside the browser package because it has nothing to do with browsers.
+"""
+
+import io
+import logging
+import mimetypes
+import time
+from pathlib import Path
+from typing import Any
+
+from PIL import Image, UnidentifiedImageError
+from pydantic_ai import BinaryContent, RunContext, ToolReturn
+
+from code_muse.messaging import emit_error, emit_info, emit_success
+from code_muse.tools.common import generate_group_id
+
+logger = logging.getLogger(__name__)
+
+# Bigger than this on either edge and we resize to save tokens.
+MAX_IMAGE_EDGE = 2048
+DEFAULT_MAX_HEIGHT = 768  # kept for backward-compat in tool signature
+
+
+def _validate_and_prepare_image(
+    image_bytes: bytes,
+    source_path: str | None = None,
+    max_edge: int | None = None,
+) -> dict[str, Any]:
+    """Verify image bytes, determine the real MIME type, and optionally resize.
+
+    The MIME type is determined from the decoded image content, not from the
+    file extension. If the image is resized, the output is normalized to PNG so
+    the returned bytes and MIME type stay in sync like civilized software.
+    """
+    try:
+        with Image.open(io.BytesIO(image_bytes)) as verified_image:
+            verified_image.verify()
+    except UnidentifiedImageError as exc:
+        raise ValueError("File is not a valid image") from exc
+    except Exception as exc:
+        raise ValueError(f"Failed to verify image: {exc}") from exc
+
+    with Image.open(io.BytesIO(image_bytes)) as image:
+        image.load()
+        original_width, original_height = image.size
+        image_format = image.format
+        actual_media_type = Image.MIME.get(image_format or "")
+
+        if not actual_media_type or not actual_media_type.startswith("image/"):
+            raise ValueError("Could not determine a valid image MIME type")
+
+        guessed_media_type = None
+        if source_path:
+            guessed_media_type, _ = mimetypes.guess_type(source_path)
+
+        largest_edge = max(original_width, original_height)
+        was_resized = False
+        output_bytes = image_bytes
+        output_media_type = actual_media_type
+        output_width = original_width
+        output_height = original_height
+
+        if max_edge and largest_edge > max_edge:
+            ratio = max_edge / largest_edge
+            output_width = max(1, int(round(original_width * ratio)))
+            output_height = max(1, int(round(original_height * ratio)))
+            resized = image.resize(
+                (output_width, output_height), Image.Resampling.LANCZOS
+            )
+            output = io.BytesIO()
+            resized.save(output, format="PNG", optimize=True)
+            output.seek(0)
+            output_bytes = output.read()
+            output_media_type = "image/png"
+            was_resized = True
+
+        return {
+            "image_bytes": output_bytes,
+            "media_type": output_media_type,
+            "actual_media_type": actual_media_type,
+            "guessed_media_type": guessed_media_type,
+            "mime_type_matches_extension": guessed_media_type
+            in (None, output_media_type),
+            "original_width": original_width,
+            "original_height": original_height,
+            "output_width": output_width,
+            "output_height": output_height,
+            "was_resized": was_resized,
+        }
+
+
+async def load_image(
+    image_path: str,
+    max_height: int = DEFAULT_MAX_HEIGHT,
+) -> ToolReturn | dict[str, Any]:
+    """Load an image from the filesystem for visual analysis."""
+    group_id = generate_group_id("load_image", image_path)
+    emit_info(f"LOAD IMAGE 🖼️ {image_path}", message_group=group_id)
+
+    try:
+        image_file = Path(image_path)
+
+        if not image_file.exists():
+            error_msg = f"Image file not found: {image_path}"
+            emit_error(error_msg, message_group=group_id)
+            return {"success": False, "error": error_msg, "image_path": image_path}
+
+        if not image_file.is_file():
+            error_msg = f"Path is not a file: {image_path}"
+            emit_error(error_msg, message_group=group_id)
+            return {"success": False, "error": error_msg, "image_path": image_path}
+
+        image_bytes = image_file.read_bytes()
+        prepared_image = _validate_and_prepare_image(
+            image_bytes,
+            source_path=str(image_file),
+            max_edge=MAX_IMAGE_EDGE,
+        )
+
+        emit_success(f"Loaded image: {image_path}", message_group=group_id)
+
+        return ToolReturn(
+            return_value=f"Image loaded from: {image_path}",
+            content=[
+                f"Here's the image from {image_file.name}:",
+                BinaryContent(
+                    data=prepared_image["image_bytes"],
+                    media_type=prepared_image["media_type"],
+                ),
+                "Please analyze what you see in this image.",
+            ],
+            metadata={
+                "success": True,
+                "image_path": image_path,
+                "media_type": prepared_image["media_type"],
+                "actual_media_type": prepared_image["actual_media_type"],
+                "guessed_media_type": prepared_image["guessed_media_type"],
+                "mime_type_matches_extension": prepared_image[
+                    "mime_type_matches_extension"
+                ],
+                "was_resized": prepared_image["was_resized"],
+                "original_size": [
+                    prepared_image["original_width"],
+                    prepared_image["original_height"],
+                ],
+                "output_size": [
+                    prepared_image["output_width"],
+                    prepared_image["output_height"],
+                ],
+                "max_height": max_height,
+                "max_edge": MAX_IMAGE_EDGE,
+                "timestamp": time.time(),
+            },
+        )
+
+    except Exception as e:
+        error_msg = f"Failed to load image: {str(e)}"
+        emit_error(error_msg, message_group=group_id)
+        logger.exception("Error loading image")
+        return {"success": False, "error": error_msg, "image_path": image_path}
+
+
+def register_load_image(agent):
+    """Register the image loading tool."""
+
+    @agent.tool
+    async def load_image_for_analysis(
+        context: RunContext,
+        image_path: str,
+    ) -> ToolReturn | dict[str, Any]:
+        """Load an image file so you can see and analyze it.
+
+        Args:
+            image_path: Path to the image file.
+
+        Returns:
+            ToolReturn with the image, or error dict.
+        """
+        return await load_image(image_path=image_path)

code_muse/tools/meetin_proxy/__init__.py

@@ -0,0 +1,243 @@
+"""mitmproxy — MITM proxy for provider traffic capture and analysis."""
+
+import json
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic_ai import RunContext
+
+from code_muse.messaging import emit_error, emit_info, emit_success
+from code_muse.tools.meetin_proxy.proxy_manager import get_manager
+
+
+class MitmproxyConfig(BaseModel):
+    """Configuration for mitmproxy capture session."""
+
+    target_domain: str = ""
+    listen_port: int | None = None
+    max_req_body: int = 10000
+    max_res_body: int = 100000
+
+
+class MitmproxyResult(BaseModel):
+    """Result of a mitmproxy capture session."""
+
+    proxy_port: int
+    flow_count: int
+    output_path: str
+    captured_at: str
+
+
+class MitmproxyStatus(BaseModel):
+    """Current status of the mitmproxy instance."""
+
+    running: bool
+    port: int | None = None
+    target: str = ""
+    flow_count: int = 0
+
+
+def register_mitmproxy(agent):
+    """Register the mitmproxy tool."""
+
+    @agent.tool
+    async def mitmproxy(
+        context: RunContext,
+        command: str,
+        target_domain: str = "",
+        listen_port: int | None = None,
+        max_req_body: int = 10000,
+        max_res_body: int = 100000,
+        duration_seconds: int | None = None,
+    ) -> MitmproxyResult | MitmproxyStatus | dict[str, Any]:
+        """Control the mitmproxy MITM capture proxy.
+
+        Commands:
+            start   — Spawn mitmdump, set proxy env vars, begin capture.
+            stop    — Stop mitmdump, read flows, unset env vars, return results.
+            status  — Return current proxy status.
+            capture — Convenience: start, optionally wait, then stop and return data.
+
+        Args:
+            command: Action to perform — "start", "stop", "status", or "capture".
+            target_domain: Only capture traffic matching this domain.
+            listen_port: Port to listen on (auto-selected if omitted).
+            max_req_body: Maximum request body bytes to record.
+            max_res_body: Maximum response body bytes to record.
+            duration_seconds: For "capture", seconds to wait before auto-stopping.
+
+        Returns:
+            MitmproxyResult on stop/capture, MitmproxyStatus on status/start.
+        """
+        manager = get_manager()
+        config = MitmproxyConfig(
+            target_domain=target_domain,
+            listen_port=listen_port,
+            max_req_body=max_req_body,
+            max_res_body=max_res_body,
+        )
+
+        if command == "start":
+            try:
+                port = manager.start(
+                    target=config.target_domain,
+                    port=config.listen_port,
+                    max_req_body=config.max_req_body,
+                    max_res_body=config.max_res_body,
+                )
+                return MitmproxyStatus(
+                    running=True,
+                    port=port,
+                    target=config.target_domain,
+                    flow_count=0,
+                )
+            except Exception as exc:
+                emit_error(f"Failed to start mitmproxy: {exc}")
+                return {
+                    "success": False,
+                    "error": str(exc),
+                    "command": command,
+                }
+
+        elif command == "stop":
+            data = manager.stop()
+            status = manager.status()
+            # After stop, status will show not running
+            flow_count = 0
+            output_path = ""
+            captured_at = ""
+            if data and isinstance(data, dict):
+                meta = data.get("meta", {})
+                flow_count = meta.get("total_flows", 0)
+                captured_at = meta.get("captured_at", "")
+                # Write a safe temp copy so caller can inspect flows
+                import tempfile
+
+                fd, output_path = tempfile.mkstemp(
+                    prefix="mitmproxy_result_", suffix=".json"
+                )
+                with open(fd, "w", encoding="utf-8") as f:
+                    json.dump(data, f, indent=2, default=str)
+
+            if data:
+                emit_success(
+                    f"mitmproxy stopped. Captured {flow_count} flow(s). "
+                    f"Data written to {output_path}"
+                )
+                return MitmproxyResult(
+                    proxy_port=status.get("port") or 0,
+                    flow_count=flow_count,
+                    output_path=output_path,
+                    captured_at=captured_at,
+                )
+            else:
+                emit_info("mitmproxy stopped. No capture data found.")
+                return MitmproxyResult(
+                    proxy_port=status.get("port") or 0,
+                    flow_count=0,
+                    output_path="",
+                    captured_at="",
+                )
+
+        elif command == "status":
+            st = manager.status()
+            return MitmproxyStatus(
+                running=st["running"],
+                port=st["port"],
+                target=st["target"],
+                flow_count=st["flow_count"],
+            )
+
+        elif command == "capture":
+            try:
+                port = manager.start(
+                    target=config.target_domain,
+                    port=config.listen_port,
+                    max_req_body=config.max_req_body,
+                    max_res_body=config.max_res_body,
+                )
+            except Exception as exc:
+                emit_error(f"Failed to start mitmproxy: {exc}")
+                return {
+                    "success": False,
+                    "error": str(exc),
+                    "command": command,
+                }
+
+            if duration_seconds and duration_seconds > 0:
+                emit_info(
+                    f"mitmproxy capturing for {duration_seconds}s on port {port}..."
+                )
+                await _async_sleep(duration_seconds)
+            else:
+                emit_info(
+                    f"mitmproxy capturing on port {port}. "
+                    "Call mitmproxy with command='stop' when finished."
+                )
+                return MitmproxyStatus(
+                    running=True,
+                    port=port,
+                    target=config.target_domain,
+                    flow_count=0,
+                )
+
+            data = manager.stop()
+            flow_count = 0
+            output_path = ""
+            captured_at = ""
+            if data and isinstance(data, dict):
+                meta = data.get("meta", {})
+                flow_count = meta.get("total_flows", 0)
+                captured_at = meta.get("captured_at", "")
+                import tempfile
+
+                fd, output_path = tempfile.mkstemp(
+                    prefix="mitmproxy_result_", suffix=".json"
+                )
+                with open(fd, "w", encoding="utf-8") as f:
+                    json.dump(data, f, indent=2, default=str)
+
+            if data:
+                emit_success(
+                    f"mitmproxy capture complete. {flow_count} flow(s). "
+                    f"Data: {output_path}"
+                )
+                return MitmproxyResult(
+                    proxy_port=port,
+                    flow_count=flow_count,
+                    output_path=output_path,
+                    captured_at=captured_at,
+                )
+            else:
+                emit_info("mitmproxy capture complete. No data captured.")
+                return MitmproxyResult(
+                    proxy_port=port,
+                    flow_count=0,
+                    output_path="",
+                    captured_at="",
+                )
+
+        else:
+            return {
+                "success": False,
+                "error": (
+                    f"Unknown command: {command}. Use start, stop, status, or capture."
+                ),
+            }
+
+    return mitmproxy
+
+
+async def _async_sleep(seconds: float) -> None:
+    """Async sleep helper."""
+    import asyncio
+
+    await asyncio.sleep(seconds)
+
+
+# Register startup callback for mitmproxy availability check
+from code_muse.tools.meetin_proxy.register_callbacks import (  # noqa: E402
+    register as _register_mitmproxy_callbacks,
+)
+
+_register_mitmproxy_callbacks()

code_muse/tools/meetin_proxy/capture_addon.py

@@ -0,0 +1,82 @@
+"""mitmproxy capture addon — injected into mitmdump via -s flag."""
+
+import json
+import os
+from datetime import UTC, datetime
+
+
+class TrafficCapture:
+    def __init__(self):
+        self.flows = []
+        self.target = os.environ.get("MITMPROXY_TARGET", "")
+        self.output = os.environ.get("MITMPROXY_OUTPUT", "/tmp/mitmproxy_capture.json")
+        self.max_req_body = int(os.environ.get("MITMPROXY_MAX_REQ_BODY", "10000"))
+        self.max_res_body = int(os.environ.get("MITMPROXY_MAX_RES_BODY", "100000"))
+
+    def _should_capture(self, flow) -> bool:
+        if not self.target:
+            return True
+        return self.target in flow.request.pretty_host
+
+    def response(self, flow):
+        if not self._should_capture(flow):
+            return
+        try:
+            req_body = flow.request.get_text(strict=False) or ""
+            res_body = flow.response.get_text(strict=False) or ""
+        except Exception:
+            req_body = f"<binary: {len(flow.request.content)} bytes>"
+            res_body = f"<binary: {len(flow.response.content)} bytes>"
+
+        if len(req_body) > self.max_req_body:
+            req_body = req_body[: self.max_req_body] + "\n... <truncated>"
+        if len(res_body) > self.max_res_body:
+            res_body = res_body[: self.max_res_body] + "\n... <truncated>"
+
+        entry = {
+            "url": flow.request.pretty_url,
+            "method": flow.request.method,
+            "host": flow.request.pretty_host,
+            "path": flow.request.path,
+            "request_headers": dict(flow.request.headers),
+            "request_body": req_body,
+            "status_code": flow.response.status_code,
+            "response_headers": dict(flow.response.headers),
+            "response_body": res_body,
+            "timestamp": datetime.now(UTC).isoformat(),
+            "content_type": flow.response.headers.get("content-type", ""),
+        }
+        self.flows.append(entry)
+
+    def error(self, flow):
+        if not self._should_capture(flow):
+            return
+        entry = {
+            "url": flow.request.pretty_url if flow.request else "N/A",
+            "method": flow.request.method if flow.request else "N/A",
+            "error": str(flow.error.msg) if flow.error else "Unknown error",
+            "timestamp": datetime.now(UTC).isoformat(),
+        }
+        self.flows.append(entry)
+
+    def done(self):
+        out_dir = os.path.dirname(self.output)
+        if out_dir:
+            os.makedirs(out_dir, exist_ok=True)
+        with open(self.output, "w") as f:
+            json.dump(
+                {
+                    "meta": {
+                        "captured_at": datetime.now(UTC).isoformat(),
+                        "target_filter": self.target,
+                        "total_flows": len(self.flows),
+                    },
+                    "flows": self.flows,
+                },
+                f,
+                indent=2,
+                default=str,
+            )
+
+
+addons = [TrafficCapture()]