ForcePilot/backend/package/yuxi/channel/extensions/tlon/tlon_api.py

import logging
from urllib.parse import urlparse

import httpx

from yuxi.channel.extensions.tlon.auth import authenticate

logger = logging.getLogger(__name__)

MEMEX_TRUSTED_HOSTS = [
    "tlon.network",
    "test.tlon.systems",
]


async def upload_file(api, blob: bytes, file_name: str,
                      content_type: str) -> str | None:
    ship_url = api._config.url
    cookie = api._config.cookie

    try:
        storage_config = await api.scry("/storage/configuration.json")
    except Exception as e:
        logger.warning("[tlon] Failed to scry storage config: %s", e)
        return None

    service = storage_config.get("service", "")

    is_hosted_url = any(
        host in ship_url for host in MEMEX_TRUSTED_HOSTS
    )

    if is_hosted_url and service == "presigned-url":
        return await _upload_memex(ship_url, cookie, api._config.ship, blob, file_name, content_type)
    else:
        return await _upload_s3_presigned(api, storage_config, blob, file_name, content_type)


async def _upload_memex(ship_url: str, cookie: str, ship: str,
                        blob: bytes, file_name: str,
                        content_type: str) -> str | None:
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            token_response = await client.get(
                f"{ship_url}/~/scry/genuine/secret.json",
                headers={"Cookie": cookie},
            )
            token_response.raise_for_status()
            token_data = token_response.json()
            token = token_data.get("secret", "")

            upload_response = await client.put(
                f"https://memex.tlon.network/v1/{ship}/upload",
                json={"filename": file_name, "contentType": content_type},
                headers={
                    "Authorization": f"Bearer {token}",
                    "Cookie": cookie,
                },
            )
            upload_response.raise_for_status()
            upload_data = upload_response.json()
            upload_url = upload_data.get("url", "")

            _assert_trusted_memex_url(upload_url)

            await client.put(
                upload_url,
                content=blob,
                headers={"Content-Type": content_type},
            )

            return upload_data.get("publicUrl", "")
    except Exception as e:
        logger.warning("[tlon] Memex upload failed: %s", e)
        return None


async def _upload_s3_presigned(api, storage_config: dict,
                               blob: bytes, file_name: str,
                               content_type: str) -> str | None:
    try:
        creds = await api.scry("/storage/credentials.json")
    except Exception as e:
        logger.warning("[tlon] Failed to scry storage credentials: %s", e)
        return None

    try:
        import boto3
        from botocore.config import Config as BotoConfig

        endpoint = creds.get("endpoint", "")
        region = storage_config.get("region", "us-east-1")
        access_key = creds.get("accessKeyId", "")
        secret_key = creds.get("secretAccessKey", "")
        bucket = storage_config.get("buckets", [""])[0] if storage_config.get("buckets") else ""

        s3 = boto3.client(
            "s3",
            endpoint_url=endpoint,
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            region_name=region,
            config=BotoConfig(signature_version="s3v4"),
        )

        presigned = s3.generate_presigned_url(
            "put_object",
            Params={
                "Bucket": bucket,
                "Key": file_name,
                "ContentType": content_type,
            },
            ExpiresIn=3600,
        )

        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.put(
                presigned,
                content=blob,
                headers={"Content-Type": content_type},
            )
            response.raise_for_status()

        public_url_base = storage_config.get("publicUrlBase", endpoint)
        return f"{public_url_base.rstrip('/')}/{bucket}/{file_name}"
    except ImportError:
        logger.warning("[tlon] boto3 not available for S3 upload")
        return None
    except Exception as e:
        logger.warning("[tlon] S3 upload failed: %s", e)
        return None


def _assert_trusted_memex_url(url: str) -> None:
    parsed = urlparse(url)
    hostname = parsed.hostname or ""

    is_trusted = any(hostname.endswith(host) or hostname == host
                     for host in MEMEX_TRUSTED_HOSTS)

    if not is_trusted:
        raise ValueError(f"Untrusted Memex URL: {url}")

    if parsed.port and parsed.port not in (80, 443):
        raise ValueError(f"Untrusted Memex port: {parsed.port}")