hy3d/hunyuan3dweb/cos_upload.py

#!/usr/bin/env python3
"""
Pure Python COS upload for Tencent Hunyuan 3D.
Uses the API to get temporary credentials, then directly uploads to COS.
"""

import hashlib
import hmac
import mimetypes
from pathlib import Path
from urllib.parse import quote

import requests

from . import Hunyuan3DAPIComplete
from .api import load_cookies_from_file


def _cam_safe_url_encode(s):
    return quote(str(s), safe="")


def _get_object_keys(obj):
    return sorted(obj.keys(), key=lambda x: x.lower())


def _obj2str(obj, lower_case_key=False):
    items = []
    for key in _get_object_keys(obj):
        val = obj[key]
        if val is None:
            val = ""
        else:
            val = str(val)
        encoded_key = _cam_safe_url_encode(key).lower() if lower_case_key else _cam_safe_url_encode(key)
        encoded_val = _cam_safe_url_encode(val) or ""
        items.append(f"{encoded_key}={encoded_val}")
    return "&".join(items)


def _cos_v1_auth(method, pathname, query_params, headers, secret_id, secret_key, key_time):
    """Generate COS V1 authorization header."""
    sign_key = hmac.new(secret_key.encode("utf-8"), key_time.encode("utf-8"), hashlib.sha1).hexdigest()

    query_str = _obj2str(query_params, True)
    headers_str = _obj2str(headers, True)
    format_string = f"{method}\n{pathname}\n{query_str}\n{headers_str}\n"

    format_string_sha1 = hashlib.sha1(format_string.encode("utf-8")).hexdigest()
    string_to_sign = f"sha1\n{key_time}\n{format_string_sha1}\n"

    signature = hmac.new(sign_key.encode("utf-8"), string_to_sign.encode("utf-8"), hashlib.sha1).hexdigest()

    q_header_list = ";".join(_get_object_keys(headers)).lower()
    q_url_param_list = ";".join(_get_object_keys(query_params)).lower()

    return "&".join([
        "q-sign-algorithm=sha1",
        f"q-ak={secret_id}",
        f"q-sign-time={key_time}",
        f"q-key-time={key_time}",
        f"q-header-list={q_header_list}",
        f"q-url-param-list={q_url_param_list}",
        f"q-signature={signature}",
    ])


def _guess_content_type(image_path):
    content_type, _ = mimetypes.guess_type(str(image_path))
    return content_type or "application/octet-stream"


def upload_file_to_cos(image_path, upload_info, use_accelerate=True):
    """
    Upload a local file to Tencent COS using temporary credentials from get_upload_info.

    Args:
        image_path: Local file path
        upload_info: Response from get_upload_info API
        use_accelerate: Use COS global accelerate endpoint

    Returns:
        requests.Response
    """
    bucket = upload_info["bucketName"]
    region = upload_info["region"]
    location = upload_info["location"]
    secret_id = upload_info["encryptTmpSecretId"]
    secret_key = upload_info["encryptTmpSecretKey"]
    token = upload_info["encryptToken"]
    start_time = upload_info["startTime"]
    expired_time = upload_info["expiredTime"]

    host = f"{bucket}.cos.accelerate.myqcloud.com" if use_accelerate else f"{bucket}.cos.{region}.myqcloud.com"
    key_time = f"{start_time};{expired_time}"
    file_size = Path(image_path).stat().st_size
    content_type = _guess_content_type(image_path)

    headers_for_sign = {
        "content-length": file_size,
        "content-type": content_type,
        "host": host,
    }

    auth = _cos_v1_auth(
        method="put",
        pathname=f"/{location}",
        query_params={},
        headers=headers_for_sign,
        secret_id=secret_id,
        secret_key=secret_key,
        key_time=key_time,
    )

    upload_url = f"https://{host}/{location}"

    with open(image_path, "rb") as f:
        resp = requests.put(
            upload_url,
            data=f,
            headers={
                "Authorization": auth,
                "x-cos-security-token": token,
                "Content-Type": content_type,
                "Content-Length": str(file_size),
                "Host": host,
            },
            timeout=120,
        )

    return resp


def upload_image(image_path, use_accelerate=True):
    """
    Full pipeline: get upload credentials from Hunyuan3D API -> upload to COS.

    Args:
        image_path: Local image file path
        use_accelerate: Use COS global accelerate endpoint

    Returns:
        str: resourceUrl for subsequent 3D generation
    """
    try:
        cookies = load_cookies_from_file()
    except FileNotFoundError as exc:
        raise RuntimeError("Cookies file not found. Please run 'hunyuan3dweb-login' first.") from exc
    api = Hunyuan3DAPIComplete(cookies=cookies)

    filename = Path(image_path).name
    upload_info = api.get_upload_info(filename)

    resp = upload_file_to_cos(image_path, upload_info, use_accelerate)
    resp.raise_for_status()

    return upload_info["resourceUrl"]