#!/usr/bin/python3
"""Forced-command wrapper restricting what a backup puller may do over SSH.

Installed on the *source* machine and pinned via command="..." in the backup
user's authorized_keys. It accepts a small fixed verb vocabulary instead of
arbitrary zfs command lines, so there is no shell and nothing to quote-escape.
"""

import datetime
import os
import re
import shlex
import subprocess
import sys
import syslog

ZFS = "/usr/sbin/zfs"
ALLOW_FILE = "/etc/zfs-backup-allow.conf"

# Matches the timestamp formats zfs-auto-snapshot produces: 20260719_1351,
# 20250101-120000, 20250101120000. Deliberately narrow -- it rules out '%'
# (zfs destroy range syntax), wildcards, and path separators.
SNAPSHOT_RE = re.compile(r"^[0-9]{8}[_-]?[0-9]{4,6}$")


class Rejected(Exception):
    pass


def load_allowed_datasets(path=ALLOW_FILE):
    datasets = []
    with open(path) as f:
        for line in f:
            line = line.split("#", 1)[0].strip()
            if line:
                datasets.append(line)
    if not datasets:
        raise Rejected(f"no datasets configured in {path}")
    return datasets


def check_dataset(dataset, allowed):
    if dataset not in allowed:
        raise Rejected(f"dataset not allowed: {dataset}")
    return dataset


def check_snapshot(name):
    if not SNAPSHOT_RE.match(name):
        raise Rejected(f"invalid snapshot name: {name}")
    return name


def parse_snapshot_time(name):
    """Parse one of our timestamp names, or return None if it isn't ours."""
    for fmt in ("%Y%m%d_%H%M", "%Y%m%d-%H%M%S", "%Y%m%d%H%M%S"):
        try:
            return datetime.datetime.strptime(name, fmt)
        except ValueError:
            continue
    return None


def newest_snapshot(dataset):
    """The newest snapshot on `dataset` whose name is one of ours.

    Snapshots in other naming schemes are ignored -- they are not ours to
    reason about, and they must not mask our own newest.
    """
    result = subprocess.run(
        [ZFS, "list", "-H", "-o", "name", "-t", "snapshot", "-d", "1", dataset],
        stdout=subprocess.PIPE,
    )
    if result.returncode != 0:
        raise Rejected(f"could not list snapshots of {dataset}")

    newest, newest_time = None, None
    for line in result.stdout.decode().splitlines():
        _, _, snapshot = line.partition("@")
        when = parse_snapshot_time(snapshot)
        if when is not None and (newest_time is None or when > newest_time):
            newest, newest_time = snapshot, when
    return newest


def check_not_newest(dataset, name, lookup):
    """Refuse to destroy the most recent snapshot we have of a dataset.

    Nothing in the backup pipeline should ever ask for this: the newest shared
    snapshot is what the next incremental sends from. Enforcing it here rather
    than trusting the caller means a compromised backup machine cannot walk a
    dataset's history off the end one snapshot at a time.
    """
    if lookup(dataset) == name:
        raise Rejected(f"refusing to destroy the newest snapshot of {dataset}: {name}")
    return name


def build_command(argv, allowed, newest_lookup=newest_snapshot):
    verb, args = argv[0], argv[1:]

    if verb == "ping":
        require_args(verb, args, 0)
        return None

    if verb == "list-snapshots":
        require_args(verb, args, 0)
        return [ZFS, "list", "-j", "-o", "name", "-t", "snapshot", "-d", "1", *allowed]

    if verb == "snapshot":
        require_args(verb, args, 2)
        dataset = check_dataset(args[0], allowed)
        name = check_snapshot(args[1])
        return [ZFS, "snapshot", "-r", f"{dataset}@{name}"]

    if verb == "send":
        require_args(verb, args, 3)
        dataset = check_dataset(args[0], allowed)
        from_snap = check_snapshot(args[1])
        to_snap = check_snapshot(args[2])
        return [ZFS, "send", "-R", "-e", "-h", "-L", "-c", "-I", f"@{from_snap}",
                f"{dataset}@{to_snap}"]

    if verb == "send-full":
        require_args(verb, args, 2)
        dataset = check_dataset(args[0], allowed)
        to_snap = check_snapshot(args[1])
        return [ZFS, "send", "-R", "-e", "-h", "-L", "-c", f"{dataset}@{to_snap}"]

    if verb == "destroy-snapshot":
        require_args(verb, args, 2)
        dataset = check_dataset(args[0], allowed)
        name = check_not_newest(dataset, check_snapshot(args[1]), newest_lookup)
        return [ZFS, "destroy", "-r", f"{dataset}@{name}"]

    raise Rejected(f"unknown verb: {verb}")


def require_args(verb, args, count):
    if len(args) != count:
        raise Rejected(f"{verb} takes {count} argument(s), got {len(args)}")


def main():
    syslog.openlog("zfs-backup-command", syslog.LOG_PID, syslog.LOG_AUTHPRIV)

    # Our own argv comes from the command="..." line in authorized_keys, which
    # the client cannot influence -- it only supplies SSH_ORIGINAL_COMMAND. So
    # an allowlist path here is server-side configuration, not client input.
    allow_file = sys.argv[1] if len(sys.argv) > 1 else ALLOW_FILE

    original = os.environ.get("SSH_ORIGINAL_COMMAND", "")
    peer = os.environ.get("SSH_CONNECTION", "?").split(" ")[0]

    try:
        argv = shlex.split(original)
        if not argv:
            raise Rejected("no command given (interactive login is not permitted)")
        allowed = load_allowed_datasets(allow_file)
        command = build_command(argv, allowed)
    except (Rejected, ValueError) as e:
        syslog.syslog(syslog.LOG_WARNING, f"denied from {peer}: {e}: {original!r}")
        print(f"zfs-backup-command: denied: {e}", file=sys.stderr)
        return 1
    except OSError as e:
        syslog.syslog(syslog.LOG_ERR, f"config error: {e}")
        print(f"zfs-backup-command: {e}", file=sys.stderr)
        return 1

    if command is None:
        syslog.syslog(syslog.LOG_INFO, f"allowed from {peer}: ping")
        print("pong")
        return 0

    syslog.syslog(syslog.LOG_INFO, f"allowed from {peer}: {shlex.join(command)}")
    return subprocess.run(command).returncode


if __name__ == "__main__":
    sys.exit(main())
