#!/usr/bin/env python3
"""One S3 request, SigV4-signed with the key in the environment.

    S3_ENDPOINT=http://127.0.0.1:7390 S3_KEY=ak S3_SECRET=sk \
        s3req.py PUT /bucket/key [query] [--body FILE] [-H 'Name: value'] [-o FILE]

Prints the HTTP status on stdout, and for an error response the S3 error
code and message on stderr. With S3_KEY empty or unset the request is
anonymous. Python standard library only, so it runs on a bare EL9 node.
"""

import base64
import datetime
import hashlib
import hmac
import http.client
import os
import re
import sys
import urllib.parse

UNRESERVED = "-_.~"


def quote(s, safe=""):
    return urllib.parse.quote(s, safe=UNRESERVED + safe)


def canonical_query(query):
    pairs = urllib.parse.parse_qsl(query, keep_blank_values=True)
    return "&".join(f"{quote(k)}={quote(v)}" for k, v in sorted(pairs))


def sign(method, host, path, query, headers, body, key, secret, region):
    now = datetime.datetime.now(datetime.timezone.utc)
    amz_date = now.strftime("%Y%m%dT%H%M%SZ")
    day = now.strftime("%Y%m%d")
    headers["host"] = host
    headers["x-amz-date"] = amz_date
    headers["x-amz-content-sha256"] = hashlib.sha256(body).hexdigest()

    signed = sorted(headers)
    canonical = "\n".join([
        method,
        quote(path, safe="/"),
        canonical_query(query),
        "".join(f"{h}:{headers[h].strip()}\n" for h in signed),
        ";".join(signed),
        headers["x-amz-content-sha256"],
    ])
    scope = f"{day}/{region}/s3/aws4_request"
    to_sign = "\n".join(["AWS4-HMAC-SHA256", amz_date, scope,
                         hashlib.sha256(canonical.encode()).hexdigest()])

    k = ("AWS4" + secret).encode()
    for part in (day, region, "s3", "aws4_request"):
        k = hmac.new(k, part.encode(), hashlib.sha256).digest()
    signature = hmac.new(k, to_sign.encode(), hashlib.sha256).hexdigest()
    headers["authorization"] = (f"AWS4-HMAC-SHA256 Credential={key}/{scope}, "
                                f"SignedHeaders={';'.join(signed)}, Signature={signature}")


def main(argv):
    args, headers, body, out = [], {}, b"", None
    it = iter(argv)
    for a in it:
        if a == "--body":
            path = next(it)
            body = sys.stdin.buffer.read() if path == "-" else open(path, "rb").read()
        elif a == "-H":
            name, _, value = next(it).partition(":")
            headers[name.strip().lower()] = value.strip()
        elif a == "-o":
            out = next(it)
        else:
            args.append(a)
    if len(args) not in (2, 3):
        sys.exit(__doc__)
    method, path = args[0].upper(), args[1]
    query = args[2] if len(args) == 3 else ""

    endpoint = urllib.parse.urlsplit(os.environ.get("S3_ENDPOINT", "http://127.0.0.1:7390"))
    region = os.environ.get("S3_REGION", "us-east-1")
    key, secret = os.environ.get("S3_KEY", ""), os.environ.get("S3_SECRET", "")

    if body:
        headers["content-md5"] = base64.b64encode(hashlib.md5(body).digest()).decode()
    if key:
        sign(method, endpoint.netloc, path, query, headers, body, key, secret, region)

    conn = http.client.HTTPConnection(endpoint.hostname, endpoint.port, timeout=30)
    target = quote(path, safe="/") + ("?" + query if query else "")
    conn.request(method, target, body=body if method in ("PUT", "POST") else None, headers=headers)
    resp = conn.getresponse()
    data = resp.read()
    if out:
        with open(out, "wb") as f:
            f.write(data)
    print(resp.status)
    if resp.status >= 300 and data:
        code = re.search(rb"<Code>(.*?)</Code>", data)
        msg = re.search(rb"<Message>(.*?)</Message>", data)
        print((code.group(1) if code else b"").decode(), "-",
              (msg.group(1) if msg else data[:200]).decode(errors="replace"), file=sys.stderr)


if __name__ == "__main__":
    main(sys.argv[1:])