#!/usr/bin/env python3
"""Validate OSS object access with credentials from an ECS RAM role."""

import argparse
import base64
import hashlib
import hmac
import json
import os
import time
import urllib.parse
import urllib.request
from email.utils import formatdate


def load_credentials(role_name: str) -> dict[str, str]:
    role = urllib.parse.quote(role_name, safe="")
    url = f"http://100.100.100.200/latest/meta-data/ram/security-credentials/{role}"
    with urllib.request.urlopen(url, timeout=10) as response:
        credentials = json.load(response)
    if credentials.get("Code") != "Success":
        raise RuntimeError("ECS RAM-role metadata did not return Success")
    return credentials


def object_request(
    method: str,
    bucket: str,
    endpoint: str,
    object_key: str,
    credentials: dict[str, str],
    body: bytes | None = None,
) -> bytes:
    date = formatdate(usegmt=True)
    content_type = "application/octet-stream" if body is not None else ""
    token = credentials["SecurityToken"]
    canonical_headers = f"x-oss-security-token:{token}\n"
    canonical_resource = f"/{bucket}/{object_key}"
    string_to_sign = (
        f"{method}\n\n{content_type}\n{date}\n"
        f"{canonical_headers}{canonical_resource}"
    )
    digest = hmac.new(
        credentials["AccessKeySecret"].encode(),
        string_to_sign.encode(),
        hashlib.sha1,
    ).digest()
    signature = base64.b64encode(digest).decode()
    encoded_key = urllib.parse.quote(object_key, safe="/")
    request = urllib.request.Request(
        f"http://{bucket}.{endpoint}/{encoded_key}",
        data=body,
        method=method,
        headers={
            "Authorization": f"OSS {credentials['AccessKeyId']}:{signature}",
            "Date": date,
            "x-oss-security-token": token,
            **({"Content-Type": content_type} if content_type else {}),
        },
    )
    with urllib.request.urlopen(request, timeout=30) as response:
        return response.read()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--bucket", required=True)
    parser.add_argument("--endpoint", required=True)
    parser.add_argument("--role-name", required=True)
    args = parser.parse_args()

    credentials = load_credentials(args.role_name)
    object_key = f"cdc/smoke/ecs-ram-role-{int(time.time())}-{os.getpid()}"
    payload = b"date-integration OSS RAM-role smoke\n"
    try:
        object_request("PUT", args.bucket, args.endpoint, object_key, credentials, payload)
        downloaded = object_request(
            "GET", args.bucket, args.endpoint, object_key, credentials
        )
        if downloaded != payload:
            raise RuntimeError("OSS smoke object content did not round-trip")
    finally:
        object_request("DELETE", args.bucket, args.endpoint, object_key, credentials)

    print("Alibaba Cloud OSS ECS RAM-role write/read/delete smoke passed.")


if __name__ == "__main__":
    main()