FabledGrimoire/src/grimoire/services/backup_s3.py

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

The signing is checked against AWS's published SigV4 examples in
tests/test_services_backup_s3.py.

Grimoire puts and lists, and reads a backup back when an admin restores from
the bucket; it never deletes. The bucket keeps the history and thins it on
its own policy, so a key that may only add and read objects is enough — and
is the safer key to hand a service, because a stolen one cannot wipe the
history it writes into.

Written by hand rather than with boto3: two 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 asyncio
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 tens of 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][/prefix], 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
    if not code and 400 <= response.status_code < 500:
        # S3 refuses with an <Error><Code> body. A 4xx without one came from
        # something else at that address — most often a web server on the
        # default port, when the S3 service listens on a port of its own.
        said = ("the address answered, but not as S3 — check the scheme and "
                "the port (an S3 service usually listens on its own port)")
    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, streamed from disk."""
    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 get_file(target: S3Target, key: str, path: Path, *,
                   transport: httpx.AsyncBaseTransport | None = None,
                   now: datetime | None = None) -> None:
    """Download one object to `path`, streamed to disk."""
    url = target.url(key)
    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=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)
    except httpx.HTTPError as exc:
        raise _reach_error(target, exc) from None


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