dtf-system/ops/db_backup.py

"""Daily off-server backup of the database to its own bucket.

    python -m ops.db_backup serve                 # the stack's backup service
    python -m ops.db_backup once                  # one backup now
    python -m ops.db_backup list                  # what the bucket holds
    python -m ops.db_backup verify  --identity F  # restore the latest into a scratch database, then drop it
    python -m ops.db_backup restore KEY --identity F --into NAME

Each backup is `pg_dump --format=custom`, checked with `pg_restore --list`,
encrypted with age to BACKUP_AGE_RECIPIENT (a public key) and uploaded to
BACKUP_BUCKET with credentials of its own. The server holds only the public
key: it can write backups but not read them. The private key (the identity)
stays with the owner, off the server, and is needed only to restore.

Old backups are removed by the bucket's lifecycle rule, not by this job, so
its credential never needs to delete; a bucket lock keeps a stolen one from
deleting either. Every run is recorded in dtf_local.backups, which the Kanban
shows on its Integrations tab.
"""
import argparse
import hashlib
import os
import subprocess
import sys
import tempfile
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path
from uuid import uuid4

import boto3
import psycopg
from botocore.config import Config
from psycopg import sql
from psycopg.conninfo import conninfo_to_dict

from app.bootstrap import admin_connect
from app.core.secrets import load as load_secret_files

BRASILIA = timezone(timedelta(hours=-3))
PREFIX = 'db/'
# After a failure the next attempt comes this much later, not a day later.
RETRY = timedelta(hours=1)
# A backup older than this means the daily run has stopped working.
STALE = timedelta(hours=26)
TABLES = ('orders', 'quotes', 'uploads', 'accounts', 'movements', 'payment_intents')


def configured():
    return all(os.environ.get(k) for k in
               ('BACKUP_S3_ENDPOINT', 'BACKUP_BUCKET', 'BACKUP_ACCESS_KEY_ID',
                'BACKUP_SECRET_ACCESS_KEY', 'BACKUP_AGE_RECIPIENT'))


def bucket():
    client = boto3.client('s3', endpoint_url=os.environ['BACKUP_S3_ENDPOINT'],
                          aws_access_key_id=os.environ['BACKUP_ACCESS_KEY_ID'],
                          aws_secret_access_key=os.environ['BACKUP_SECRET_ACCESS_KEY'],
                          region_name=os.environ.get('BACKUP_REGION', 'auto'),
                          config=Config(signature_version='s3v4', s3={'addressing_style': 'path'},
                                        retries={'max_attempts': 5, 'mode': 'standard'}))
    return client, os.environ['BACKUP_BUCKET']


def admin_params(database=None):
    """The administrator's connection, optionally to another database."""
    if os.environ.get('DATABASE_ADMIN_HOST'):
        params = {'host': os.environ['DATABASE_ADMIN_HOST'], 'user': os.environ['DATABASE_ADMIN_USER'],
                  'password': os.environ['DATABASE_ADMIN_PASSWORD'], 'dbname': os.environ['DATABASE_ADMIN_NAME']}
    else:
        params = conninfo_to_dict(os.environ['DATABASE_ADMIN_URL'])
    if database:
        params['dbname'] = database
    return params


def libpq_env(database=None):
    """pg_dump and pg_restore read the credentials from the environment, so the
    password never appears in a process list."""
    names = {'host': 'PGHOST', 'port': 'PGPORT', 'user': 'PGUSER', 'password': 'PGPASSWORD', 'dbname': 'PGDATABASE'}
    return {**os.environ, **{names[k]: str(v) for k, v in admin_params(database).items() if k in names}}


def run(command, **kwargs):
    result = subprocess.run(command, capture_output=True, **kwargs)
    if result.returncode:
        detail = result.stderr.decode(errors='replace').strip().splitlines()
        raise RuntimeError(f'{command[0]} failed: ' + (detail[-1] if detail else f'exit {result.returncode}'))
    return result


def sha256(path):
    digest = hashlib.sha256()
    with open(path, 'rb') as stream:
        for block in iter(lambda: stream.read(1 << 20), b''):
            digest.update(block)
    return digest.hexdigest()


def record(started, status, key=None, size=None, detail=None):
    with admin_connect() as c:
        c.execute('''INSERT INTO dtf_local.backups(id,started_at,finished_at,status,object_key,bytes,detail)
                     VALUES(%s,%s,now(),%s,%s,%s,%s)''',
                  (uuid4(), started, status, key, size, detail))


def run_once():
    """Dump, check, encrypt and upload one backup. Returns its object key."""
    started = datetime.now(timezone.utc)
    key = PREFIX + started.strftime('%Y/%m/dtf-%Y%m%dT%H%M%SZ') + '.dump.age'
    try:
        client, name = bucket()
        with tempfile.TemporaryDirectory(dir=os.environ.get('BACKUP_TMP')) as work:
            dump, sealed = Path(work, 'dtf.dump'), Path(work, 'dtf.dump.age')
            run(['pg_dump', '--format=custom', '--compress=6', '--file', str(dump)], env=libpq_env())
            # A truncated or damaged archive fails here, before it is kept.
            run(['pg_restore', '--list', str(dump)])
            run(['age', '--encrypt', '--recipient', os.environ['BACKUP_AGE_RECIPIENT'],
                 '--output', str(sealed), str(dump)])
            size = sealed.stat().st_size
            client.upload_file(str(sealed), name, key, ExtraArgs={'Metadata': {
                'sha256': sha256(sealed), 'dump-sha256': sha256(dump), 'format': 'pg-custom+age'}})
        record(started, 'ok', key, size)
        print(f'Backup {key} uploaded ({size} bytes).', flush=True)
        return key
    except Exception as error:
        # The message names a command or a provider error, never a credential.
        detail = str(error)[:500]
        try:
            record(started, 'failed', key, None, detail)
        finally:
            print(f'Backup failed: {detail}', file=sys.stderr, flush=True)
        raise


def last_ok():
    with admin_connect() as c:
        row = c.execute("SELECT max(finished_at) FROM dtf_local.backups WHERE status='ok'").fetchone()
    return row[0]


def next_run(now, hour):
    """The next `hour` o'clock in Brasília after `now`."""
    local = now.astimezone(BRASILIA)
    at = local.replace(hour=hour, minute=0, second=0, microsecond=0)
    if at <= local:
        at += timedelta(days=1)
    return at.astimezone(timezone.utc)


def serve():
    if not configured():
        print('Backup not configured: set BACKUP_S3_ENDPOINT, BACKUP_BUCKET, BACKUP_ACCESS_KEY_ID, '
              'BACKUP_SECRET_ACCESS_KEY and BACKUP_AGE_RECIPIENT. Idle.', flush=True)
        while True:
            time.sleep(3600)
    hour = int(os.environ.get('BACKUP_HOUR', '3'))
    # Wait for the schema (db-init) before the first query.
    for _ in range(60):
        try:
            previous = last_ok()
            break
        except psycopg.Error:
            time.sleep(10)
    else:
        previous = last_ok()
    now = datetime.now(timezone.utc)
    # A fresh deployment, or one that missed a day, backs up straight away.
    due = now if previous is None or now - previous > STALE else next_run(now, hour)
    while True:
        time.sleep(max(0, (due - datetime.now(timezone.utc)).total_seconds()))
        try:
            run_once()
            due = next_run(datetime.now(timezone.utc), hour)
        except Exception:
            due = datetime.now(timezone.utc) + RETRY


def listing():
    client, name = bucket()
    keys = []
    for page in client.get_paginator('list_objects_v2').paginate(Bucket=name, Prefix=PREFIX):
        keys += [(o['Key'], o['Size'], o['LastModified']) for o in page.get('Contents', [])]
    return sorted(keys)


def counts(database):
    with psycopg.connect(**admin_params(database)) as c:
        return {t: c.execute(sql.SQL('SELECT count(*) FROM dtf_local.{}').format(sql.Identifier(t))).fetchone()[0]
                for t in TABLES}


def restore(key, identity, into):
    """Restore one backup into a database that is new or holds no DTF data."""
    client, name = bucket()
    with admin_connect() as c:
        c.autocommit = True
        if into == c.execute('SELECT current_database()').fetchone()[0]:
            raise SystemExit('Refusing to restore over the database the stack is using.')
        exists = c.execute('SELECT 1 FROM pg_database WHERE datname=%s', (into,)).fetchone()
        if not exists:
            c.execute(sql.SQL('CREATE DATABASE {}').format(sql.Identifier(into)))
    if exists:
        with psycopg.connect(**admin_params(into)) as c:
            if c.execute("SELECT 1 FROM pg_namespace WHERE nspname='dtf_local'").fetchone():
                raise SystemExit(f'{into} already holds DTF data; choose an empty or new database.')
    with tempfile.TemporaryDirectory(dir=os.environ.get('BACKUP_TMP')) as work:
        sealed, dump = Path(work, 'dtf.dump.age'), Path(work, 'dtf.dump')
        client.download_file(name, key, str(sealed))
        expected = client.head_object(Bucket=name, Key=key)['Metadata'].get('sha256')
        if expected and sha256(sealed) != expected:
            raise SystemExit('The downloaded backup does not match its checksum.')
        run(['age', '--decrypt', '--identity', identity, '--output', str(dump), str(sealed)])
        run(['pg_restore', '--exit-on-error', '--no-owner', '--no-privileges', '--dbname', into, str(dump)],
            env=libpq_env(into))
    return counts(into)


def drop(database):
    with admin_connect() as c:
        c.autocommit = True
        c.execute(sql.SQL('DROP DATABASE IF EXISTS {} WITH (FORCE)').format(sql.Identifier(database)))


def verify(identity, key=None):
    """The restore test: the latest backup into a scratch database, counted, dropped."""
    key = key or listing()[-1][0]
    scratch = 'dtf_verify_' + uuid4().hex[:12]
    try:
        restored = restore(key, identity, scratch)
    finally:
        drop(scratch)
    print(f'PASS: {key} decrypted and restored into a scratch database (now dropped). Rows: ' +
          ', '.join(f'{t} {n}' for t, n in restored.items()))
    return restored


def identity_file(value):
    """The identity as a path, or '-' to paste it: nothing is written to disk
    except a private temporary file removed when the command ends."""
    if value != '-':
        return value, None
    print('Paste the private key (AGE-SECRET-KEY-...) and press Enter:', file=sys.stderr)
    line = sys.stdin.readline().strip()
    handle = tempfile.NamedTemporaryFile('w', prefix='identity-', dir=os.environ.get('BACKUP_TMP'), delete=False)
    os.chmod(handle.name, 0o600)
    handle.write(line + '\n')
    handle.close()
    return handle.name, handle.name


def main(argv=None):
    load_secret_files()
    parser = argparse.ArgumentParser(prog='python -m ops.db_backup')
    sub = parser.add_subparsers(dest='command', required=True)
    sub.add_parser('serve')
    sub.add_parser('once')
    sub.add_parser('list')
    for name in ('verify', 'restore'):
        p = sub.add_parser(name)
        if name == 'restore':
            p.add_argument('key')
            p.add_argument('--into', required=True, help='a new or empty database')
        else:
            p.add_argument('key', nargs='?', help='default: the latest backup')
        p.add_argument('--identity', required=True, help="the private key file, or '-' to paste it")
    args = parser.parse_args(argv)
    if args.command == 'serve':
        serve()
    elif args.command == 'once':
        run_once()
    elif args.command == 'list':
        for key, size, modified in listing():
            print(f'{modified:%Y-%m-%d %H:%M} UTC  {size:>12}  {key}')
    else:
        path, temporary = identity_file(args.identity)
        try:
            if args.command == 'verify':
                verify(path, args.key)
            else:
                rows = restore(args.key, path, args.into)
                print(f'Restored {args.key} into {args.into}. Rows: ' + ', '.join(f'{t} {n}' for t, n in rows.items()))
        finally:
            if temporary:
                os.unlink(temporary)


if __name__ == '__main__':
    main()