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}")