ForcePilot/backend/package/yuxi/channels/adapters/msteams/file_upload.py

"""Microsoft Teams 文件上传增强。

FileConsentCard 大文件交互流程 + SharePoint 上传策略 + SSRF 三层防护。
"""

from __future__ import annotations

import ipaddress
import re
import time
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

from yuxi.utils.logging_config import logger

FILE_CONSENT_MAX_SIZE = 4 * 1024 * 1024

_SSRF_ALLOWED_DOMAINS = frozenset(
    {
        "graph.microsoft.com",
        "graph.microsoft.us",
        "login.microsoftonline.com",
        "login.microsoftonline.us",
        "api.botframework.com",
        "api.botframework.us",
        "smba.trafficmanager.net",
        "sharepoint.com",
        "sharepoint-df.com",
        "onedrive.com",
        "office.com",
        "office.net",
    }
)

_SSRF_DOMAIN_PARENT_PATTERN = re.compile(
    r"^(.+\.)?(" + "|".join(map(re.escape, _SSRF_ALLOWED_DOMAINS)) + r")$",
    re.IGNORECASE,
)

_SSRF_PRIVATE_RANGES = [
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("169.254.0.0/16"),
    ipaddress.ip_network("0.0.0.0/8"),
    ipaddress.ip_network("fc00::/7"),
    ipaddress.ip_network("::1/128"),
]


def _is_ssrf_safe_url(url: str) -> tuple[bool, str]:
    """三层 SSRF 防护：协议 → 域名白名单 → DNS 私有地址检测。

    Returns:
        (is_safe, reason) — safe 为 True 表示可安全请求。
    """
    parsed = urlparse(url)

    if parsed.scheme != "https":
        return False, f"Protocol must be HTTPS, got '{parsed.scheme}'"

    hostname = parsed.hostname or ""
    if not hostname:
        return False, "No hostname in URL"

    if not _SSRF_DOMAIN_PARENT_PATTERN.match(hostname):
        return False, f"Domain '{hostname}' not in SSRF allowlist"

    try:
        addr = ipaddress.ip_address(hostname)
        for net in _SSRF_PRIVATE_RANGES:
            if addr in net:
                return False, f"IP '{hostname}' is in private range {net}"
    except ValueError:
        pass

    return True, ""


def _validate_upload_url(url: str) -> bool:
    safe, reason = _is_ssrf_safe_url(url)
    if not safe:
        logger.warning(f"MSTeams SSRF blocked: {reason} (url={url[:120]})")
    return safe


def _validate_graph_url(url: str) -> bool:
    parsed = urlparse(url)
    hostname = (parsed.hostname or "").lower()
    if hostname in ("graph.microsoft.com", "graph.microsoft.us"):
        return True
    logger.warning(f"MSTeams SSRF blocked: Graph URL host={hostname}")
    return False


def needs_file_consent(file_size: int, chat_type: str = "direct") -> bool:
    if chat_type == "direct":
        return file_size > FILE_CONSENT_MAX_SIZE
    return file_size > FILE_CONSENT_MAX_SIZE * 5


def choose_upload_strategy(
    file_size: int,
    chat_type: str,
    sharepoint_site_id: str = "",
) -> str:
    if chat_type == "direct" and file_size <= FILE_CONSENT_MAX_SIZE:
        return "direct"
    if sharepoint_site_id and chat_type in ("group", "channel"):
        return "sharepoint"
    if chat_type in ("group",) and file_size > FILE_CONSENT_MAX_SIZE:
        return "onedrive"
    return "direct"


def build_file_consent_card(
    filename: str,
    file_size: int,
    description: str = "",
    accept_context: dict[str, Any] | None = None,
    decline_context: dict[str, Any] | None = None,
) -> dict[str, Any]:
    size_mb = file_size / (1024 * 1024)
    body_text = f"发送文件 **{filename}** ({size_mb:.1f} MB)"
    if description:
        body_text += f"\n\n{description}"

    return {
        "type": "AdaptiveCard",
        "version": "1.5",
        "body": [
            {
                "type": "TextBlock",
                "size": "Large",
                "weight": "Bolder",
                "text": "文件发送确认",
            },
            {
                "type": "TextBlock",
                "text": body_text,
                "wrap": True,
            },
        ],
        "actions": [
            {
                "type": "Action.Submit",
                "title": "接受",
                "style": "positive",
                "data": {"action": "file_consent_accept", **(accept_context or {})},
            },
            {
                "type": "Action.Submit",
                "title": "拒绝",
                "style": "destructive",
                "data": {"action": "file_consent_decline", **(decline_context or {})},
            },
        ],
    }


async def upload_to_sharepoint(
    token: str,
    file_data: bytes,
    filename: str,
    site_id: str = "",
    drive_id: str = "",
    folder_path: str = "ForcePilot Uploads",
) -> dict[str, Any]:
    import aiohttp

    graph_url = "https://graph.microsoft.com/v1.0"
    if not _validate_graph_url(graph_url):
        return {"error": "SSRF blocked: invalid Graph URL"}

    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }

    if site_id and drive_id:
        upload_path = f"/sites/{site_id}/drives/{drive_id}/root:/{folder_path}/{filename}:/createUploadSession"
    else:
        upload_path = f"/me/drive/root:/{folder_path}/{filename}:/createUploadSession"

    upload_session_url = f"{graph_url}{upload_path}"

    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                upload_session_url,
                headers=headers,
                json={"item": {"@microsoft.graph.conflictBehavior": "rename"}},
            ) as resp:
                if resp.status not in (200, 201):
                    body = await resp.text()
                    logger.warning(f"SharePoint upload session failed: HTTP {resp.status} - {body[:200]}")
                    return {"error": f"HTTP {resp.status}", "detail": body[:500]}
                upload_session = await resp.json()

            upload_url = upload_session.get("uploadUrl", "")
            if not upload_url:
                return {"error": "No upload URL returned"}

            if not _is_ssrf_safe_url(upload_url)[0]:
                return {"error": "SSRF blocked: upload URL not allowed"}

            total_size = len(file_data)
            chunk_size = 4 * 1024 * 1024
            offset = 0

            while offset < total_size:
                end = min(offset + chunk_size, total_size)
                chunk = file_data[offset:end]
                range_header = f"bytes {offset}-{end - 1}/{total_size}"

                async with session.put(
                    upload_url,
                    data=chunk,
                    headers={"Content-Range": range_header, "Content-Length": str(len(chunk))},
                ) as resp:
                    if resp.status in (200, 201):
                        result = await resp.json()
                        web_url = result.get("webUrl", "")
                        return {"webUrl": web_url, "id": result.get("id", ""), "name": result.get("name", filename)}
                    if resp.status == 202:
                        offset = end
                        continue
                    body = await resp.text()
                    logger.warning(f"SharePoint chunk upload failed: HTTP {resp.status}")
                    return {"error": f"HTTP {resp.status}", "detail": body[:500]}

            return {"error": "Upload incomplete"}

    except Exception as e:
        logger.error(f"SharePoint upload error: {e}")
        return {"error": str(e)}


def create_shareable_link(token: str, item_id: str, link_type: str = "view") -> dict[str, Any]:
    return {
        "type": "share",
        "item_id": item_id,
        "link_type": link_type,
        "create_link_endpoint": f"https://graph.microsoft.com/v1.0/me/drive/items/{item_id}/createLink",
    }


def build_teams_file_info_card(
    filename: str,
    file_size: int,
    content_url: str = "",
    site_url: str = "",
) -> dict[str, Any]:
    size_str = f"{file_size / (1024 * 1024):.1f} MB" if file_size >= 1024 * 1024 else f"{file_size / 1024:.1f} KB"

    return {
        "contentType": "application/vnd.microsoft.teams.file.info",
        "content": {
            "name": filename,
            "size": file_size,
            "siteUrl": site_url,
        },
        "metadata": {
            "filename": filename,
            "size_str": size_str,
            "uploaded_at": "",
        },
    }


_PENDING_UPLOADS_DIRNAME = "msteams-pending-uploads"


class PendingUploadStore:
    """跨进程 FileConsentCard 上传状态持久化。

    用于 Gateway 和 CLI 不同进程间的 FileConsentCard 回调共享。
    """

    def __init__(self, storage_dir: str | None = None):
        self._storage_dir = Path(storage_dir or str(Path.home() / ".yuxi" / _PENDING_UPLOADS_DIRNAME))
        self._storage_dir.mkdir(parents=True, exist_ok=True)

    def _file_path(self, upload_id: str) -> Path:
        return self._storage_dir / f"{upload_id}.json"

    def save_pending(
        self,
        upload_id: str,
        filename: str,
        file_data_b64: str,
        conversation_id: str,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        import json

        entry = {
            "upload_id": upload_id,
            "filename": filename,
            "file_data_b64": file_data_b64,
            "conversation_id": conversation_id,
            "created_at": time.time(),
            "metadata": metadata or {},
        }
        try:
            self._file_path(upload_id).write_text(
                json.dumps(entry, ensure_ascii=False),
                encoding="utf-8",
            )
        except OSError as e:
            logger.error(f"MSTeams pending upload: failed to save {upload_id}: {e}")

    def load_pending(self, upload_id: str) -> dict[str, Any] | None:
        import json

        fp = self._file_path(upload_id)
        if not fp.exists():
            return None
        try:
            return json.loads(fp.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"MSTeams pending upload: failed to load {upload_id}: {e}")
            return None

    def remove_pending(self, upload_id: str) -> None:
        try:
            fp = self._file_path(upload_id)
            if fp.exists():
                fp.unlink()
        except OSError as e:
            logger.warning(f"MSTeams pending upload: failed to remove {upload_id}: {e}")

    def list_pending_ids(self) -> list[str]:
        try:
            files = list(self._storage_dir.glob("*.json"))
        except OSError:
            return []
        return [f.stem for f in files]

    def cleanup_expired(self, max_age_seconds: int = 3600) -> int:
        now = time.time()
        cleaned = 0
        for fp in list(self._storage_dir.glob("*.json")):
            try:
                age = now - fp.stat().st_mtime
                if age > max_age_seconds:
                    fp.unlink()
                    cleaned += 1
            except OSError:
                pass
        if cleaned:
            logger.info(f"MSTeams pending upload: cleaned {cleaned} expired entries")
        return cleaned