231 lines
7.4 KiB
Python
231 lines
7.4 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import re
|
|
import time
|
|
from typing import Any, TYPE_CHECKING
|
|
|
|
import httpx
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from .client import UrbitClient
|
|
|
|
|
|
_MEMEX_UPLOAD_TIMEOUT = 120.0
|
|
_S3_UPLOAD_TIMEOUT = 300.0
|
|
_UNSAFE_FILENAME_RE = re.compile(r"[/\\:*?\"<>|]")
|
|
|
|
|
|
def sanitize_filename(filename: str) -> str:
|
|
name = os.path.basename(filename) or "file"
|
|
name = _UNSAFE_FILENAME_RE.sub("_", name)
|
|
name = name.strip(". ")
|
|
if not name:
|
|
name = f"file_{hashlib.md5(filename.encode()).hexdigest()[:8]}"
|
|
if len(name) > 255:
|
|
base, ext = os.path.splitext(name)
|
|
name = base[: 255 - len(ext)] + ext
|
|
return name
|
|
|
|
|
|
async def get_storage_configuration(client: UrbitClient) -> dict[str, Any] | None:
|
|
from .scry import scry_get
|
|
|
|
result = await scry_get(client, "storage", "/configuration.json")
|
|
if isinstance(result, dict):
|
|
return result
|
|
return None
|
|
|
|
|
|
async def get_storage_credentials(client: UrbitClient) -> dict[str, Any] | None:
|
|
from .scry import scry_get
|
|
|
|
result = await scry_get(client, "storage", "/credentials.json")
|
|
if isinstance(result, dict):
|
|
return result
|
|
return None
|
|
|
|
|
|
async def upload_file(
|
|
client: UrbitClient,
|
|
file_data: bytes,
|
|
filename: str,
|
|
content_type: str = "application/octet-stream",
|
|
) -> str | None:
|
|
safe_filename = sanitize_filename(filename)
|
|
if safe_filename != filename:
|
|
logger.debug(f"[Urbit] Filename sanitized: {filename!r} → {safe_filename!r}")
|
|
|
|
storage_config = await get_storage_configuration(client)
|
|
if not storage_config:
|
|
logger.warning("[Urbit] No storage configuration found, upload skipped")
|
|
return None
|
|
|
|
backend = storage_config.get("backend", "")
|
|
|
|
if backend == "memex" or not backend:
|
|
return await _upload_via_memex(client, file_data, safe_filename, content_type)
|
|
elif backend == "s3":
|
|
return await _upload_via_s3(client, file_data, safe_filename, content_type, storage_config)
|
|
else:
|
|
logger.warning(f"[Urbit] Unknown storage backend: {backend}")
|
|
return None
|
|
|
|
|
|
async def _upload_via_memex(
|
|
client: UrbitClient,
|
|
file_data: bytes,
|
|
filename: str,
|
|
content_type: str,
|
|
) -> str | None:
|
|
try:
|
|
upload_url = f"{client.ship_url}/~/memex/upload"
|
|
files = {"file": (filename, file_data, content_type)}
|
|
r = await client.http.post(
|
|
upload_url,
|
|
files=files,
|
|
timeout=httpx.Timeout(_MEMEX_UPLOAD_TIMEOUT),
|
|
)
|
|
r.raise_for_status()
|
|
result = r.json()
|
|
public_url = result.get("url", "") or result.get("publicUrl", "")
|
|
if public_url:
|
|
logger.info(f"[Urbit] Memex upload success: {filename} → {public_url}")
|
|
return public_url
|
|
logger.warning("[Urbit] Memex upload response missing URL")
|
|
return None
|
|
except Exception as e:
|
|
logger.warning(f"[Urbit] Memex upload failed for {filename}: {e}")
|
|
return None
|
|
|
|
|
|
async def _upload_via_s3(
|
|
client: UrbitClient,
|
|
file_data: bytes,
|
|
filename: str,
|
|
content_type: str,
|
|
storage_config: dict[str, Any],
|
|
) -> str | None:
|
|
credentials = await get_storage_credentials(client)
|
|
if not credentials:
|
|
logger.warning("[Urbit] No S3 credentials found")
|
|
return None
|
|
|
|
endpoint = credentials.get("endpoint", "")
|
|
access_key = credentials.get("accessKeyId", "")
|
|
secret_key = credentials.get("secretAccessKey", "")
|
|
bucket = storage_config.get("bucket", "moltbot-uploads")
|
|
|
|
if not all([endpoint, access_key, secret_key]):
|
|
logger.warning("[Urbit] Incomplete S3 credentials")
|
|
return None
|
|
|
|
file_hash = hashlib.sha256(file_data).hexdigest()[:16]
|
|
ext = os.path.splitext(filename)[1] or ".bin"
|
|
key = f"uploads/{int(time.time())}-{file_hash}{ext}"
|
|
|
|
try:
|
|
import hmac
|
|
from datetime import datetime, UTC
|
|
|
|
date_short = datetime.now(UTC).strftime("%Y%m%d")
|
|
date_long = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
|
region = credentials.get("region", "us-east-1")
|
|
service = "s3"
|
|
|
|
credential = f"{access_key}/{date_short}/{region}/{service}/aws4_request"
|
|
|
|
canonical_request = _build_canonical_request("PUT", key, content_type, file_data, bucket)
|
|
string_to_sign = _build_string_to_sign(date_long, date_short, region, service, canonical_request)
|
|
signing_key = _build_signing_key(secret_key, date_short, region, service)
|
|
signature = hmac.new(signing_key, string_to_sign.encode(), "sha256").hexdigest()
|
|
|
|
authorization = (
|
|
f"AWS4-HMAC-SHA256 "
|
|
f"Credential={credential}, "
|
|
f"SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date, "
|
|
f"Signature={signature}"
|
|
)
|
|
|
|
headers = {
|
|
"Content-Type": content_type,
|
|
"x-amz-content-sha256": hashlib.sha256(file_data).hexdigest(),
|
|
"x-amz-date": date_long,
|
|
"Authorization": authorization,
|
|
}
|
|
|
|
upload_url = f"{endpoint}/{bucket}/{key}"
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(_S3_UPLOAD_TIMEOUT)) as s3_client:
|
|
r = await s3_client.put(upload_url, content=file_data, headers=headers)
|
|
r.raise_for_status()
|
|
|
|
public_url = f"{endpoint}/{bucket}/{key}"
|
|
logger.info(f"[Urbit] S3 upload success: {filename} → {public_url}")
|
|
return public_url
|
|
|
|
except ImportError:
|
|
logger.warning("[Urbit] S3 upload skipped: hmac module unavailable")
|
|
return None
|
|
except Exception as e:
|
|
logger.warning(f"[Urbit] S3 upload failed for {filename}: {e}")
|
|
return None
|
|
|
|
|
|
def _build_canonical_request(
|
|
method: str,
|
|
key: str,
|
|
content_type: str,
|
|
data: bytes,
|
|
bucket: str,
|
|
) -> str:
|
|
payload_hash = hashlib.sha256(data).hexdigest()
|
|
canonical_uri = f"/{key}"
|
|
canonical_querystring = ""
|
|
canonical_headers = (
|
|
f"content-type:{content_type}\nhost:{bucket}\nx-amz-content-sha256:{payload_hash}\nx-amz-date:{'TIMESTAMP'}\n"
|
|
)
|
|
signed_headers = "content-type;host;x-amz-content-sha256;x-amz-date"
|
|
return f"{method}\n{canonical_uri}\n{canonical_querystring}\n{canonical_headers}\n{signed_headers}\n{payload_hash}"
|
|
|
|
|
|
def _build_string_to_sign(
|
|
date_long: str,
|
|
date_short: str,
|
|
region: str,
|
|
service: str,
|
|
canonical_request: str,
|
|
) -> str:
|
|
import hashlib
|
|
|
|
scope = f"{date_short}/{region}/{service}/aws4_request"
|
|
return f"AWS4-HMAC-SHA256\n{date_long}\n{scope}\n{hashlib.sha256(canonical_request.encode()).hexdigest()}"
|
|
|
|
|
|
def _build_signing_key(secret_key: str, date_short: str, region: str, service: str) -> bytes:
|
|
import hmac
|
|
|
|
k_date = hmac.new(f"AWS4{secret_key}".encode(), date_short.encode(), "sha256").digest()
|
|
k_region = hmac.new(k_date, region.encode(), "sha256").digest()
|
|
k_service = hmac.new(k_region, service.encode(), "sha256").digest()
|
|
return hmac.new(k_service, b"aws4_request", "sha256").digest()
|
|
|
|
|
|
def assert_trusted_upload_url(url: str) -> bool:
|
|
from .probe import validate_urbit_url
|
|
|
|
is_valid, warnings = validate_urbit_url(url)
|
|
if not is_valid:
|
|
logger.warning(f"[Urbit] Untrusted upload URL: {url} - {warnings}")
|
|
return False
|
|
return True
|
|
|
|
|
|
def assert_safe_upload_result_url(url: str) -> bool:
|
|
if not url.startswith(("http://", "https://")):
|
|
logger.warning(f"[Urbit] Unsafe upload result URL: {url}")
|
|
return False
|
|
return True
|