btv/btv

#! /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("crypto", "keyfile")
	
	if not os.path.exists(f):
		raise RuntimeError("keyfile does not exist")
	
	s = os.stat(f)
	
	if s.st_size < 5:
		raise RuntimeError("password in the keyfile is too short")
	
	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 __post_init__(self):
		self.path = os.path.join(cfg.get("snap", "dir"), self.name)
		self.subvolumes = set()
		self.notes = set()
	
	@classmethod
	def load(cls, path):
		with open(os.path.join(path, ".meta")) as f:
			meta = json.load(f)
			snap = cls(meta["name"], meta["rank"])
			snap.subvolumes.update(meta["subvolumes"])
			snap.notes.update(meta["notes"])
			snap.path = path
			snap.check()
			
			return snap
	
	def dump(self):
		if not os.path.exists(self.path):
			os.makedirs(self.path)
		
		with open(os.path.join(self.path, ".meta"), "w") as f:
			meta = {
				"name": self.name,
				"rank": self.rank,
				"subvolumes": sorted(self.subvolumes),
				"notes": sorted(self.notes)
			}
			
			json.dump(meta, f)
	
	def check(self):
		## verify all declared subvolumes exist
		contents = set(os.listdir(self.path))
		
		if self.subvolumes - contents:
			self.notes.add("missing subvolumes")
	
	def drop(self):
		for subvolume in self.subvolumes:
			btrfs_snap = os.path.join(self.path, subvolume)
			
			if not os.path.exists(btrfs_snap):
				print(" !! subvolume not found: %s" %(btrfs_snap))
				return 1
			
			error = os.system(shlex.join(("btrfs", "subvolume", "delete", btrfs_snap)))
			
			if error:
				print(" !! error %d" %(error))
				return error
		
		os.remove(os.path.join(self.path, ".meta"))
		os.rmdir(self.path)
		
		return 0

def get_snapshot_by_name(name):
	"""
	Returns a Snapshot with the requested name iff such exists. Otherwise,
	returns None.
	"""
	
	rootdir = cfg.get("snap", "dir")
	
	try:
		return Snapshot.load(os.path.join(rootdir, name))
	except FileNotFoundError:
		pass

def list_snapshots(min_rank=0):
	"""
	Returns a list of all existing Snapshots sorted oldest to newest.
	"""
	
	snaps = []
	rootdir = cfg.get("snap", "dir")
	
	for name in os.listdir(rootdir):
		s = get_snapshot_by_name(name)
		
		if s and s.rank >= min_rank:
			snaps.append(s)
	
	return snaps

def delete_snapshot(snap):
	"""
	Drops a snapshot referenced by the object.
	"""
	

def serialize(snap, outdir, key, snap_from=None):
	"""
	Serializes a snapshot (differentially if snap_from is passed, fully
	otherwise) into a new directory under outdir.
	
	Key is the path to an openssl-style passphrase file.
	
	Snap and snap_from are Snapshot objects.
	
	Returns 0 on success.
	"""
	
	## prepare directories
	##
	if snap_from:
		name = "%s diff from %s" %(snap.name, snap_from.name)
	else:
		name = "%s full" %(snap.name)
	
	directory = os.path.join(outdir, name)
	os.makedirs(directory)
	
	## preflight checks
	##
	if snap_from:
		# I don't want to handle (dis)appearing subvolumes
		if not snap.subvolumes == snap_from.subvolumes:
			print(" !! snapshots have different subvolumes; please do a full stream")
			print("    %s: %s" %(snap.name, " ".join(snap.subvolumes)))
			print("    %s: %s" %(snap_from.name, " ".join(snap_from.subvolumes)))
			
			return 1
	
	## serialize each subvolume or subvolume pair
	##
	for subvolume in snap.subvolumes:
		if snap_from:
			btrfs_send = 'btrfs send -p "%s" "%s"' %(
				os.path.join(snap_from.path, subvolume),
				os.path.join(snap.path, subvolume)
			)
		else:
			btrfs_send = 'btrfs send "%s"' %(
				os.path.join(snap.path, subvolume)
			)
		
		error = os.system('%s | zstd | openssl enc -e -aes-256-cbc -pbkdf2 -salt -pass "file:%s" > "%s.btrfs.zst.aes"' %(
			btrfs_send,
			cfg.get("crypto", "keyfile"),
			os.path.join(directory, subvolume)
		))
		
		if error:
			print(" !! failed to serialize %s" %(subvolume))
			return error
	
	## calculate checksums and add a self-check executable
	## FIXME calculate this on the fly, re-reading is expensive
	previous_wd = os.getcwd()
	os.chdir(directory)
	os.system('sha512sum "%s" > manifest.sha512' %(
		'" "'.join("%s.btrfs.zst.aes" %(s) for s in snap.subvolumes)
	))
	
	with open("check-integrity.sh", "w") as f:
		f.write("#! /bin/sh\n\nsha512sum --check manifest.sha512\n")
	
	## fix permissions and ownership of created objects
	os.chmod("check-integrity.sh", 0o555)
	
	outdir_stat = os.stat(outdir)
	os.chown(directory, outdir_stat.st_uid, outdir_stat.st_gid)
	
	for file in os.listdir(directory):
		path = os.path.join(directory, file)
		os.chown(path, outdir_stat.st_uid, outdir_stat.st_gid)
	
	return 0

# ------------------------------------------------------------------------------
# Verbs

def do_create(args):
	"""
	Bulk of the work happens here.
	
	- makes snapshots
	- calls serialization if needed
	- calls gc if needed
	"""
	
	snapshot = Snapshot(
		datetime.datetime.now().strftime("%Y-%m-%d.%H%M%S"),
		0
	)
	
	## determine the rank of the new snapshot
	##
	snaps_since_rank_1 = 1
	snaps_since_rank_2 = 1
	
	for snap in list_snapshots():
		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
	
	## promote the snapshot
	if snaps_since_rank_2 >= cfg.getint("snap", "rank_2_interval"):
		if "--process" in args:
			snapshot.rank = 2
		else:
			print("!!! Rank 2 snapshot is due, please enable --process")
			snapshot.rank = 1
	elif snaps_since_rank_1 >= cfg.getint("snap", "rank_1_interval"):
		snapshot.rank = 1
	
	## create the snapshot itself
	##
	os.makedirs(snapshot.path)
	ignore_prefix = cfg.get("subvol", "ignore_prefix")
	
	with chdir(cfg.get("subvol", "dir")):
		for subvolume in os.listdir():
			if subvolume.startswith(ignore_prefix):
				continue
			
			os.system(shlex.join((
				"btrfs", "subvolume", "snapshot", "-r",
				subvolume,
				os.path.join(snapshot.path, subvolume)
			)))
			
			snapshot.subvolumes.add(subvolume)
	
	## do optional processing
	##
	if snapshot.rank == 2:
		
		## snapshot serialization
		# if there's a previous Rank 2 snapshot, compute a diff against it
		# if not or if the process fails, demote this snap to Rank 1
		prev_r2_snaps = list_snapshots(2)
		
		if prev_r2_snaps:
			prev_r2_snap = prev_r2_snaps[-1]
			print(">>> doing a diff against snapshot %s" %(prev_r2_snap.name))
			
			error = serialize(
				snapshot,
				cfg.get("storage", "dir"),
				cfg.get("crypto", "keyfile"),
				prev_r2_snap
			)
			
			if error:
				print("!!! serialization failed, demoting snapshot to Rank 1")
				snapshot.rank = 1
		
		else:
			print("!!! no previous Rank 2 snapshot, please create one using btv stream")
			snapshot.rank = 1
		
		## garbage collection
		do_gc()
	
	## save all snapshot metadata
	snapshot.dump()
	
	print(">>> Snapshot created: rank %d %s" %(snapshot.rank, snapshot.name))

def do_list(args):
	"""
	Print a list of existing snapshots.
	"""
	
	print("Rank Name")
	
	for snap in list_snapshots():
		print("   %d %s %s" %(snap.rank, snap.name, ", ".join(snap.notes)))

def do_drop(args):
	"""
	Drop a snapshot in args[0].
	"""
	
	snap = get_snapshot_by_name(args[0])
	
	if not snap:
		print("!!! %s is not a snapshot" %(args[0]))
		sys.exit(2)
	
	error = snap.drop()
	
	if error:
		print("!!! failed to drop snapshot %s" %(snap.name))

def do_stream(args):
	"""
	Stream the snapshot args[0] into dir args[1].
	
	The affected snapshot is then promoted to Rank 2.
	"""
	
	snapshot = get_snapshot_by_name(args[0])
	
	if not snapshot:
		print("!!! %s is not a snapshot" %(args[0]))
		sys.exit(2)
	
	directory = args[1]
	
	if not os.path.isdir(directory):
		print("!!! %s is not a directory" %(directory))
		sys.exit(2)
	
	print(">>> Serializing %s into %s" %(snapshot.name, directory))
	
	error = serialize(
		snapshot,
		directory,
		cfg.get("crypto", "keyfile")
	)
	
	if error:
		print("!!! serialization failed")
	else:
		snapshot.rank = 2
		snapshot.notes.add("fully streamed")
		snapshot.dump()
		
		print("<<< %s serialized and promoted to Rank %d" %(snapshot.name, snapshot.rank))

def do_gc(args=None):
	"""
	Drops old snapshots.
	"""
	
	counts = [0, 0, 0] # Rank 0, 1, 2
	limits = [
		cfg.getint("gc", "rank_0_count"),
		cfg.getint("gc", "rank_1_count"),
		cfg.getint("gc", "rank_2_count")
	]
	
	for snap in reversed(list_snapshots()):
		counts[snap.rank] += 1
		
		if counts[snap.rank] > limits[snap.rank]:
			print(" >> delete Rank %d %s" %(snap.rank, snap.name))
			
			snap.drop()

verb_router = {
	"snapshot": do_create,
	"list": do_list,
	"drop": do_drop,
	"gc": do_gc,
	"stream": do_stream
}

# ------------------------------------------------------------------------------
# Entry point

def print_help():
	print("""Usage:
  {cmd} VERB [args]

Verbs:
  snapshot [--process]
    Create a snapshot. If --process is passed, do all optional work.
  
  stream SNAPSHOT OUTPUT_DIR
    Streams the SNAPSHOT into OUTPUT_DIR as a full (milestone) bin.
  
  list
    List snapshots and their rank.
  
  drop SNAPSHOT
    Drops (removes) SNAPSHOT. Use the full name as provided by {cmd} list.
  
  gc
    Drops old local snapshots based on garbage collector settings.
""".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()
	
	try:
		verb_router[verb](sys.argv[2:])
	except UsageError:
		print()
		print_help()
		sys.exit(1)