zfs-backup/zfs-backup.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import annotations

from datetime import datetime
from os import getuid
from pathlib import Path
from re import search
from subprocess import PIPE, Popen, call, check_output
from sys import exit
from typing import Any
from uuid import UUID

import config

EXIT_ERROR = 1

ALREADY_DECRYPTED = "Cannot decrypt, already decrypted"
ALREADY_IMPORTED = "Cannot import, already imported"
CANNOT_FIND_BACKUP_DRIVE = "Could not find a backup drive"
DELETE_IN_LOCAL_POOL = "ALERT! Tried to delete in local pool!"
MAPPER_ENTRY_ALREADY_EXISTS = "mapper entry already exists"
NOT_DECRYPTED = "Cannot encrypt, not decrypted"
NOT_IMPORTED = "Cannot export, not imported"
NOT_VALID_ZFS_PATH = "Not a valid ZFSPath"
RUN_AS_ROOT = "Run as root."
SAME_DATASET = "Cannot send incremental snapshots if start and end snapshot are not based on the same dataset"

DISK_BY_UUID = Path("/dev/disk/by-uuid")
MAPPER_PATH = Path("/dev/mapper")


def main():
    _verify_running_as_root()
    local_pool = Pool(config.local_pool_to_backup)
    external_pool = ExternalPool.find_from_dict(config.backup_pools)
    manager = Manager(local_pool, external_pool)
    manager.backup()


# TODO tests
# TODO docstrings
# TODO --skip-import
# TODO --skip-post-backup-scrub
# TODO print time estimation while scrub is in progress
# TODO logging


class Manager:
    """Controls the overall backup process."""

    _map_type: dict[Dataset, Dataset | None]

    def __init__(self, local_pool: Pool, external_pool: ExternalPool):
        self._local_pool = local_pool
        self._external_pool = external_pool

    def backup(self):
        """Backup workflow."""
        self._external_pool.import_()
        try:
            self._external_pool.scrub(skip_if_recently_scrubbed=True)
            self._backup_all_datasets()
            self._external_pool.clean_old_snapshots()
            self._external_pool.scrub(skip_if_recently_scrubbed=False)
        finally:
            self._external_pool.export()

    def _backup_all_datasets(self):
        """Send snapshots of local datasets to the backup pool."""
        print("Backing up datasets")
        for (
            local_dataset,
            remote_dataset,
        ) in self._get_user_filtered_dataset_backup_map().items():
            if remote_dataset is None:
                stream = self._send_absolute(local_dataset)
            else:
                stream = self._send_incremental(local_dataset, remote_dataset)
            ZFS.receive(stream, self._external_pool)

    def _get_user_filtered_dataset_backup_map(self) -> _map_type:
        return self._user_confirm_new_datasets(self._get_raw_dataset_map())

    def _user_confirm_new_datasets(self, dataset_map: _map_type) -> _map_type:
        """For datasets that exist only locally, ask the user if he wants to back it up.

        If he doesn't, remove it from the dict.
        """
        return {
            local_dataset: remote_dataset
            for local_dataset, remote_dataset in dataset_map.items()
            if remote_dataset is not None or self._ask_user_confirmation(local_dataset)
        }

    def _get_raw_dataset_map(self) -> _map_type:
        """Return map from local to remote datasets.

        If the local dataset has no remote match, set the remote dataset to None.
        """
        return {
            local_dataset: self._search_matching_dataset_in_remote_pool(
                local_dataset, self._external_pool
            )
            for local_dataset in self._datasets_to_backup
        }

    def _send_incremental(
        self, local_dataset: Dataset, remote_dataset: Dataset
    ) -> Popen:
        """Send incremental stream between two datasets to remote."""
        start_snapshot = self._get_newest_common_snapshot_with_backup_tags(
            local_dataset, remote_dataset
        )
        end_snapshot = self._get_newest_snapshot_with_backup_tag(local_dataset)
        stream = ZFS.send_incremental(start_snapshot, end_snapshot)
        return stream

    def _send_absolute(self, local_dataset: Dataset) -> Popen:
        """Send absolute stream to remote."""
        snapshot = self._get_newest_snapshot_with_backup_tag(local_dataset)
        return ZFS.send_absolute(snapshot)

    def _get_newest_snapshot_with_backup_tag(self, dataset: Dataset) -> Snapshot:
        """From all snapshots in the dataset, get the newest one that matches the backup tags in the config."""
        snapshots = self._find_snapshots_with_backup_tag(dataset)
        snapshots.sort()
        end_snapshot = snapshots[-1]
        return end_snapshot

    @property
    def _datasets_to_backup(self) -> list[Dataset]:
        """Return local datasets to consider for a backup."""
        return [
            dataset
            for dataset in self._local_pool.datasets
            if dataset.qualified_name not in config.do_not_backup
        ]

    @staticmethod
    def _find_snapshots_with_backup_tag(dataset: Dataset) -> list[Snapshot]:
        """Find snapshots of a dataset that carry the backup tags."""
        regex = _get_regex_matching_snapshots_with_tags(
            [config.snapshot_tag, config.snapshot_interval]
        )
        return [
            snapshot for snapshot in dataset.snapshots if snapshot.matches_regex(regex)
        ]

    @staticmethod
    def _search_matching_dataset_in_remote_pool(
        dataset: Dataset, pool: Pool
    ) -> Dataset | None:
        """For a dataset in the local pool, get the matching dataset in the remote pool.

        If there is no match, return None.
        """
        dataset_to_search = dataset.replace_pool(pool.name)
        for dataset in pool.datasets:
            if dataset == dataset_to_search:
                return dataset

    @classmethod
    def _get_newest_common_snapshot_with_backup_tags(
        cls, local_dataset: Dataset, remote_dataset: Dataset
    ) -> Snapshot:
        """In a local and a remote dataset, find the newest snapshot with backup tags which exists in both datasets.

        Return the local snapshot."""
        common_snapshots = [
            local_snap
            for local_snap in list(cls._find_snapshots_with_backup_tag(local_dataset))
            for remote_snap in list(cls._find_snapshots_with_backup_tag(remote_dataset))
            if local_snap.snapshot_name == remote_snap.snapshot_name
        ]
        common_snapshots.sort()
        return common_snapshots[-1]

    @staticmethod
    def _ask_user_confirmation(dataset: Dataset) -> bool:
        msg = f"Dataset {dataset.qualified_name} does not exist in backup pool. Backup? [y/n] "
        reply = ""
        while not reply:
            reply = input(msg).lower()
            if reply == "y":
                return True
            elif reply == "n":
                return False
            else:
                reply = ""


def _get_regex_matching_snapshots_with_tags(tags: list[str]):
    return "@" + ".*".join(tags)


class ZFSPath:
    def __init__(self, elements: list[str], snapshot_name: str | None = None):
        self._elements = elements
        self._snapshot_name = snapshot_name
        self._is_snapshot = False
        if snapshot_name:
            self._is_snapshot = True
        self._sanity_check()

    def __repr__(self):
        path = "/".join(self._elements)
        if self._is_snapshot:
            return "@".join([path, self._snapshot_name])
        return path

    def __eq__(self, other: ZFSPath):
        return (
            self._elements == other._elements
            and self._is_snapshot == other._is_snapshot
            and self._snapshot_name == other._snapshot_name
        )

    @classmethod
    def from_string(cls, string: str) -> ZFSPath:
        snapshot_name = None
        if "@" in string:
            string, snapshot_name = string.split("@")
        elements: list[str] = string.split("/")
        return ZFSPath(elements, snapshot_name)

    @property
    def snapshot_name(self) -> str:
        return self._snapshot_name

    @property
    def is_snapshot(self) -> bool:
        return self._is_snapshot

    @property
    def pool(self) -> str:
        return self._elements[0]

    @property
    def dataset(self) -> ZFSPath:
        return ZFSPath(self._elements)

    def replace_pool(self, pool: str) -> ZFSPath:
        return ZFSPath([pool] + self._elements[1:].copy(), self.snapshot_name)

    @property
    def name_without_pool(self) -> str:
        name = "/".join(self._elements[1:])
        if self.is_snapshot:
            name = "@".join([name, self.snapshot_name])
        return name

    def _sanity_check(self):
        for item in self._elements:
            if len(item) == 0:
                raise ValueError(NOT_VALID_ZFS_PATH)


def _assert_type(instance: Any, object_type: Any, message: str) -> None:
    if not isinstance(instance, object_type):
        raise TypeError(message)


class CommandInterface:
    _BINARY: Path

    @classmethod
    def _run_command(cls, *args: str) -> None:
        CommandRunner.run([str(cls._BINARY), *args])

    @classmethod
    def _get_output(cls, *args: str) -> str:
        return CommandRunner.get_output([str(cls._BINARY), *args])

    @classmethod
    def _open_stream(cls, *args: str) -> Popen:
        return CommandRunner.send_stream([str(cls._BINARY), *args])

    @classmethod
    def _receive_stream(cls, *args: str, stream: Popen) -> None:
        CommandRunner.receive_from_stream([str(cls._BINARY), *args], stream=stream)


class ZPool(CommandInterface):
    _BINARY = Path("/usr/bin/zpool")

    class _Subcommands:
        # -N: no mount
        # -d: directory to search the pool in
        import_ = ["import", "-N", "-d"]
        export = ["export"]
        scrub = ["scrub"]
        status = ["status"]
        wait = ["wait", "-t", "scrub"]

    @classmethod
    def import_from_directory(cls, pool: Pool, directory: Path) -> None:
        _assert_type(pool, Pool, "Can only import Pool objects")
        cls._run_command(*cls._Subcommands.import_, str(directory), pool.name)

    @classmethod
    def export(cls, pool: Pool) -> None:
        _assert_type(pool, Pool, "Can only export Pool objects")
        cls._run_command(*cls._Subcommands.export, pool.name)

    @classmethod
    def scrub(
        cls,
        pool: Pool,
        *,
        skip_if_recently_scrubbed: bool,
        wait_for_finish: bool = True,
    ) -> None:
        _assert_type(pool, Pool, "Can only scrub pools.")
        cls._scrub_if_necessary(pool, skip_if_recently_scrubbed)
        if wait_for_finish:
            cls._run_command(*cls._Subcommands.wait, pool.name)
        if not cls._is_healthy(pool):
            raise IOError(f"Pool {pool.name} is not healthy.")

    @classmethod
    def _scrub_if_necessary(cls, pool: Pool, skip_if_recently_scrubbed: bool) -> None:
        def do_scrub():
            cls._run_command(*cls._Subcommands.scrub, pool.name)

        if cls._scrub_in_progress(pool):
            return
        if not skip_if_recently_scrubbed:
            do_scrub()
        else:
            if not cls._recently_scrubbed(pool):
                do_scrub()

    @classmethod
    def _recently_scrubbed(cls, pool: Pool) -> bool:
        return datetime.now() - cls._last_scrub(pool) < config.recent_scrub_timedelta

    @classmethod
    def _scrub_in_progress(cls, pool: Pool) -> bool:
        return "scrub in progress" in cls._get_pool_status_output(pool)

    @classmethod
    def _is_healthy(cls, pool: Pool) -> bool:
        return "ONLINE" in cls._get_pool_status_output(pool)

    @classmethod
    def _get_pool_status_output(cls, pool: Pool) -> str:
        return cls._get_output(*cls._Subcommands.status, pool.name)

    @classmethod
    def _last_scrub(cls, pool: Pool) -> datetime:
        return Time.get_last_scrub_from_status_output(cls._get_pool_status_output(pool))


class ZFS(CommandInterface):

    """Interface of `zfs` commands"""

    class _Subcommands:
        class List:
            # -H   Scripting mode, omit headers
            # -d1  maximum depth of 1
            # -o   name output column
            # -t   object type
            # -r   recursive
            _base = ["list", "-H", "-d1", "-o", "name"]
            dataset = [*_base, "-t", "filesystem"]
            snapshot = [*_base, "-t", "snapshot", "-r"]

        class Send:
            # -R  Replicate filesystem
            # -I  send all intermediary snapshots
            absolute = ["send", "-R"]
            incremental = [*absolute, "-I"]

        destroy = ["destroy"]
        # -d  Discard the first element of the "send" snapshot's file system name
        # -F  Force a rollback of the file system to the most recent snapshot before performing the "receive".
        # -u  File system that is associated with the received stream is not mounted.
        receive = ["receive", "-d", "-F", "-u"]

    _BINARY = Path("/usr/bin/zfs")

    @classmethod
    def get_first_level_datasets(cls, pool: Pool) -> list[Dataset]:
        _assert_type(pool, Pool, "Can only get datasets from Pool objects")
        datasets_str = cls._get_output(*cls._Subcommands.List.dataset, pool.name)
        # we omit the first return value, which is the pool, not the dataset
        datasets_list = datasets_str.split("\n")[1:]
        return [Dataset(ZFSPath.from_string(d)) for d in datasets_list]

    @classmethod
    def get_snapshots(cls, dataset: Dataset) -> list[Snapshot]:
        _assert_type(dataset, Dataset, "Can only get snapshots from Dataset object")
        command = [*cls._Subcommands.List.snapshot, dataset.qualified_name]
        snapshot_names: list[str] = cls._get_output(*command).split("\n")
        return [Snapshot(ZFSPath.from_string(name)) for name in snapshot_names]

    @classmethod
    def destroy_snapshot(cls, snapshot: Snapshot):
        _assert_type(snapshot, Snapshot, "Can only destroy snapshots.")
        if config.local_pool_to_backup == snapshot.pool:
            raise Exception(DELETE_IN_LOCAL_POOL)
        cls._run_command(*cls._Subcommands.destroy, str(snapshot))

    @classmethod
    def send_incremental(
        cls, old_snapshot: Snapshot, new_snapshot: Snapshot
    ) -> Popen[str]:
        cls._pre_send_sanity_checks(old_snapshot, new_snapshot)
        print(f"  {old_snapshot} -> {new_snapshot}")
        return cls._open_stream(
            *cls._Subcommands.Send.incremental,
            str(old_snapshot),
            str(new_snapshot),
        )

    @classmethod
    def _pre_send_sanity_checks(cls, old_snapshot: Snapshot, new_snapshot: Snapshot):
        for snapshot in [old_snapshot, new_snapshot]:
            _assert_type(
                snapshot,
                Snapshot,
                f"Cannot send stream, object is not a Snapshot: {snapshot}",
            )
        if old_snapshot.dataset != new_snapshot.dataset:
            raise ValueError(SAME_DATASET)
        if old_snapshot.newer_than(new_snapshot):
            raise ValueError("Old snapshot is newer than new snapshot.")
        if old_snapshot == new_snapshot:
            raise ValueError("Cannot send stream, snapshots are identical.")

    @classmethod
    def receive(cls, stream: Popen[str], pool: Pool):
        _assert_type(pool, Pool, "Can only receive into pools.")
        cls._receive_stream(*cls._Subcommands.receive, pool.name, stream=stream)

    @classmethod
    def send_absolute(cls, snapshot: Snapshot) -> Popen:
        print(f"  {snapshot} -> [new]")
        return cls._open_stream(*cls._Subcommands.Send.absolute, str(snapshot))


class Time:
    @staticmethod
    def _get_date_format() -> str:
        """date format as used by zpool status

        %a: day of week, short version
        %b: month, short version
        %e: day of month, space-padded
        %H, %M, %S: hour, minute, second, zero-padded
        %Y: year
        """
        return "%a %b %e %H:%M:%S %Y"

    @classmethod
    def _extract_datetime_segment(cls, line_with_datetime: str) -> str:
        return line_with_datetime.split(" errors on ")[-1]

    @classmethod
    def _get_line_with_datetime_from_status_output(cls, output: str) -> str:
        lines = output.split("\n")
        for line in lines:
            if line.startswith("  scan: scrub repaired"):
                return line.strip()
        else:
            raise ValueError("Could not find the correct line")

    @classmethod
    def get_last_scrub_from_status_output(cls, output: str) -> datetime:
        line_with_datetime = cls._get_line_with_datetime_from_status_output(output)
        date_str = cls._extract_datetime_segment(line_with_datetime)
        return datetime.strptime(date_str, cls._get_date_format())


class Disk:
    def __init__(self, uuid: UUID, name: str):
        self.uuid = uuid
        self._decrypted: bool = False
        self._was_already_decrypted: bool = False
        self._name = name

    def decrypt(self) -> None:
        if self._decrypted:
            raise ValueError(ALREADY_DECRYPTED, self._name)
        self._was_already_decrypted = Cryptsetup.decrypt(self._path, self._mapper_entry)
        self._decrypted = True

    @property
    def _path(self) -> Path:
        return DISK_BY_UUID / str(self.uuid)

    @property
    def _mapper_entry(self) -> str:
        return f"crypt-{self._name}"

    def encrypt(self):
        if not self._decrypted:
            raise ValueError(NOT_DECRYPTED, self._name)
        if self._was_already_decrypted:
            return
        Cryptsetup.encrypt(self._mapper_entry)
        self._decrypted = False


class Cryptsetup(CommandInterface):
    class _Subcommands:
        close = "close"
        status = "status"
        open = "open"

    _BINARY = Path("/usr/bin/cryptsetup")

    @classmethod
    def encrypt(cls, mapper_entry: str) -> None:
        cls._run_command(cls._Subcommands.close, mapper_entry)

    @classmethod
    def _status(cls, name: str) -> str:
        return cls._get_output(cls._Subcommands.status, name)

    @classmethod
    def decrypt(cls, path: Path, mapper_entry: str) -> bool:
        """Decrypt a disk, return whether it was already decrypted."""
        if cls._mapper_entry_in_use(mapper_entry):
            if cls._already_decrypted(path, mapper_entry):
                return True
            print(MAPPER_ENTRY_ALREADY_EXISTS)
            exit(EXIT_ERROR)
        cls._run_command(cls._Subcommands.open, str(path), mapper_entry)
        return False

    @classmethod
    def _mapper_entry_in_use(cls, mapper_entry: str) -> bool:
        return (MAPPER_PATH / mapper_entry).exists()

    @classmethod
    def _get_device_by_uuid_from_status_output(cls, status: str) -> Path:
        device = cls._get_device_name(status)
        result = cls._resolve_uuid_of_device(device)
        return result

    @classmethod
    def _resolve_uuid_of_device(cls, device):
        for item in DISK_BY_UUID.iterdir():
            if item.readlink() == device:
                result = item
                break
        else:
            print(f"Could not resolve UUID of cryptsetup device {device}")
            exit(EXIT_ERROR)
        # noinspection PyUnboundLocalVariable
        return result

    @classmethod
    def _get_device_name(cls, status):
        for line in status.split("\n"):
            if line.startswith("  device:"):
                device = Path(line.split(":")[-1].strip())
                break
        else:
            print("panic")
            exit(EXIT_ERROR)
        # noinspection PyUnboundLocalVariable
        return device

    @classmethod
    def _already_decrypted(cls, path: Path, mapper_entry: str) -> bool:
        decrypted_disk_path = cls._get_device_by_uuid_from_status_output(
            cls._status(mapper_entry)
        )
        return decrypted_disk_path == path


class Pool:
    def __init__(self, name: str):
        self._name = name
        self._path: ZFSPath = ZFSPath.from_string(name)

    @property
    def datasets(self) -> list[Dataset]:
        return ZFS.get_first_level_datasets(self)

    @property
    def name(self) -> str:
        return self._name


class ExternalPool(Pool):
    def __init__(self, name: str, disk: Disk):
        super().__init__(name)
        self._imported: bool = False
        self._disk = disk

    def import_(self) -> None:
        if self._imported:
            raise ValueError(ALREADY_IMPORTED, self.name)
        self._disk.decrypt()
        ZPool.import_from_directory(self, MAPPER_PATH)
        self._imported = True

    def export(self) -> None:
        if not self._imported:
            raise ValueError(NOT_IMPORTED, self.name)
        self._export_pool()
        self._disk.encrypt()
        self._imported = False
        print(f"Exported {self.name}. Disk {self._disk.uuid} can be removed.")

    def _export_pool(self) -> None:
        ZPool.export(self)

    def scrub(self, skip_if_recently_scrubbed: bool) -> None:
        print(f"Scrubbing {self._name}")
        ZPool.scrub(self, skip_if_recently_scrubbed=skip_if_recently_scrubbed)

    @classmethod
    def find_from_dict(cls, backup_pools: dict[str, UUID]) -> ExternalPool:
        present_uuids = [item.name for item in DISK_BY_UUID.iterdir()]
        for pool_name, disk_uuid in backup_pools.items():
            if str(disk_uuid) in present_uuids:
                disk = Disk(disk_uuid, pool_name)
                print(f"Found disk {pool_name}, {disk_uuid}")
                return cls(pool_name, disk)
        else:
            print(CANNOT_FIND_BACKUP_DRIVE)
            exit(EXIT_ERROR)

    def clean_old_snapshots(self):
        print("Cleaning old snapshots")
        for dataset in self.datasets:
            dataset.clean_old_snapshots()


class Dataset:
    def __init__(self, zfs_path: ZFSPath):
        self._path = zfs_path

    def __repr__(self):
        return self.qualified_name

    def __eq__(self, other: Dataset):
        return self._path == other._path

    @classmethod
    def from_string(cls, qualified_name: str) -> Dataset:
        return cls(ZFSPath.from_string(qualified_name))

    @property
    def qualified_name(self) -> str:
        return str(self._path)

    @property
    def name_without_pool(self) -> str:
        return self._path.name_without_pool

    @property
    def snapshots(self) -> list[Snapshot]:
        """get qualified snapshot names that are direct children of the dataset"""
        return ZFS.get_snapshots(self)

    def replace_pool(self, name: str) -> Dataset:
        return Dataset(self._path.replace_pool(name))

    def clean_old_snapshots(self):
        for interval, number in config.keep_snapshots_per_interval.keys():
            self._clean_old_snapshots_for_interval(interval, number)

    def _clean_old_snapshots_for_interval(self, interval: str, max_count: int):
        regex = _get_regex_matching_snapshots_with_tags([config.snapshot_tag, interval])
        snapshots = [
            snapshot for snapshot in self.snapshots if snapshot.matches_regex(regex)
        ]
        old_snapshots = self._get_n_oldest_snapshots(snapshots, max_count)
        for snapshot in old_snapshots:
            snapshot.destroy()

    @staticmethod
    def _get_n_oldest_snapshots(
        snapshots: list[Snapshot], count: int
    ) -> list[Snapshot]:
        snapshots.sort()
        return snapshots[:count]


class Snapshot:
    def __init__(self, path: ZFSPath):
        self._path = path

    @classmethod
    def from_string(cls, qualified_name: str) -> Snapshot:
        return cls(ZFSPath.from_string(qualified_name))

    def __lt__(self, other: Snapshot) -> bool:
        return self.snapshot_name < other.snapshot_name

    def __repr__(self):
        return str(self._path)

    def newer_than(self, other: Snapshot) -> bool:
        return self > other

    @property
    def snapshot_name(self) -> str:
        return self._path.snapshot_name

    def matches_regex(self, regex: str) -> bool:
        return bool(search(regex, str(self._path)))

    def destroy(self):
        ZFS.destroy_snapshot(self)

    @property
    def pool(self) -> str:
        return self._path.pool

    @property
    def dataset(self) -> Dataset:
        return Dataset(self._path.dataset)


class CommandRunner:
    @staticmethod
    def _sanity_check(cmdline: list[str]) -> None:
        _assert_type(cmdline, list, "Command line must be a list")
        for item in cmdline:
            _assert_type(item, str, "Only strings allowed in command line")

    @staticmethod
    def _to_list(command: str | list[str]) -> list[str]:
        if isinstance(command, str):
            return [command]
        return command

    @classmethod
    def run(cls, cmdline: str | list[str]) -> None:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        call(cmdline)

    @classmethod
    def get_output(cls, cmdline: str | list[str]) -> str:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        return check_output(cmdline).strip().decode()

    @classmethod
    def receive_from_stream(cls, cmdline: str | list[str], stream: Popen) -> None:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        check_output(cmdline, stdin=stream.stdout)

    @classmethod
    def send_stream(cls, cmdline: str | list[str]) -> Popen:
        cmdline = cls._to_list(cmdline)
        cls._sanity_check(cmdline)
        return Popen(cmdline, stdout=PIPE)


def _verify_running_as_root():
    if getuid() > 0:
        print(RUN_AS_ROOT)
        exit(EXIT_ERROR)


if __name__ == "__main__":
    main()