wehub-resource-sync
d0e4308def
Validate YAML Workflows / Validate YAML Configuration Files (push) Waiting to run
132 lines
5.2 KiB
Python
Executable File
132 lines
5.2 KiB
Python
Executable File
"""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
|