#!/usr/bin/env python3
"""Download a backup set from Garage S3 to a local archive directory."""

import argparse
import datetime
import hashlib
import hmac
import os
import sys
import urllib.request
import xml.etree.ElementTree as ET
from urllib.parse import urlparse, quote


def _aws_quote(value: str) -> str:
    return quote(value, safe="-_.~")


def _canonical_query(params: dict) -> str:
    if not params:
        return ""
    items = []
    for key in sorted(params.keys()):
        val = params[key]
        items.append(f"{_aws_quote(str(key))}={_aws_quote(str(val))}")
    return "&".join(items)


def _sign(key: bytes, msg: str) -> bytes:
    return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()


def _signature_key(
    secret_key: str, date_stamp: str, region: str, service: str
) -> bytes:
    k_date = _sign(("AWS4" + secret_key).encode("utf-8"), date_stamp)
    k_region = _sign(k_date, region)
    k_service = _sign(k_region, service)
    k_signing = _sign(k_service, "aws4_request")
    return k_signing


def _canonical_uri(path: str) -> str:
    if not path.startswith("/"):
        path = "/" + path
    parts = [_aws_quote(p) for p in path.split("/")]
    return "/".join(parts)


def _sign_request(
    method: str,
    host: str,
    path: str,
    params: dict,
    access_key: str,
    secret_key: str,
    region: str,
    service: str,
) -> tuple[dict[str, str], str, str]:
    """Sign an AWS request and return headers, canonical URI, and canonical query string."""
    now = datetime.datetime.now(datetime.timezone.utc)
    amz_date = now.strftime("%Y%m%dT%H%M%SZ")
    date_stamp = now.strftime("%Y%m%d")
    payload_hash = hashlib.sha256(b"").hexdigest()

    canonical_uri = _canonical_uri(path)
    canonical_querystring = _canonical_query(params)

    canonical_headers = (
        f"host:{host}\n"
        f"x-amz-content-sha256:{payload_hash}\n"
        f"x-amz-date:{amz_date}\n"
    )
    signed_headers = "host;x-amz-content-sha256;x-amz-date"

    canonical_request = "\n".join(
        [
            method,
            canonical_uri,
            canonical_querystring,
            canonical_headers,
            signed_headers,
            payload_hash,
        ]
    )

    algorithm = "AWS4-HMAC-SHA256"
    credential_scope = f"{date_stamp}/{region}/{service}/aws4_request"
    hashed_request = hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()
    string_to_sign = "\n".join([algorithm, amz_date, credential_scope, hashed_request])

    signing_key = _signature_key(secret_key, date_stamp, region, service)
    signature = hmac.new(
        signing_key, string_to_sign.encode("utf-8"), hashlib.sha256
    ).hexdigest()

    authorization_header = (
        f"{algorithm} Credential={access_key}/{credential_scope}, "
        f"SignedHeaders={signed_headers}, Signature={signature}"
    )

    headers = {
        "x-amz-date": amz_date,
        "x-amz-content-sha256": payload_hash,
        "Authorization": authorization_header,
    }

    return headers, canonical_uri, canonical_querystring


def _aws_request(
    method: str,
    endpoint: str,
    path: str,
    params: dict,
    access_key: str,
    secret_key: str,
    region: str,
):
    parsed = urlparse(endpoint)
    host = parsed.netloc

    headers, canonical_uri, canonical_querystring = _sign_request(
        method, host, path, params, access_key, secret_key, region, "s3"
    )

    url = f"{parsed.scheme}://{host}{canonical_uri}"
    if canonical_querystring:
        url = f"{url}?{canonical_querystring}"

    req = urllib.request.Request(url, method=method, headers=headers)
    return urllib.request.urlopen(req)


def _list_objects(
    endpoint: str,
    bucket: str,
    prefix: str,
    access_key: str,
    secret_key: str,
    region: str,
):
    token = ""
    while True:
        params = {"list-type": "2"}
        if prefix:
            params["prefix"] = prefix
        if token:
            params["continuation-token"] = token
        with _aws_request(
            "GET", endpoint, f"/{bucket}", params, access_key, secret_key, region
        ) as resp:
            data = resp.read()
        root = ET.fromstring(data)
        for entry in root.findall(".//{*}Contents"):
            key = entry.findtext("{*}Key")
            if key:
                yield key
        is_truncated = root.findtext(".//{*}IsTruncated")
        if str(is_truncated).lower() != "true":
            break
        token = root.findtext(".//{*}NextContinuationToken") or ""
        if not token:
            break


def _safe_join(base: str, rel: str) -> str:
    rel = rel.lstrip("/")
    path = os.path.normpath(os.path.join(base, rel))
    base_norm = os.path.normpath(base)
    if not path.startswith(base_norm):
        raise ValueError("Refusing to write outside archive directory")
    return path


def main() -> int:
    parser = argparse.ArgumentParser(description="Archive Garage S3 backups locally")
    parser.add_argument("--endpoint", required=True)
    parser.add_argument("--region", required=True)
    parser.add_argument("--access-key", required=True)
    parser.add_argument("--secret-key", required=True)
    parser.add_argument("--bucket", required=True)
    parser.add_argument("--prefix", default="")
    parser.add_argument("--contains", default="")
    parser.add_argument("--dest", required=True)
    args = parser.parse_args()

    os.makedirs(args.dest, exist_ok=True)
    prefix = args.prefix
    if prefix and not prefix.endswith("/"):
        prefix = prefix + "/"

    keys = []
    for key in _list_objects(
        args.endpoint,
        args.bucket,
        prefix,
        args.access_key,
        args.secret_key,
        args.region,
    ):
        if args.contains and args.contains not in key:
            continue
        keys.append(key)

    if not keys:
        print("No matching objects found in bucket.", file=sys.stderr)
        return 2

    for key in keys:
        rel = key
        if prefix and key.startswith(prefix):
            rel = key[len(prefix) :]
        if not rel:
            rel = os.path.basename(key)
        dest_path = _safe_join(args.dest, rel)
        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
        with _aws_request(
            "GET",
            args.endpoint,
            f"/{args.bucket}/{key}",
            {},
            args.access_key,
            args.secret_key,
            args.region,
        ) as resp:
            with open(dest_path, "wb") as fh:
                while True:
                    chunk = resp.read(1024 * 1024)
                    if not chunk:
                        break
                    fh.write(chunk)
        print(f"Archived {key} -> {dest_path}")

    return 0


if __name__ == "__main__":
    sys.exit(main())