FabledSteward/steward/backup/s3.py

"""A small S3 client for sending backups: AWS Signature Version 4 over httpx,
path-style addressing (endpoint/bucket/key), the three calls a backup needs.

Fabled Archives speaks S3 with path-style buckets and append-only keys by
default (System 180 ruling: backups go there with an append-only key), so
Steward puts, lists and reads back (to check or restore a backup); it never
deletes. Retention is the target's business. A key that may not read says
so as AccessDenied, and the backup can be uploaded by hand instead.

Written by hand rather than with boto3: four calls don't justify a large
dependency, and signing is a pure function tests can hold to AWS's own
published examples.
"""
from __future__ import annotations

import hashlib
import hmac
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import quote

import httpx

EMPTY_SHA256 = hashlib.sha256(b"").hexdigest()
TIMEOUT_SECONDS = 30.0
# A backup is megabytes; a slow link still gets this long to send one.
UPLOAD_TIMEOUT_SECONDS = 15 * 60
CHUNK = 1 << 20


class S3Error(Exception):
    """Why a call to the bucket failed, said to the person."""


@dataclass(frozen=True)
class S3Target:
    endpoint: str          # http(s)://host[:port], no bucket
    bucket: str
    access_key: str
    secret_key: str
    region: str = "us-east-1"
    verify_tls: bool = True

    def url(self, key: str = "") -> str:
        base = f"{self.endpoint.rstrip('/')}/{self.bucket}"
        return f"{base}/{key}" if key else f"{base}/"


@dataclass(frozen=True)
class StoredObject:
    key: str
    size: int
    last_modified: str


# ── Signing ──────────────────────────────────────────────────────────────────

def _encode(text: str, keep_slash: bool = False) -> str:
    return quote(text, safe="-_.~" + ("/" if keep_slash else ""))


def _hmac(key: bytes, msg: str) -> bytes:
    return hmac.new(key, msg.encode(), hashlib.sha256).digest()


def sign(method: str, url: str, headers: dict[str, str], payload_sha256: str,
         access_key: str, secret_key: str, region: str, at: datetime,
         service: str = "s3") -> dict[str, str]:
    """The request's headers with x-amz-date, x-amz-content-sha256 and
    Authorization added. Every header given is signed, Host always."""
    u = httpx.URL(url)
    stamp = at.astimezone(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    day = stamp[:8]
    out = dict(headers)
    out["x-amz-date"] = stamp
    out["x-amz-content-sha256"] = payload_sha256
    signed = {k.lower(): " ".join(str(v).split()) for k, v in out.items()}
    signed["host"] = u.netloc.decode()
    names = sorted(signed)

    query = sorted((_encode(k), _encode(v)) for k, v in u.params.multi_items())
    canonical = "\n".join([
        method.upper(),
        _encode(u.path or "/", keep_slash=True),
        "&".join(f"{k}={v}" for k, v in query),
        "".join(f"{n}:{signed[n]}\n" for n in names),
        ";".join(names),
        payload_sha256,
    ])
    scope = f"{day}/{region}/{service}/aws4_request"
    to_sign = "\n".join(["AWS4-HMAC-SHA256", stamp, scope,
                         hashlib.sha256(canonical.encode()).hexdigest()])
    key = _hmac(_hmac(_hmac(_hmac(f"AWS4{secret_key}".encode(), day), region), service), "aws4_request")
    signature = hmac.new(key, to_sign.encode(), hashlib.sha256).hexdigest()
    out["Authorization"] = (f"AWS4-HMAC-SHA256 Credential={access_key}/{scope}, "
                            f"SignedHeaders={';'.join(names)}, Signature={signature}")
    return out


# ── Calls ────────────────────────────────────────────────────────────────────

def _says(response: httpx.Response) -> str:
    """An S3 error response in words: its Code and Message when it has them."""
    try:
        root = ET.fromstring(response.content)
        code = root.findtext("Code") or ""
        message = root.findtext("Message") or ""
    except ET.ParseError:
        code, message = "", ""
    hints = {
        "SignatureDoesNotMatch": "the secret key is wrong",
        "InvalidAccessKeyId": "the access key isn't known there",
        "NoSuchBucket": "there is no such bucket",
        "AccessDenied": "the key isn't allowed to do that",
    }
    said = hints.get(code) or message or response.reason_phrase
    return f"HTTP {response.status_code}{f' {code}' if code else ''}: {said}."


def _reach_error(target: S3Target, exc: httpx.HTTPError) -> S3Error:
    if isinstance(exc, httpx.TimeoutException):
        return S3Error(f"{target.endpoint} didn't answer in time.")
    return S3Error(f"Couldn't reach {target.endpoint}: {exc}.")


def file_sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(CHUNK):
            h.update(chunk)
    return h.hexdigest()


async def put_file(target: S3Target, key: str, path: Path, *,
                   transport: httpx.AsyncBaseTransport | None = None,
                   now: datetime | None = None) -> None:
    """Upload one file as one object."""
    import asyncio
    digest = await asyncio.to_thread(file_sha256, path)
    size = path.stat().st_size
    url = target.url(key)
    headers = sign("PUT", url, {"Content-Length": str(size), "Content-Type": "application/x-tar"},
                   digest, target.access_key, target.secret_key, target.region,
                   now or datetime.now(timezone.utc))

    async def body():
        with open(path, "rb") as f:
            while chunk := f.read(CHUNK):
                yield chunk

    try:
        async with httpx.AsyncClient(verify=target.verify_tls, timeout=UPLOAD_TIMEOUT_SECONDS,
                                     transport=transport) as client:
            response = await client.put(url, headers=headers, content=body())
    except httpx.HTTPError as exc:
        raise _reach_error(target, exc) from None
    if response.status_code // 100 != 2:
        raise S3Error(f"Upload of {key} refused: {_says(response)}")


async def list_objects(target: S3Target, prefix: str = "", limit: int = 1000, *,
                       transport: httpx.AsyncBaseTransport | None = None,
                       now: datetime | None = None) -> list[StoredObject]:
    """Up to `limit` objects under `prefix` (ListObjectsV2, one page)."""
    params = {"list-type": "2", "max-keys": str(limit)}
    if prefix:
        params["prefix"] = prefix
    url = str(httpx.URL(target.url(), params=params))
    headers = sign("GET", url, {}, EMPTY_SHA256, target.access_key, target.secret_key,
                   target.region, now or datetime.now(timezone.utc))
    try:
        async with httpx.AsyncClient(verify=target.verify_tls, timeout=TIMEOUT_SECONDS,
                                     transport=transport) as client:
            response = await client.get(url, headers=headers)
    except httpx.HTTPError as exc:
        raise _reach_error(target, exc) from None
    if response.status_code != 200:
        raise S3Error(f"Listing {target.bucket} refused: {_says(response)}")
    try:
        root = ET.fromstring(response.content)
    except ET.ParseError:
        raise S3Error(f"{target.endpoint} answered a listing with something that isn't S3's XML.") from None
    out = []
    for item in root.iter():
        if item.tag.rsplit("}", 1)[-1] != "Contents":
            continue
        fields = {child.tag.rsplit("}", 1)[-1]: (child.text or "") for child in item}
        out.append(StoredObject(fields.get("Key", ""), int(fields.get("Size") or 0),
                                fields.get("LastModified", "")))
    return out


async def get_file(target: S3Target, key: str, path: Path, *,
                   transport: httpx.AsyncBaseTransport | None = None,
                   now: datetime | None = None) -> int:
    """Download one object to `path`, streamed; returns its size. A partial
    download is removed."""
    url = target.url(key)
    headers = sign("GET", url, {}, EMPTY_SHA256, target.access_key, target.secret_key,
                   target.region, now or datetime.now(timezone.utc))
    size = 0
    try:
        async with httpx.AsyncClient(verify=target.verify_tls, timeout=UPLOAD_TIMEOUT_SECONDS,
                                     transport=transport) as client:
            async with client.stream("GET", url, headers=headers) as response:
                if response.status_code != 200:
                    await response.aread()
                    raise S3Error(f"Reading {key} refused: {_says(response)}")
                with open(path, "wb") as f:
                    async for chunk in response.aiter_bytes(CHUNK):
                        f.write(chunk)
                        size += len(chunk)
    except httpx.HTTPError as exc:
        path.unlink(missing_ok=True)
        raise _reach_error(target, exc) from None
    except S3Error:
        path.unlink(missing_ok=True)
        raise
    return size