Moved rollbacker into the abstraction

Also organized into directories.

Issue #32

src/code/configs.py

@@ -22,7 +22,7 @@
 
 from . import human_interval
 from . import os_utils
-from . import snap_mechanisms
+from .mechanisms import snap_mechanisms
 
 from typing import Iterable, Iterator, Optional
 

src/code/global_flags.py

@@ -14,6 +14,11 @@
 
 import dataclasses
 
+# Format of time used for files.
+TIME_FORMAT = r"%Y%m%d%H%M%S"
+# Length of the string produced by this format.
+TIME_FORMAT_LEN = 14
+
 
 @dataclasses.dataclass
 class _Flags:

src/code/main.py

@@ -20,10 +20,10 @@
 from . import colored_logs
 from . import configs
 from . import global_flags
-from . import rollbacker
 from . import os_utils
-from . import snap_mechanisms
+from . import rollbacker
 from . import snap_operator
+from .mechanisms import snap_mechanisms
 
 
 def _parse_args() -> argparse.Namespace:

src/code/mechanisms/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

src/code/mechanisms/abstract_mechanism.py

@@ -0,0 +1,48 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import enum
+
+
+# The type of snapshot is maintained in two places -
+# 1. Config
+# 2. Snapshot
+class SnapType(enum.Enum):
+    UNKNOWN = "UNKNOWN"
+    BTRFS = "BTRFS"
+
+
+class SnapMechanism(abc.ABC):
+    """Interface with necessary methods to implement snapshotting system.
+
+    Implementations may be based on btrfs, rsync, bcachefs, etc.
+    """
+
+    @abc.abstractmethod
+    def verify_volume(self, mount_point: str) -> bool:
+        """Confirms that the source path can be snapshotted."""
+
+    @abc.abstractmethod
+    def create(self, source: str, destination: str):
+        """Creates a snapshot of source in a destination path."""
+
+    @abc.abstractmethod
+    def delete(self, destination: str):
+        """Deletes an existing snapshot."""
+
+    @abc.abstractmethod
+    def rollback_gen(self, source_dests: list[tuple[str, str]]) -> list[str]:
+        """Returns shell lines which when executed will result in a rollback of snapshots."""
+        # It is okay to leave unimplemented, and raise NotImplementedError().

src/code/snap_mechanisms.py → src/code/mechanisms/btrfs_mechanism.py

@@ -1,38 +1,23 @@
-import abc
-import enum
-import functools
-import logging
-
-from . import global_flags
-from . import os_utils
-
-# TODO: Move btrfs based rollback to this abstraction.
-# TODO: Move btrfs sync to this abstraction.
-
-# The type of snapshot is maintained in two places -
-# 1. Config
-# 2. Snapshot
-class SnapType(enum.Enum):
-    UNKNOWN = "UNKNOWN"
-    BTRFS = "BTRFS"
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
+import logging
 
-class SnapMechanism(abc.ABC):
-    """Interface with necessary methods to implement snapshotting system.
-
-    Example implementations may be based on btrfs, rsync, bcachefs, etc.
-    """
-    @abc.abstractmethod
-    def verify_volume(self, mount_point: str) -> bool:
-        pass
-
-    @abc.abstractmethod
-    def create(self, source: str, destination: str):
-        pass
-
-    @abc.abstractmethod
-    def delete(self, destination: str):
-        pass
+from . import abstract_mechanism
+from . import rollback_btrfs
+from .. import global_flags
+from .. import os_utils
 
 
 def _execute_sh(cmd: str):
@@ -42,16 +27,7 @@ def _execute_sh(cmd: str):
         os_utils.execute_sh(cmd)
 
 
-class _BtrfsSnapMechanism(SnapMechanism):
-    # @staticmethod
-    # def sync(mount_paths: set[str]) -> None:
-    #     for mount_path in sorted(mount_paths):
-    #         if global_flags.FLAGS.dryrun:
-    #             os_utils.eprint(f"Would sync {mount_path}")
-    #             continue
-    #         os_utils.eprint("Syncing ...", flush=True)
-    #         os_utils.execute_sh(f"btrfs subvolume sync {mount_path}")
-
+class BtrfsSnapMechanism(abstract_mechanism.SnapMechanism):
     def verify_volume(self, mount_point: str) -> bool:
         # Based on https://stackoverflow.com/a/32865333/196462
         fstype = os_utils.execute_sh(
@@ -88,10 +64,10 @@ def delete(self, destination: str):
             logging.error("Unable to delete; are you running as root?")
             raise
 
-
-@functools.cache
-def get(snap_type: SnapType) -> SnapMechanism:
-    """Singleton implementation."""
-    if snap_type == SnapType.BTRFS:
-        return _BtrfsSnapMechanism()
-    raise RuntimeError(f"Unknown snap_type {snap_type}")
+    def rollback_gen(self, source_dests: list[tuple[str, str]]) -> list[str]:
+        for source, _ in source_dests:
+            if not self.verify_volume(source):
+                raise RuntimeError(
+                    f"Mount point may no longer be a btrfs volume: {source}"
+                )
+        return rollback_btrfs.rollback_gen(source_dests)

src/code/mechanisms/rollback_btrfs.py

@@ -0,0 +1,141 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import dataclasses
+import datetime
+import os
+
+from .. import global_flags
+
+from typing import Iterable, Optional
+
+
+# This will be cleaned up if it exists by rollback script.
+_PACMAN_LOCK_FILE = "/var/lib/pacman/db.lck"
+
+
+@dataclasses.dataclass
+class _MountAttributes:
+    # The device e.g. /mnt/volume/subv under which we have the subv or its parent is mounted.
+    device: str
+    subvol_name: str
+
+
+def _get_mount_attributes(
+    mount_point: str, mtab_lines: Iterable[str]
+) -> _MountAttributes:
+    # For a mount point, this denotes the longest path that was seen in /etc/mtab.
+    # This is therefore the point where that directory is mounted.
+    longest_match_to_mount_point = ""
+    # Which line matches the mount point.
+    matched_line: str = ""
+    for this_line in mtab_lines:
+        this_tokens = this_line.split()
+        if mount_point.startswith(this_tokens[1]):
+            if len(this_tokens[1]) > len(longest_match_to_mount_point):
+                longest_match_to_mount_point = this_tokens[1]
+                matched_line = this_line
+
+    if not matched_line:
+        raise ValueError(f"Mount point not found: {mount_point}")
+    tokens = matched_line.split()
+
+    if tokens[2] != "btrfs":
+        raise ValueError(f"Mount point is not btrfs: {mount_point} ({tokens[2]})")
+    mount_param: str = ""
+    for mount_param in tokens[3].split(","):
+        subvol_prefix = "subvol="
+        if mount_param.startswith(subvol_prefix):
+            break
+    else:
+        raise RuntimeError(f"Could not find subvol= in {matched_line!r}")
+    subvol_name = mount_param[len(subvol_prefix) :]
+    if tokens[1] != mount_point:
+        nested_subvol = mount_point.removeprefix(tokens[1])
+        assert nested_subvol.startswith("/")
+        subvol_name = nested_subvol
+    return _MountAttributes(device=tokens[0], subvol_name=subvol_name)
+
+
+def _get_mount_attributes_from_mtab(mount_point: str) -> _MountAttributes:
+    return _get_mount_attributes(mount_point, open("/etc/mtab"))
+
+
+def _get_now_str():
+    return datetime.datetime.now().strftime(global_flags.TIME_FORMAT)
+
+
+def rollback_gen(source_dests: list[tuple[str, str]]) -> list[str]:
+    if not source_dests:
+        return ["# No snapshot matched to rollback."]
+
+    sh_lines = [
+        "#!/bin/bash",
+        "# Save this to a script, review and run as root to perform the rollback.",
+        "",
+        "set -uexo pipefail",
+        "",
+    ]
+
+    # Mount all required volumes at root.
+    mount_points: dict[str, str] = {}
+    for source, _ in source_dests:
+        live_subvolume = _get_mount_attributes_from_mtab(source)
+        if live_subvolume.device not in mount_points:
+            mount_pt = f"/run/mount/_yabsnap_internal_{len(mount_points)}"
+            mount_points[live_subvolume.device] = mount_pt
+            sh_lines += [
+                f"mkdir -p {mount_pt}",
+                f"mount {live_subvolume.device} {mount_pt} -o subvolid=5",
+            ]
+
+    now_str = _get_now_str()
+
+    def drop_root_slash(s: str) -> str:
+        if s[0] != "/":
+            raise RuntimeError(f"Could not drop initial / from {s!r}")
+        if "/" in s[1:]:
+            raise RuntimeError(f"Unexpected / after the first one in subvolume {s!r}")
+        return s[1:]
+
+    sh_lines.append("")
+    backup_paths: list[str] = []
+    current_dir: Optional[str] = None
+    for source, dest in source_dests:
+        live_subvolume = _get_mount_attributes_from_mtab(source)
+        backup_subvolume = _get_mount_attributes_from_mtab(os.path.dirname(dest))
+        # The snapshot must be on the same block device as the original (target) volume.
+        assert backup_subvolume.device == live_subvolume.device
+        mount_pt = mount_points[live_subvolume.device]
+        if current_dir != mount_pt:
+            sh_lines += [f"cd {mount_pt}", ""]
+            current_dir = mount_pt
+        live_path = drop_root_slash(live_subvolume.subvol_name)
+        backup_path = f"{drop_root_slash(backup_subvolume.subvol_name)}/rollback_{now_str}_{live_path}"
+        backup_path_after_reboot = (
+            f"{os.path.dirname(dest)}/rollback_{now_str}_{live_path}"
+        )
+        # sh_lines.append(f'[[ -e {backup_path} ]] && btrfs subvolume delete {backup_path}')
+        sh_lines.append(f"mv {live_path} {backup_path}")
+        backup_paths.append(backup_path_after_reboot)
+        sh_lines.append(f"btrfs subvolume snapshot {dest} {live_path}")
+        if os.path.isfile(dest + _PACMAN_LOCK_FILE):
+            sh_lines.append(f"rm {live_path}{_PACMAN_LOCK_FILE}")
+        sh_lines.append("")
+    sh_lines += ["echo Please reboot to complete the rollback.", "echo"]
+    sh_lines.append("echo After reboot you may delete -")
+    for backup_path in backup_paths:
+        sh_lines.append(f'echo "# sudo btrfs subvolume delete {backup_path}"')
+
+    return sh_lines