FabledCurator/backend/app/services/backup_s3.py

"""A small S3 client for sending backups off the box: AWS Signature Version 4,
path-style addressing (endpoint/bucket/key), and the calls a backup needs —
put an object (whole, or in parts when its size isn't known up front), read
one back, list a prefix.

FabledCurator puts, lists and reads; 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: a stolen one cannot
wipe the history it writes into.

Built on the standard library's http.client rather than boto3: three calls
don't justify a large dependency, and signing is a pure function the tests
hold to AWS's own published examples. Every call is synchronous, because the
uploads run in Celery workers; the API runs the Test destination call in a
thread.
"""
from __future__ import annotations

import hashlib
import hmac
import http.client
import os
import socket
import ssl
import time
import xml.etree.ElementTree as ET
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit

EMPTY_SHA256 = hashlib.sha256(b"").hexdigest()
TIMEOUT_SECONDS = 30.0
# A single object can be a multi-gigabyte video or archive; a slow link still
# gets this long between bytes before the call gives up.
TRANSFER_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('/')}/{quote(self.bucket, safe='')}"
        return f"{base}/{quote(key, safe='/-_.~')}" if key else f"{base}/"


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


@dataclass(frozen=True)
class ListPage:
    objects: list[StoredObject]
    next_token: str | None     # pass back as `continuation` for the next page


# ── 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 Host, x-amz-date, x-amz-content-sha256 and
    Authorization added. Every header given is signed, Host always."""
    u = urlsplit(url)
    stamp = at.astimezone(UTC).strftime("%Y%m%dT%H%M%SZ")
    day = stamp[:8]
    out = dict(headers)
    out["Host"] = u.netloc
    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()}
    names = sorted(signed)

    query = sorted(
        (_encode(k), _encode(v)) for k, v in parse_qsl(u.query, keep_blank_values=True)
    )
    canonical = "\n".join([
        method.upper(),
        _encode(unquote(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


# ── Transport ────────────────────────────────────────────────────────────────

def _connect(target: S3Target, url: str, timeout: float) -> http.client.HTTPConnection:
    u = urlsplit(url)
    if u.scheme == "https":
        ctx = ssl.create_default_context()
        if not target.verify_tls:
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
        return http.client.HTTPSConnection(
            u.hostname, u.port, timeout=timeout, context=ctx, blocksize=CHUNK,
        )
    if u.scheme != "http":
        raise S3Error(f"{target.endpoint} isn't an http:// or https:// address.")
    return http.client.HTTPConnection(u.hostname, u.port, timeout=timeout, blocksize=CHUNK)


def _request_target(url: str) -> str:
    u = urlsplit(url)
    return (u.path or "/") + (f"?{u.query}" if u.query else "")


def _reach_error(target: S3Target, exc: Exception) -> S3Error:
    if isinstance(exc, TimeoutError):
        return S3Error(f"{target.endpoint} didn't answer in time.")
    if isinstance(exc, socket.gaierror):
        return S3Error(f"Couldn't reach {target.endpoint}: the host name doesn't resolve.")
    if isinstance(exc, ConnectionRefusedError):
        return S3Error(
            f"Couldn't reach {target.endpoint}: nothing is listening there "
            "(check the port)."
        )
    if isinstance(exc, ssl.SSLError):
        return S3Error(f"Couldn't reach {target.endpoint}: TLS failed ({exc.reason}).")
    return S3Error(f"Couldn't reach {target.endpoint}: {exc}.")


def _says(status: int, reason: str, content: bytes) -> str:
    """An S3 error response in words: its Code and Message when it has them."""
    try:
        root = ET.fromstring(content) if content else None
        code = (root.findtext("Code") or "") if root is not None else ""
        message = (root.findtext("Message") or "") if root is not None else ""
    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",
        "XAmzContentSHA256Mismatch": "the file's bytes don't match its recorded checksum",
    }
    said = hints.get(code) or message or reason
    if not code and 400 <= status < 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 {status}{f' {code}' if code else ''}: {said}."


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

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()


def put_file(target: S3Target, key: str, path: Path, *,
             content_type: str = "application/octet-stream",
             payload_sha256: str | None = None,
             pacer: Pacer | None = None,
             now: datetime | None = None) -> None:
    """Upload one file as one object, streamed from disk. `payload_sha256` is
    the file's sha256 when the caller already knows it (the library stores
    every file's), which saves reading the file twice; the bucket checks the
    bytes against it, so a file that changed on disk is refused rather than
    stored under the wrong name."""
    digest = payload_sha256 or file_sha256(path)
    size = path.stat().st_size
    url = target.url(key)
    headers = sign("PUT", url, {"Content-Length": str(size), "Content-Type": content_type},
                   digest, target.access_key, target.secret_key, target.region,
                   now or datetime.now(UTC))
    conn = _connect(target, url, TRANSFER_TIMEOUT_SECONDS)
    try:
        with open(path, "rb") as f:
            conn.request("PUT", _request_target(url),
                         body=pacer.file_chunks(f) if pacer else f, headers=headers)
            response = conn.getresponse()
            content = response.read()
    except (OSError, http.client.HTTPException) as exc:
        raise _reach_error(target, exc) from None
    finally:
        conn.close()
    if response.status // 100 != 2:
        raise S3Error(f"Upload of {key} refused: "
                      f"{_says(response.status, response.reason, content)}")


class ObjectMissing(S3Error):
    """The bucket holds nothing at the key asked for."""


def get_file(target: S3Target, key: str, path: Path, *,
             sha256: str | None = None, size: int | None = None,
             pacer: Pacer | None = None,
             now: datetime | None = None) -> None:
    """Download one object to `path`, streamed to disk. The bytes land under a
    `.partial` name and are renamed into place when complete, so a broken
    transfer never leaves a truncated file at `path`. With `sha256` (and
    `size`) given, the bytes are checked on arrival and a copy that doesn't
    match is thrown away and reported. Raises ObjectMissing when the bucket
    has no such object."""
    url = target.url(key)
    headers = sign("GET", url, {}, EMPTY_SHA256, target.access_key, target.secret_key,
                   target.region, now or datetime.now(UTC))
    partial = path.with_name(path.name + ".partial")
    digest = hashlib.sha256()
    got = 0
    conn = _connect(target, url, TRANSFER_TIMEOUT_SECONDS)
    try:
        conn.request("GET", _request_target(url), headers=headers)
        response = conn.getresponse()
        if response.status == 404:
            response.read()
            raise ObjectMissing(f"{key} is not in the {target.bucket} bucket.")
        if response.status != 200:
            raise S3Error(f"Reading {key} refused: "
                          f"{_says(response.status, response.reason, response.read())}")
        with open(partial, "wb") as f:
            while chunk := response.read(CHUNK):
                if pacer:
                    pacer.take(len(chunk))
                digest.update(chunk)
                got += len(chunk)
                f.write(chunk)
        if size is not None and got != size:
            raise S3Error(f"{key} arrived as {got} bytes, not the {size} expected.")
        if sha256 and digest.hexdigest() != sha256:
            raise S3Error(f"{key} arrived with a different fingerprint "
                          f"({digest.hexdigest()[:12]}…), so it wasn't kept.")
        os.replace(partial, path)
    except (OSError, http.client.HTTPException) as exc:
        partial.unlink(missing_ok=True)
        raise _reach_error(target, exc) from None
    except S3Error:
        partial.unlink(missing_ok=True)
        raise
    finally:
        conn.close()


def list_objects(target: S3Target, prefix: str = "", limit: int = 1000, *,
                 continuation: str | None = None,
                 now: datetime | None = None) -> ListPage:
    """One page (up to `limit`) of the objects under `prefix`, ListObjectsV2."""
    params = {"list-type": "2", "max-keys": str(limit)}
    if prefix:
        params["prefix"] = prefix
    if continuation:
        params["continuation-token"] = continuation
    url = f"{target.url()}?{urlencode(params, quote_via=quote)}"
    headers = sign("GET", url, {}, EMPTY_SHA256, target.access_key, target.secret_key,
                   target.region, now or datetime.now(UTC))
    conn = _connect(target, url, TIMEOUT_SECONDS)
    try:
        conn.request("GET", _request_target(url), headers=headers)
        response = conn.getresponse()
        content = response.read()
    except (OSError, http.client.HTTPException) as exc:
        raise _reach_error(target, exc) from None
    finally:
        conn.close()
    if response.status != 200:
        raise S3Error(f"Listing {target.bucket} refused: "
                      f"{_says(response.status, response.reason, content)}")
    try:
        root = ET.fromstring(content)
    except ET.ParseError:
        raise S3Error(
            f"{target.endpoint} answered a listing with something that isn't S3's XML."
        ) from None
    objects, next_token = [], None
    for item in root.iter():
        tag = item.tag.rsplit("}", 1)[-1]
        if tag == "NextContinuationToken":
            next_token = item.text or None
        if tag != "Contents":
            continue
        fields = {child.tag.rsplit("}", 1)[-1]: (child.text or "") for child in item}
        objects.append(StoredObject(fields.get("Key", ""), int(fields.get("Size") or 0),
                                    fields.get("LastModified", "")))
    return ListPage(objects, next_token)


# ── Multipart: one object sent in parts ─────────────────────────────────────

# Each part is held in memory while it is sent, so this bounds the memory a
# streamed upload uses. S3 allows 10,000 parts, which at this size is 640 GiB
# for one object, and every part but the last must be at least 5 MiB.
PART_SIZE = 64 * CHUNK
# A part the bucket couldn't take (it didn't answer, or said it was busy) is
# sent again after a rest that grows each time: 15 s, then 60 s.
PART_RETRY_WAITS = (15, 60)
_BUSY = frozenset({429, 500, 502, 503, 504})
_MAX_RETRY_AFTER = 300


class Pacer:
    """Holds the bytes a job moves to an average rate, across every request
    it makes, so a backup or a restore leaves room on the link for everything
    else."""

    def __init__(self, mbps: int | None):
        self.rate = mbps * 1_000_000 / 8 if mbps else None   # bytes per second
        self.sent = 0
        self.started = time.monotonic()

    def take(self, n: int) -> None:
        """Count `n` bytes moved, pausing while the job is ahead of its rate."""
        if self.rate:
            self.sent += n
            ahead = self.sent / self.rate - (time.monotonic() - self.started)
            if ahead > 0:
                time.sleep(ahead)

    def chunks(self, data: bytes) -> Iterator[memoryview]:
        view = memoryview(data)
        for at in range(0, len(view), CHUNK):
            piece = view[at:at + CHUNK]
            self.take(len(piece))
            yield piece

    def file_chunks(self, f) -> Iterator[bytes]:
        while piece := f.read(CHUNK):
            self.take(len(piece))
            yield piece


def _send(target: S3Target, method: str, url: str, body: bytes = b"",
          headers: dict[str, str] | None = None, *, timeout: float = TIMEOUT_SECONDS,
          pacer: Pacer | None = None,
          now: datetime | None = None) -> tuple[int, str, http.client.HTTPMessage, bytes]:
    """Sign and send one request with an in-memory body, paced when a pacer
    is given. Returns (status, reason, headers, content)."""
    digest = hashlib.sha256(body).hexdigest()
    signed = sign(method, url, {"Content-Length": str(len(body)), **(headers or {})},
                  digest, target.access_key, target.secret_key, target.region,
                  now or datetime.now(UTC))
    conn = _connect(target, url, timeout)
    try:
        conn.request(method, _request_target(url),
                     body=pacer.chunks(body) if pacer and body else body, headers=signed)
        response = conn.getresponse()
        content = response.read()
    except (OSError, http.client.HTTPException) as exc:
        raise _reach_error(target, exc) from None
    finally:
        conn.close()
    return response.status, response.reason, response.headers, content


def _xml_field(content: bytes, name: str) -> str | None:
    try:
        root = ET.fromstring(content)
    except ET.ParseError:
        return None
    for item in root.iter():
        if item.tag.rsplit("}", 1)[-1] == name:
            return item.text
    return None


def _is_error_body(content: bytes) -> bool:
    try:
        return bool(content) and ET.fromstring(content).tag.rsplit("}", 1)[-1] == "Error"
    except ET.ParseError:
        return False


def _retry_after(headers: http.client.HTTPMessage | None, default: int) -> int:
    try:
        asked = int((headers or {}).get("Retry-After") or default)
    except ValueError:
        asked = default
    return max(0, min(asked, _MAX_RETRY_AFTER))


class Upload:
    """One object written as a series of parts, for data whose size isn't known
    when it starts, such as a database dump streamed from pg_dump.

    The object appears in the bucket only when `complete()` succeeds. Until
    then the parts are invisible, so a stream that fails partway leaves
    nothing behind that looks like a finished backup. `abort()` asks the
    bucket to drop the parts at once; if the abort is refused or never
    arrives, S3 services expire an unfinished upload on their own (Fabled
    Archives after a day)."""

    def __init__(self, target: S3Target, key: str, *,
                 content_type: str = "application/octet-stream",
                 pacer: Pacer | None = None):
        self.target = target
        self.key = key
        self.size = 0
        self.pacer = pacer
        self._etags: list[str] = []
        status, reason, _, content = _send(
            target, "POST", f"{target.url(key)}?uploads", headers={"Content-Type": content_type},
        )
        upload_id = _xml_field(content, "UploadId") if status == 200 else None
        if not upload_id:
            raise S3Error(f"Starting the upload of {key} refused: "
                          f"{_says(status, reason, content)}")
        self.upload_id = upload_id

    def _url(self, **params: str) -> str:
        query = urlencode({**params, "uploadId": self.upload_id}, quote_via=quote)
        return f"{self.target.url(self.key)}?{query}"

    @property
    def parts(self) -> int:
        return len(self._etags)

    def add(self, data: bytes) -> None:
        """Send the next part, resting and sending it again when the bucket
        doesn't answer or says it is busy."""
        n = len(self._etags) + 1
        for wait in (*PART_RETRY_WAITS, None):
            headers = None
            try:
                status, reason, headers, content = _send(
                    self.target, "PUT", self._url(partNumber=str(n)), data,
                    timeout=TRANSFER_TIMEOUT_SECONDS, pacer=self.pacer,
                )
            except S3Error as exc:
                if wait is None:
                    raise S3Error(f"Part {n} of {self.key}: {exc}") from None
                time.sleep(wait)
                continue
            etag = headers.get("ETag") if status == 200 else None
            if etag:
                self._etags.append(etag)
                self.size += len(data)
                return
            if status not in _BUSY or wait is None:
                raise S3Error(f"Part {n} of {self.key} refused: "
                              f"{_says(status, reason, content)}")
            time.sleep(_retry_after(headers, wait))

    def complete(self) -> None:
        """Join the parts into the object. S3 can refuse a completion with an
        error body under a 200 status, so the body is read too."""
        if not self._etags:
            raise S3Error(f"Nothing was sent for {self.key}.")
        body = "<CompleteMultipartUpload>" + "".join(
            f"<Part><PartNumber>{n}</PartNumber><ETag>{etag}</ETag></Part>"
            for n, etag in enumerate(self._etags, start=1)
        ) + "</CompleteMultipartUpload>"
        status, reason, _, content = _send(
            self.target, "POST", self._url(), body.encode(),
            headers={"Content-Type": "application/xml"}, timeout=TRANSFER_TIMEOUT_SECONDS,
        )
        if status != 200 or _is_error_body(content):
            raise S3Error(f"Finishing the upload of {self.key} refused: "
                          f"{_says(status, reason, content)}")

    def abort(self) -> None:
        """Drop the parts sent so far. Never raises: it runs while another
        failure is already being reported, and the bucket expires the parts on
        its own if this doesn't reach it."""
        try:
            _send(self.target, "DELETE", self._url())
        except Exception:  # noqa: BLE001 — best effort, see above
            pass


def put_file_in_parts(target: S3Target, key: str, path: Path, *, sha256: str,
                      content_type: str = "application/octet-stream",
                      pacer: Pacer | None = None) -> int:
    """Upload a large file as a multipart object, read from disk a part at a
    time. The whole file's sha256 is computed on the way and must equal
    `sha256` before the upload is completed; otherwise it is aborted, so a file
    that changed on disk is never stored under its old name. Returns the
    number of bytes sent."""
    upload = Upload(target, key, content_type=content_type, pacer=pacer)
    digest = hashlib.sha256()
    try:
        with open(path, "rb") as f:
            while part := f.read(PART_SIZE):
                digest.update(part)
                upload.add(part)
        if digest.hexdigest() != sha256:
            raise S3Error(f"{path} no longer matches its recorded checksum; not stored.")
        upload.complete()
    except BaseException:
        upload.abort()
        raise
    return upload.size


def object_present(target: S3Target, key: str, size: int) -> bool:
    """Whether `key` is already in the bucket at `size` bytes. Used when a put
    is refused: an append-only bucket refuses to replace an object that is
    already there, which for a content-addressed key means the work is done."""
    try:
        page = list_objects(target, prefix=key, limit=1)
    except S3Error:
        return False
    return any(o.key == key and o.size == size for o in page.objects)