WIP, rewrite in Python

btv

@@ -0,0 +1,245 @@
+#! /usr/bin/env python3
+# encoding: utf-8
+
+import configparser
+from contextlib import contextmanager
+from dataclasses import dataclass
+import datetime
+import json
+import os
+import shlex
+import sys
+
+CONFIG = "/etc/btv/config.ini"
+
+# ------------------------------------------------------------------------------
+# Global
+cfg = configparser.ConfigParser()
+cfg.read(CONFIG)
+
+class UsageError(Exception):
+	pass
+
+# ------------------------------------------------------------------------------
+# Aux
+
+@contextmanager
+def chdir(new_dir):
+	previous_dir = os.getcwd()
+	os.chdir(new_dir)
+	
+	try:
+		yield
+	finally:
+		os.chdir(previous_dir)
+
+# ------------------------------------------------------------------------------
+# Functions
+
+def verify_keyfile():
+	"""
+	Throws a RuntimeError if the keyfile does not exist, is empty, or has the
+	wrong set of permissions.
+	"""
+	
+	f = cfg.get("main", "keyfile")
+	
+	if not os.path.exists(f):
+		raise RuntimeError("does not exist")
+	
+	s = os.stat(f)
+	
+	if s.st_uid != 0:
+		raise RuntimeError("keyfile is not owned by root")
+	
+	if s.st_mode & 0b100111111 != 256:
+		raise RuntimeError("keyfile is readable by others")
+
+def verify_user():
+	"""
+	Throws a RuntimeError if the user running the program is not root.
+	"""
+	
+	if os.getuid() != 0:
+		raise RuntimeError("this can only be used by root")
+
+@dataclass
+class Snapshot:
+	name: str
+	rank: int
+
+def get_snapshot(name):
+	"""
+	Returns a Snapshot with the requested name iff such exists. Otherwise,
+	returns None.
+	"""
+	
+	rootdir = cfg.get("snap", "dir")
+	
+	try:
+		with open(os.path.join(rootdir, name, ".meta")) as f:
+			meta = json.load(f)
+			return Snapshot(
+				name,
+				meta["rank"]
+			)
+	except FileNotFoundError:
+		pass
+
+def list_snapshots():
+	"""
+	Returns a list of all existing Snapshots sorted oldest to newest.
+	"""
+	
+	snaps = []
+	rootdir = cfg.get("snap", "dir")
+	
+	for snapdir in os.listdir(rootdir):
+		snaps.append(get_snapshot(snapdir))
+	
+	return snaps
+
+# ------------------------------------------------------------------------------
+# Verbs
+
+def do_create(args, context):
+	"""
+	Bulk of the work happens here.
+	
+	- makes snapshots
+	- calls serialization if needed
+	- calls gc if needed
+	"""
+	
+	snap_path = os.path.join(
+		cfg.get("snap", "dir"),
+		datetime.datetime.now().strftime("%Y-%m-%d.%H%M%S")
+	)
+	
+	# determine the rank of the new snapshot
+	snaps = list_snapshots()
+	snaps_since_rank_1 = 1
+	snaps_since_rank_2 = 1
+	
+	for snap in snaps:
+		if snap.rank == 1:
+			snaps_since_rank_1 = 1
+		else:
+			snaps_since_rank_1 += 1
+		
+		if snap.rank == 2:
+			snaps_since_rank_1 = 1
+			snaps_since_rank_2 = 1
+		else:
+			snaps_since_rank_2 += 1
+	
+	if snaps_since_rank_2 >= cfg.getint("snap", "rank_2_interval"):
+		rank = 2
+	elif snaps_since_rank_1 >= cfg.getint("snap", "rank_1_interval"):
+		rank = 1
+	else:
+		rank = 0
+	
+	print(">>> Creating Rank %d snapshot" %(rank), snap_path)
+	
+	os.makedirs(snap_path)
+	
+	ignore_prefix = cfg.get("subvol", "ignore_prefix")
+	
+	with chdir(cfg.get("subvol", "dir")):
+		for sub in os.listdir():
+			if sub.startswith(ignore_prefix):
+				continue
+			
+			os.system(shlex.join((
+				"btrfs", "subvolume", "snapshot", "-r",
+				sub,
+				os.path.join(snap_path, sub)
+			)))
+		
+		with open(os.path.join(snap_path, ".meta"), "w") as f:
+			json.dump({
+				"rank": rank
+			}, f)
+	
+	if rank == 2:
+		print("Rank 2, serialize and GC now")
+
+def do_list(args, context):
+	"""
+	Print a list of existing snapshots.
+	"""
+	
+	print("R Name")
+	
+	for snap in list_snapshots():
+		print("%d %s" %(snap.rank, snap.name))
+
+def do_drop(args, context):
+	"""
+	Drop a snapshot in args[0].
+	"""
+	
+	snap_name = args[0]
+	snap = get_snapshot(snap_name)
+	
+	if not snap:
+		print("!!! %s is not a snapshot" %(snap_name))
+		sys.exit(2)
+	
+	snap_path = os.path.join(cfg.get("snap", "dir"), snap_name)
+	
+	for dir in os.listdir(snap_path):
+		if dir[0] != ".":
+			btrfs_snap = os.path.join(snap_path, dir)
+			os.system(shlex.join(("btrfs", "subvolume", "delete", btrfs_snap)))
+	
+	os.remove(os.path.join(snap_path, ".meta"))
+	os.rmdir(snap_path)
+
+verb_router = {
+	"snapshot": (do_create, None),
+	"list": (do_list, None),
+	"drop": (do_drop, None),
+	# ~ "prune": (do_prune, None),
+	# ~ "stream-full": (do_stream, "full")
+}
+
+# ------------------------------------------------------------------------------
+# Entry point
+
+def print_help():
+	print("""Usage:
+  {cmd} VERB [args]
+
+Verbs:
+  snapshot [--hooks]
+    Create a snapshot. If --hooks is passed, do all optional work.
+  
+  list
+    List snapshots and their rank.
+  
+  drop SNAPSHOT
+    Drops (removes) SNAPSHOT. Use the full name as provided by {cmd} list.
+""".format(cmd=sys.argv[0]))
+
+if __name__ == "__main__":
+	try:
+		verb = sys.argv[1].lower()
+		assert verb in verb_router
+	
+	except:
+		print_help()
+		sys.exit(1)
+	
+	verify_user()
+	verify_keyfile()
+	
+	fn, ctx = verb_router[verb]
+	
+	try:
+		fn(sys.argv[2:], ctx)
+	except UsageError:
+		print()
+		print_help()
+		sys.exit(1)

btv.zsh

@@ -0,0 +1,105 @@
+#! /bin/zsh
+
+setopt extendedglob
+
+function die() {
+	echo "!!! $1"
+	exit 1
+}
+
+### read the config file
+source /etc/btv/config
+
+### verify keyfile sanity
+if [[ ! -a "${KEYFILE}" ]]; then
+	echo "The keyfile (${KEYFILE}) does not exist."
+	exit 2
+fi
+
+KEYFILE_OWNER=$(stat -c "%u" $KEYFILE)
+KEYFILE_ACCESS=$(stat -c "%a" $KEYFILE)
+
+if [[ ${KEYFILE_OWNER} -ne 0 ]]; then
+	echo "The keyfile (${KEYFILE}) is owned by UID ${KEYFILE_OWNER}, it should be 0 (root)."
+	exit 2
+fi
+
+if [[ ${KEYFILE_ACCESS} -ne 600 ]]; then
+	echo "The keyfile (${KEYFILE}) has permissions ${KEYFILE_ACCESS}, it should be 600 (rw-------)."
+	exit 2
+fi
+
+VERB="$1"
+shift
+
+function drop_snap {
+	SNAP="${1}"
+	
+	for S in "${SNAP}/"*
+	do
+		btrfs subvolume delete "$S"
+	done
+	
+	rmdir "${SNAP}"
+}
+
+case "$VERB" in
+	(snap*)
+		T="$(date +%Y-%m-%d.%H%M%S)"
+		SNAP="${SNAP_DIR}/${T}"
+		echo "Creating snapshot ${SNAP}"
+		mkdir "${SNAP}" || die "failed to mkdir ${SNAP}"
+		
+		cd "${SUBVOL_DIR}" || die "failed to cd ${SUBVOL_DIR}"
+		
+		for SUB in ^${SUBVOL_IGNORE_PREFIX}*
+		do
+			btrfs subvolume snapshot -r "${SUB}" "${SNAP}/${SUB}" || die "failed to btrfs snap ${SUB}"
+		done
+		
+		;;
+	
+	(list)
+		cd "${SNAP_DIR}"
+		
+		for SNAP in *
+		do
+			if [[ -f "${SNAP}/streamed" ]]; then
+				printf "S"
+			else
+				printf " "
+			fi
+		done
+		
+		echo
+		echo "S = streamed to storage"
+		echo "R = scheduled for removal"
+		
+		;;
+	
+	(drop)
+		SNAP="${SNAP_DIR}/${1}"
+		
+		if [[ ! -d ${SNAP} ]]; then
+			echo "Snapshot ${1} does not exist"
+			exit 3
+		fi
+		
+		echo "Dropping $SNAP"
+		drop_snap "$SNAP"
+		;;
+	
+	(clean*)
+		echo ">>> gc"
+		;;
+	
+	(stream)
+		echo ">>> stream $1"
+		;;
+	
+	(*)
+		echo "Unknown verb: ${VERB}"
+		exit 3
+		;;
+esac
+

cfg/config.ini

@@ -0,0 +1,32 @@
+[main]
+# Openssl password file (first line is the password).
+keyfile = /etc/btv/key
+
+[subvol]
+# Directory containing all subvolumes.
+dir = /mnt/pool/subvol
+
+# Prefix of subvolumes that should be ignored.
+ignore_prefix = _
+
+[snap]
+# Directory to place local snapshots in.
+dir = /mnt/pool/snap
+
+# Every x-th snapshot is promoted to Rank 1 so it gets to stay around longer.
+# Default is 6, i.e. once per hour assuming 10-minute service timer.
+rank_1_interval=6
+
+# Every x-th snapshot is promoted to Rank 2 so it is serialized to storage.
+# Default is 24, i.e. every 4 hours provided the assumption in the previous
+# comment holds.
+rank_2_interval=24
+
+[prune]
+# How many snapshots of each Rank to keep around.
+# Note: pruning is only performed when a new Rank 2 snapshot is created.
+# Default values are 25, 36, and 42, i.e. 5 hours for Rank 0, 2 days for Rank 1,
+# and a week for Rank 2.
+rank_0_count=25
+rank_1_count=36
+rank_2_count=42

cfg/key

@@ -0,0 +1,3 @@
+guest
+
+# the first line of this file is the password