ForcePilot/backend/package/yuxi/channel/extensions/msteams/files.py

from __future__ import annotations

import ipaddress
import logging
import re
from enum import StrEnum

logger = logging.getLogger(__name__)

SSRF_BLOCKED_DOMAINS: tuple[str, ...] = (
    "169.254.",
    "metadata.google.internal",
    "metadata.tencentyun.com",
    "100.64.",
    "0.0.0.0",
)
SSRF_BLOCKED_NETS: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = (
    ipaddress.IPv4Network("127.0.0.0/8"),
    ipaddress.IPv4Network("10.0.0.0/8"),
    ipaddress.IPv4Network("172.16.0.0/12"),
    ipaddress.IPv4Network("192.168.0.0/16"),
    ipaddress.IPv4Network("169.254.0.0/16"),
    ipaddress.IPv4Network("100.64.0.0/10"),
    ipaddress.IPv4Network("0.0.0.0/8"),
    ipaddress.IPv6Network("::1/128"),
    ipaddress.IPv6Network("fe80::/10"),
)
MAX_DOWNLOAD_SIZE = 100 * 1024 * 1024


class UploadRoute(StrEnum):
    ATTACHMENT_URL = "attachment_url"
    CONNECTOR = "connector"
    ONEDRIVE = "onedrive"
    SHAREPOINT = "sharepoint"
    M365_NATIVE = "m365_native"


def is_url_microsoft_safe(url: str) -> bool:
    try:
        host = _extract_host(url)
        if not host:
            return False

        host_lower = host.lower()

        microsoft_domains = (
            "graph.microsoft.com",
            "graph.microsoft.us",
            "api.teams.microsoft.com",
            "login.microsoftonline.com",
            "login.microsoftonline.us",
            "smba.trafficmanager.net",
            "sharepoint.com",
            "sharepoint.cn",
            "onedrive.com",
            "office.com",
            "office365.com",
            "microsoft.com",
            "microsoftonline.com",
            "azure.com",
            "azure.cn",
        )
        if host_lower == "localhost":
            return False
        if host_lower in ("127.0.0.1", "::1", "0.0.0.0"):
            return False

        for domain in microsoft_domains:
            if host_lower == domain or host_lower.endswith(f".{domain}"):
                return True

        return False
    except Exception:
        return False


def _validate_url_safety(url: str) -> bool:
    try:
        host = _extract_host(url)
        if not host:
            return False

        host_lower = host.lower()

        if host_lower == "localhost":
            logger.warning("SSRF blocked: localhost")
            return False
        if host_lower in ("127.0.0.1", "::1", "0.0.0.0"):
            logger.warning("SSRF blocked: loopback %s", host_lower)
            return False

        try:
            ip = ipaddress.ip_address(host)
            for net in SSRF_BLOCKED_NETS:
                if ip in net:
                    logger.warning("SSRF blocked: private IP %s", host)
                    return False
            return True
        except ValueError:
            pass

        for domain in SSRF_BLOCKED_DOMAINS:
            if domain in host_lower:
                logger.warning("SSRF blocked: %s contains %s", host_lower, domain)
                return False

        try:
            ip = ipaddress.ip_address(host)
            for net in SSRF_BLOCKED_NETS:
                if ip in net:
                    logger.warning("SSRF blocked: resolved IP %s is private", ip)
                    return False
        except ValueError:
            pass

        return True
    except Exception:
        return False


def resolve_upload_route(
    file_size: int,
    *,
    url: str = "",
    is_team_channel: bool = False,
) -> UploadRoute:
    if url and is_url_microsoft_safe(url):
        return UploadRoute.M365_NATIVE

    if file_size < 4 * 1024 * 1024:
        return UploadRoute.ATTACHMENT_URL

    if file_size < 50 * 1024 * 1024:
        return UploadRoute.CONNECTOR

    if is_team_channel:
        return UploadRoute.SHAREPOINT

    return UploadRoute.ONEDRIVE


def _extract_host(url: str) -> str:
    url = url.strip()
    if "://" not in url:
        url = f"https://{url}"
    idx = url.index("://") + 3
    host = url[idx:]
    host = re.split(r"[/?#:@]", host)[0]
    return host.strip()


async def download_public_file(url: str) -> bytes | None:
    if not _validate_url_safety(url):
        return None

    try:
        import httpx

        async with httpx.AsyncClient(timeout=60.0, follow_redirects=True) as client:
            resp = await client.get(url)
            resp.raise_for_status()

            content_length = int(resp.headers.get("content-length", "0"))
            if content_length > MAX_DOWNLOAD_SIZE:
                logger.warning("File too large: %d bytes from %s", content_length, url)
                return None

            content = await resp.aread()
            if len(content) > MAX_DOWNLOAD_SIZE:
                logger.warning("File too large: %d bytes from %s", len(content), url)
                return None

            return content
    except Exception as e:
        logger.warning("Failed to download file from %s: %s", url, e)
        return None