chore: import upstream snapshot with attribution

2026-07-13 12:37:51 +08:00
commit d0e4308def
614 changed files with 74458 additions and 0 deletions
@@ -0,0 +1,131 @@
+"""Attachment helpers shared by HTTP routes and executors."""
+
+import logging
+import mimetypes
+import os
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from fastapi import UploadFile
+
+from entity.messages import MessageBlock, MessageBlockType
+from utils.attachments import AttachmentStore, AttachmentRecord
+
+
+class AttachmentService:
+    """Handles attachment lifecycle per session."""
+
+    def __init__(self, *, root: Path | str = Path("WareHouse")) -> None:
+        self.logger = logging.getLogger(__name__)
+        self.attachments_root = Path(root)
+        self.attachments_root.mkdir(parents=True, exist_ok=True)
+        env_flag = os.environ.get("MAC_AUTO_CLEAN_ATTACHMENTS", "0").strip().lower()
+        self.clean_on_cleanup = env_flag in {"1", "true", "yes"}
+
+    def prepare_session_workspace(self, session_id: str) -> Path:
+        return self._session_attachments_path(session_id, create=True)
+
+    def cleanup_session(self, session_id: str) -> None:
+        attachment_dir = self._session_attachments_path(session_id, create=False)
+        if not attachment_dir:
+            return
+        if self.clean_on_cleanup:
+            shutil.rmtree(attachment_dir, ignore_errors=True)
+            self.logger.info("Cleaned attachment directory for session %s", session_id)
+        else:
+            self.logger.info(
+                "Attachment cleanup disabled; preserved files for session %s", session_id
+            )
+
+    def get_attachment_store(self, session_id: str) -> AttachmentStore:
+        path = self.prepare_session_workspace(session_id)
+        return AttachmentStore(path)
+
+    @staticmethod
+    def _safe_upload_filename(raw: Optional[str]) -> str:
+        """Reduce a client-supplied upload filename to a safe basename.
+
+        The multipart ``filename`` is attacker-controlled. Joining it onto the
+        temporary upload directory verbatim allows path traversal (e.g.
+        ``../../../etc/cron.d/x``): the write target escapes the temp dir and the
+        cleanup step then unlinks the same traversed path, yielding arbitrary
+        file write and delete. Normalise both POSIX and Windows separators and
+        keep only the final path component so the write stays confined.
+        """
+        candidate = os.path.basename((raw or "").replace("\\", "/")).strip()
+        if not candidate or candidate in {".", ".."}:
+            return "upload.bin"
+        return candidate
+
+    async def save_upload_file(self, session_id: str, upload: UploadFile) -> AttachmentRecord:
+        filename = self._safe_upload_filename(upload.filename)
+        temp_dir = Path(tempfile.mkdtemp(prefix="mac_upload_"))
+        temp_path = temp_dir / filename
+        try:
+            with temp_path.open("wb") as buffer:
+                while True:
+                    chunk = await upload.read(1024 * 1024)
+                    if not chunk:
+                        break
+                    buffer.write(chunk)
+            store = self.get_attachment_store(session_id)
+            mime_type = upload.content_type or mimetypes.guess_type(filename)[0]
+            record = store.register_file(
+                temp_path,
+                kind=MessageBlockType.from_mime_type(mime_type),
+                display_name=filename,
+                mime_type=mime_type,
+                extra={
+                    "source": "user_upload",
+                    "origin": "web_upload",
+                    "session_id": session_id,
+                },
+            )
+            return record
+        finally:
+            if temp_path.exists():
+                try:
+                    temp_path.unlink()
+                except OSError:
+                    pass
+            try:
+                temp_dir.rmdir()
+            except OSError:
+                pass
+
+    def build_attachment_blocks(
+        self,
+        session_id: str,
+        attachment_ids: List[str],
+        *,
+        target_store: Optional[AttachmentStore] = None,
+    ) -> List[MessageBlock]:
+        if not attachment_ids:
+            return []
+        source_store = self.get_attachment_store(session_id)
+        source_root = source_store.root.resolve()
+        target_root = target_store.root.resolve() if target_store else None
+        blocks: List[MessageBlock] = []
+        for attachment_id in attachment_ids:
+            record = source_store.get(attachment_id)
+            if not record:
+                continue
+            if target_store:
+                copy_required = target_root != source_root
+                record = target_store.ingest_record(record, copy_file=copy_required)
+            blocks.append(record.as_message_block())
+        return blocks
+
+    def list_attachment_manifests(self, session_id: str) -> Dict[str, Any]:
+        store = self.get_attachment_store(session_id)
+        return store.export_manifest()
+
+    def _session_attachments_path(self, session_id: str, *, create: bool = True) -> Optional[Path]:
+        session_dir_name = session_id if session_id.startswith("session_") else f"session_{session_id}"
+        path = self.attachments_root / session_dir_name / "code_workspace" / "attachments"
+        if create:
+            path.mkdir(parents=True, exist_ok=True)
+            return path
+        return path if path.exists() else None