src/r2_artifacts.py

import hashlib
import hmac
import os
from datetime import datetime, timezone
from urllib.parse import quote

import requests
from fastapi import HTTPException


R2_REGION = "auto"
R2_SERVICE = "s3"
R2_SIGNED_URL_EXPIRES_SECONDS = int(os.getenv("R2_SIGNED_URL_EXPIRES_SECONDS", "600"))


def _env(name: str) -> str:
    return os.getenv(name, "").strip()


def r2_configured() -> bool:
    return all(
        _env(name)
        for name in (
            "R2_ACCOUNT_ID",
            "R2_BUCKET_NAME",
            "R2_ACCESS_KEY_ID",
            "R2_SECRET_ACCESS_KEY",
        )
    )


def require_r2_configured():
    if not r2_configured():
        raise HTTPException(status_code=503, detail="R2 artifact storage is not configured.")


def get_score_storage_keys(user_id: str, score_id: str) -> dict:
    base = f"users/{user_id}/scores/{score_id}"
    return {
        "original_musicxml": f"{base}/original.musicxml",
        "full_html": f"{base}/renders/full.html",
        "fingered_html": f"{base}/renders/fingered.html",
        "full_svg": f"{base}/renders/full.svg",
        "full_pdf": f"{base}/exports/full.pdf",
        "cache_prefix": f"{base}/cache",
    }


def _endpoint_host() -> str:
    return f"{_env('R2_ACCOUNT_ID')}.r2.cloudflarestorage.com"


def _canonical_uri(bucket: str, key: str) -> str:
    parts = [quote(part, safe="-_.~") for part in key.split("/")]
    return f"/{quote(bucket, safe='-_.~')}/{'/'.join(parts)}"


def _object_url(key: str) -> str:
    return f"https://{_endpoint_host()}{_canonical_uri(_env('R2_BUCKET_NAME'), key)}"


def _signing_key(date_stamp: str) -> bytes:
    secret = _env("R2_SECRET_ACCESS_KEY").encode("utf-8")
    k_date = hmac.new(b"AWS4" + secret, date_stamp.encode("utf-8"), hashlib.sha256).digest()
    k_region = hmac.new(k_date, R2_REGION.encode("utf-8"), hashlib.sha256).digest()
    k_service = hmac.new(k_region, R2_SERVICE.encode("utf-8"), hashlib.sha256).digest()
    return hmac.new(k_service, b"aws4_request", hashlib.sha256).digest()


def _signature(string_to_sign: str, date_stamp: str) -> str:
    return hmac.new(_signing_key(date_stamp), string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()


def _credential_scope(date_stamp: str) -> str:
    return f"{date_stamp}/{R2_REGION}/{R2_SERVICE}/aws4_request"


def _amz_dates() -> tuple[str, str]:
    now = datetime.now(timezone.utc)
    return now.strftime("%Y%m%dT%H%M%SZ"), now.strftime("%Y%m%d")


def upload_r2_artifact(*, key: str, body: str | bytes, content_type: str) -> str:
    require_r2_configured()
    body_bytes = body.encode("utf-8") if isinstance(body, str) else body
    payload_hash = hashlib.sha256(body_bytes).hexdigest()
    amz_date, date_stamp = _amz_dates()
    host = _endpoint_host()
    canonical_uri = _canonical_uri(_env("R2_BUCKET_NAME"), key)
    signed_headers = "content-type;host;x-amz-content-sha256;x-amz-date"
    canonical_headers = (
        f"content-type:{content_type}\n"
        f"host:{host}\n"
        f"x-amz-content-sha256:{payload_hash}\n"
        f"x-amz-date:{amz_date}\n"
    )
    canonical_request = "\n".join(
        [
            "PUT",
            canonical_uri,
            "",
            canonical_headers,
            signed_headers,
            payload_hash,
        ]
    )
    credential_scope = _credential_scope(date_stamp)
    string_to_sign = "\n".join(
        [
            "AWS4-HMAC-SHA256",
            amz_date,
            credential_scope,
            hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(),
        ]
    )
    authorization = (
        "AWS4-HMAC-SHA256 "
        f"Credential={_env('R2_ACCESS_KEY_ID')}/{credential_scope}, "
        f"SignedHeaders={signed_headers}, "
        f"Signature={_signature(string_to_sign, date_stamp)}"
    )
    try:
        response = requests.put(
            _object_url(key),
            data=body_bytes,
            headers={
                "Authorization": authorization,
                "Content-Type": content_type,
                "Host": host,
                "x-amz-content-sha256": payload_hash,
                "x-amz-date": amz_date,
            },
            timeout=30,
        )
    except requests.exceptions.RequestException as exc:
        raise HTTPException(status_code=503, detail="R2 upload failed. Try again.") from exc
    if response.status_code >= 400:
        raise HTTPException(status_code=503, detail=f"R2 upload failed: {response.text[:200]}")
    return key


def get_r2_artifact_download_url(key: str, *, expires_in: int | None = None) -> str:
    require_r2_configured()
    expires = int(expires_in or R2_SIGNED_URL_EXPIRES_SECONDS)
    amz_date, date_stamp = _amz_dates()
    credential_scope = _credential_scope(date_stamp)
    credential = f"{_env('R2_ACCESS_KEY_ID')}/{credential_scope}"
    host = _endpoint_host()
    signed_headers = "host"
    query_params = {
        "X-Amz-Algorithm": "AWS4-HMAC-SHA256",
        "X-Amz-Credential": credential,
        "X-Amz-Date": amz_date,
        "X-Amz-Expires": str(expires),
        "X-Amz-SignedHeaders": signed_headers,
    }
    canonical_query = "&".join(
        f"{quote(k, safe='-_.~')}={quote(v, safe='-_.~')}"
        for k, v in sorted(query_params.items())
    )
    canonical_request = "\n".join(
        [
            "GET",
            _canonical_uri(_env("R2_BUCKET_NAME"), key),
            canonical_query,
            f"host:{host}\n",
            signed_headers,
            "UNSIGNED-PAYLOAD",
        ]
    )
    string_to_sign = "\n".join(
        [
            "AWS4-HMAC-SHA256",
            amz_date,
            credential_scope,
            hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(),
        ]
    )
    signature = _signature(string_to_sign, date_stamp)
    return f"{_object_url(key)}?{canonical_query}&X-Amz-Signature={signature}"