bmount

#!/usr/bin/python3
"""
Mirror subtrees between different parts of the filesystem using bind-mounts.
"""
from __future__ import print_function

import argparse
import logging
import os
import os.path
import subprocess
import sys
import time
import traceback
import unittest

MOUNTINFO = "/proc/self/mountinfo"
if not os.path.exists(MOUNTINFO):
	raise ImportError("need Linux >= 2.6.26") # see `man proc`

DEFAULT_BINDFILE = ".bmrc"

SEPLEN = len(os.path.sep)
ROOT = os.getcwd()
while ROOT != os.path.dirname(ROOT):
	ROOT = os.path.dirname(ROOT)

def p_key(k):
	"""Comparator that forces path children to be sorted next to its parent."""
	return k.split(os.path.sep)

def p_is_dir(p):
	return p.endswith(os.path.sep)

def p_to_dir(p):
	return os.path.join(p, '')

def p_to_rel(p):
	return os.path.relpath(p, ROOT) if os.path.isabs(p) else p

def p_to_abs(p):
	return os.path.normpath(os.path.join(ROOT, p))

def p_get_device(p):
	assert not os.path.islink(p) # must never call this with a symlink
	if not os.path.exists(p): return None, None
	# run df twice because it doesn't escape spaces in paths
	# WONTFIX(infinity0): it also turns \n into ? in the output, but this is
	# too uncommon to bother accounting for. (it handles \n in the input fine.)
	output = subprocess.check_output(["df", "--output=source", p])
	source = output.decode("ascii").split('\n')
	output = subprocess.check_output(["df", "--output=target", p])
	target = output.decode("ascii").split('\n')
	return source[-2], target[-2]

def f_type(p):
	return ('_' if not os.path.lexists(p)
	  else 'l' if os.path.islink(p)
	  else 'd' if os.path.isdir(p) else 'f')

def f_make(p, t):
	assert not os.path.islink(p) # must never call this with a symlink
	par = os.path.dirname(p)
	# TODO(infinity0): try sudo if we get permission error
	if not os.path.exists(par):
		os.makedirs(par)
	if t == 'f':
		logging.debug("create file %s", p)
		open(p, 'a').close()
	else:
		logging.debug("create dir %s", p)
		os.mkdir(p)

def d_get_bindfile(dd, f=DEFAULT_BINDFILE):
	d = os.path.abspath(dd)
	while not os.path.isfile(os.path.join(d, f)):
		if d == ROOT: return os.path.normpath(os.path.join(dd, f))
		d = os.path.dirname(d)
	return os.path.join(d, f)

def read_significant_line(fp):
	while 1:
		line = fp.readline()
		if not line: return None
		line = line.rstrip('\n')
		if not line or line.startswith('#'): continue
		return line

def read_significant_lines(fp):
	while 1:
		line = read_significant_line(fp)
		if not line: break
		yield line

def run_unstrict(cmd, *args, **kwargs):
	try:
		logging.info("running %r" % cmd)
		subprocess.check_output(cmd, *args, **kwargs)
	except subprocess.CalledProcessError as e:
		logging.warn("command failed: %s" % e.output)

def encode_path(path):
	return path.encode("unicode_escape").decode("ascii")

def decode_path(path):
	return path.encode("ascii").decode("unicode_escape")


class Points(object):
	"""
	:var list pts: List of path point entries. Each entry is an absolute
	     path, where a final slash denotes a directory, and a lack of
	     one denotes a file. Whenever you see a variable called "pp",
	     you may assume interpret it like this.
	"""
	@staticmethod
	def checkConsistent(a, b):
		if p_is_dir(a):
			if b.startswith(a):
				raise ValueError("overlap: %s inside %s" % (b, a))
		else:
			if b.startswith(p_to_dir(a)):
				raise ValueError("both a file and dir: %s, %s" % (a, b))

	def __init__(self, unsafe_pts=[]):
		pts = sorted(set(unsafe_pts), key=p_key)
		for i in range(len(pts)-1):
			Points.checkConsistent(pts[i], pts[i+1])
		self.pts = pts

	def __repr__(self):
		return "Points(%r)" % self.pts

	def insert(self, pp):
		"""
		:param pp: path point entry
		"""
		self.pts.append(pp)
		self.pts.sort(key=p_key)
		idx = self.pts.index(pp)
		logging.debug("inserted %s @ %s in %s", pp, idx, self.pts)
		try:
			if idx > 0:
				Points.checkConsistent(self.pts[idx-1], self.pts[idx])
			if idx+1 < len(self.pts):
				Points.checkConsistent(self.pts[idx], self.pts[idx+1])
		except Exception:
			self.pts.pop(idx)
			raise
		return self

	def remove(self, pp):
		"""
		:param pp: path point entry
		"""
		self.pts.remove(pp)
		return self

	def get_anc(self, p):
		"""
		:param p: normalised path, no final sep
		"""
		parents = [pp for pp in self.pts if p_is_dir(pp) and p.startswith(pp) and len(p) > len(pp)]
		assert len(parents) <= 1 # due to disjoint restriction
		return parents[0] if parents else None

	def get_self(self, p):
		"""
		:param p: normalised path, no final sep
		"""
		return (p if p in self.pts else
		  p_to_dir(p) if p_to_dir(p) in self.pts else
		  None)

	def get_desc(self, p):
		"""
		:param p: normalised path, no final sep
		"""
		return [c for c in self.pts if c.startswith(p_to_dir(p))]


class Bind(object):

	def __init__(self, source, target):
		self.source = p_to_dir(os.path.abspath(source))
		self.target = p_to_dir(os.path.abspath(target))
		if self.source == self.target:
			raise ValueError("source and target cannot be same: %s" % self.source)
		self.points = Points()
		self._state = None
		self._state_atime = 0

	def _readMounts(self):
		logging.info("reading mounts from %s", MOUNTINFO)
		state = []
		#traceback.print_stack()
		with open(MOUNTINFO, 'rb') as fp:
			for line in fp.readlines():
				part1, part2 = line.split(b" - ") # see `man proc` for spec
				(_, _, _, mount_rel, fs_file, _) = part1.split(b' ', 5)
				fs_file = fs_file.decode("unicode_escape")
				mount_rel = mount_rel.decode("unicode_escape")
				(_, fs_spec, _) = part2.split(b' ', 2)
				fs_spec = fs_spec.decode("unicode_escape")
				# WONTFIX(infinity0): technically we should ignore mount points that are buried
				# under other mount points. But I don't see a significant use-case for that to
				# warrant the complexity of handling it.

				if fs_file.startswith(self.target) or p_to_dir(fs_file) == self.target:
					pp = p_to_abs(os.path.relpath(fs_file, self.target))
					pp = p_to_dir(pp) if os.path.isdir(fs_file) else pp
					if pp in state:
						logging.warn("multiple mounts on: %s", pp)
					if pp not in self.points.pts:
						src, dst = self.getEndpoints(pp)
						src_dev, base_pt = p_get_device(src)
						# Below is (I believe) the best possible heuristic given that Linux does not
						# store explicitly whether a mountpoint was done using --bind. It is
						# possible that it is even fully-correct but I don't think it's worth the
						# effort for me to think it through in that much detail. In particular, I
						# have not considered the cases where src/dst traverses symlinks.
						src_mount_rel = os.path.relpath(src, base_pt) if base_pt else None
						if fs_spec != src_dev or mount_rel != p_to_abs(src_mount_rel):
							logging.debug("ignore (seemingly) non-bmount mount point: %s" % encode_path(fs_file))
							continue
						logging.warn("extraneous mount on: %s", pp)
					state.append(pp)
		return state

	@property
	def state(self):
		atime = time.time()
		if self._state is None or atime > self._state_atime + 0.25: # update after a second
			try:
				self._state = Points(self._readMounts())
				self._state_atime = atime
			except Exception as e:
				traceback.print_exc()
				logging.warn("unable to parse current bind state: %s", e)
				logging.warn("I will continue, but avoid touching the system")
				logging.warn("You may also run umountForce to try to fix the issue")
				self._state = Points(["/invalid_state_detected"])
				# i.e. assert INVALID but avoid infinite recursion
				assert set(self._state.pts) - set(self.points.pts)
		return self._state

	@property
	def status(self):
		if self.points.pts == self.state.pts:
			return 'FULL'
		elif not self.state.pts:
			return 'NONE'
		elif set(self.state.pts) - set(self.points.pts):
			return 'INVALID'
		else:
			return 'PARTIAL'

	def report(self, msg='report', debug=False):
		l = logging.debug if debug else logging.info
		l("%s: %s: %s vs %s", msg, self.status, self.points, self.state)

	def getEndpoints(self, p):
		src = os.path.normpath(os.path.join(self.source, p_to_rel(p)))
		dst = os.path.normpath(os.path.join(self.target, p_to_rel(p)))
		return src, dst

	def _mountTry(self, pp):
		src, dst = self.getEndpoints(pp)
		run_unstrict(["sudo", "mount", "-B", src, dst])
		self._state = None # invalidate cache

	def _umountTry(self, pp):
		src, dst = self.getEndpoints(pp)
		run_unstrict(["sudo", "umount", dst])
		self._state = None # invalidate cache

	def mountAll(self):
		status = self.status
		if status == 'FULL':
			logging.info("mountAll nothing to do")
			return
		elif status == 'INVALID':
			raise ValueError("cannot run in state INVALID")

		for pp in sorted(set(self.points.pts) - set(self.state.pts), key=p_key):
			self._ensureEndpoints(pp, True)
			self._mountTry(pp)

		if self.status == 'FULL':
			self.report("mountAll complete")
		else:
			self.report("mountAll incomplete")

	def umountAll(self):
		status = self.status
		if status == 'NONE':
			logging.info("umountAll nothing to do")
			return
		elif status == 'INVALID':
			raise ValueError("cannot run in state INVALID")

		for pp in self.state.pts:
			self._umountTry(pp)

		if self.status == 'NONE':
			self.report("umountAll complete")
		else:
			self.report("umountAll incomplete")

	def umountForce(self):
		if self.status != 'INVALID':
			raise ValueError("can only run in state INVALID")

		mounts = self._readMounts()
		oldlen = len(mounts)
		while 1:
			for pp in mounts.__reversed__():
				self._umountTry(pp)
			mounts = self._readMounts()
			curlen = len(mounts)
			if curlen == 0 or curlen >= oldlen: # finished, or making no progress
				break
			oldlen = curlen
		if curlen == 0:
			self.report("umountForce complete")
			assert self.status == 'NONE'
		else:
			self.report("umountForce incomplete; requires manual intervention")

	def _insertBind(self, pp, dryrun=False):
		"""
		:param p: normalised path, no final sep
		"""
		logging.debug("insert %s", pp)
		status = self.status if not dryrun else None # don't call if not needed
		self.points.insert(pp)
		if status == 'FULL':
			logging.info("FULL, will try to insert the bind mount")
			self._mountTry(pp)
		return pp

	def _removeBind(self, pp, dryrun=False):
		"""
		:param p: normalised path, no final sep
		"""
		logging.debug("remove %s", pp)
		status = self.status if not dryrun else None # don't call if not needed
		self.points.remove(pp)
		if status == 'FULL':
			logging.info("FULL, will try to remove the bind mount")
			self._umountTry(pp)
		return pp

	def _ensureEndpoints(self, p, pPointEntry, dryrun=False):
		"""
		Ensure that the endpoints exist for bind-mounting.

		If either endpoint is a file/dir, the other must be a file/dir.
		If one does not exist, it is created of the same type as the
		other. If neither exist, both are created as directories/files .
		"""
		src, dst = self.getEndpoints(p)

		src_t = f_type(src)
		dst_t = f_type(dst)
		p = os.path.join(ROOT, p) # force absolute
		def pdIfNotSubsume(t):
			if t == 'f': return p
			if self.source.startswith(p_to_dir(dst)):
				raise ValueError("target point would subsume source tree: %s" % dst)
			return p_to_dir(p)
		# TODO(infinity0): think about what if src/dst traverses symlinks.
		# This may interfere with bmount mount point detection, in which case
		# we would have to disallow it.

		if src_t == dst_t == '_':
			f_t = 'd' if not pPointEntry or p_is_dir(p) else 'f'
			if not dryrun: f_make(src, f_t)
			if not dryrun: f_make(dst, f_t)
			return pdIfNotSubsume(f_t)
		elif src_t == 'l':
			raise ValueError("symlinks not supported due to mount(8) limitation: %s" % src)
		elif dst_t == 'l':
			raise ValueError("symlinks not supported due to mount(8) limitation: %s" % dst)
		elif src_t == '_' and dst_t != '_':
			if not dryrun: f_make(src, dst_t)
			return pdIfNotSubsume(dst_t)
		elif src_t != '_' and dst_t == '_':
			if not dryrun: f_make(dst, src_t)
			return pdIfNotSubsume(src_t)
		elif src_t == dst_t:
			return pdIfNotSubsume(src_t)
		else:
			raise ValueError("conflicting types: %s, %s" % (src, dst))

	def insert(self, p, dryrun=False, noreport=False):
		"""
		:param p: path to insert
		:return ins, rem: the bind points that were actually inserted and/or removed
		"""
		p = p_to_abs(p)
		logging.debug("resolved to %s", p)

		panc = self.points.get_anc(p)
		if panc:
			logging.info("ignore insert: ancestor already inserted: %s", panc)
			return [], []

		pp = self.points.get_self(p)
		if pp:
			logging.info("ignore insert: self already inserted: %s", pp)
			return [], []

		desc = self.points.get_desc(p)
		# umount children first

		rem = [self._removeBind(c, dryrun) for c in desc]
		ins = [self._insertBind(self._ensureEndpoints(p, False, dryrun), dryrun)]
		if not noreport: self.report("insert finished", dryrun)
		return ins, rem

	def remove(self, p, dryrun=False, noreport=False):
		"""
		:param p: path to remove
		:return ins, rem: the bind points that were actually inserted and/or removed
		"""
		p = p_to_abs(p)
		logging.debug("resolved to %s", p)

		panc = self.points.get_anc(p)
		if panc:
			raise ValueError("cannot remove: must remove ancestor instead: %s" % panc)

		pp = self.points.get_self(p)
		if pp:
			return [], [self._removeBind(pp, dryrun)]

		desc = self.points.get_desc(p)
		if not desc:
			logging.info("ignore remove: self already removed and no children: %s", pp)

		rem = [self._removeBind(c, dryrun) for c in desc]
		ins = []
		if not noreport: self.report("remove finished", dryrun)
		return ins, rem


class BindFile(object):

	def __init__(self, fn=DEFAULT_BINDFILE):
		self.fn = fn
		self.args = None
		self.bind = None

	def parseArgs(self, *args):
		if not args:
			raise ValueError("arg syntax: (src, dst), or ('import', src='/') or ('export', dst='/')")
		elif args[0] == "import":
			src, dst = args[1] if len(args) > 1 and args[1] else ROOT, os.curdir
		elif args[0] == "export":
			src, dst = os.curdir, args[1] if len(args) > 1 and args[1] else ROOT
		elif args[0] and args[1]:
			src, dst = args[0], args[1]
		else:
			raise ValueError("arg syntax: (src, dst), or ('import', src='/') or ('export', dst='/')")

		return src, dst

	def init(self, *args):
		self.parseArgs(*args) # ensure valid
		self.args = args
		with open(self.fn, 'w') as fp:
			print("\t".join(map(encode_path, self.args)), file=fp)
		logging.debug("file wrote: %s", self.fn)
		self.load()

	def load(self):
		with open(self.fn, 'r') as fp:
			args = read_significant_line(fp).rstrip("\n").split("\t")
			args = list(map(decode_path, args))
			points = read_significant_lines(fp)
			points = list(map(lambda x: decode_path(x.rstrip('\n')), points))

		src, dst = self.parseArgs(*args)

		# expand ~ and ~username
		src = os.path.expanduser(src)
		dst = os.path.expanduser(dst)

		# paths relative to file rather than cwd
		if not os.path.isabs(src):
			src = os.path.join(os.path.dirname(self.fn), src)
		if not os.path.isabs(dst):
			dst = os.path.join(os.path.dirname(self.fn), dst)

		bind = Bind(src, dst)
		logging.debug("file read: %s: %s to %s", self.fn, bind.source, bind.target)
		try:
			for point in points:
				bind.insert(point, dryrun=True, noreport=True)
		except:
			logging.warn("error loading file: %s", self.fn)
			raise

		self.args = args
		self.bind = bind
		_ = bind.state

	def save(self):
		if self.bind is None:
			raise ValueError("not loaded")

		with open(self.fn, 'w') as fp:
			print("\t".join(map(encode_path, self.args)), file=fp)
			for pp in self.bind.points.pts:
				print(encode_path(pp), file=fp)
		logging.debug("file wrote: %s", self.fn)

	def insert(self, *pps, dryrun=False):
		for pp in pps:
			self.bind.insert(pp, dryrun=dryrun)
		if not dryrun:
			self.save()

	def remove(self, *pps, dryrun=False):
		for pp in pps:
			self.bind.remove(pp, dryrun=dryrun)
		if not dryrun:
			self.save()


def main(*args):
	parser = argparse.ArgumentParser()
	parser.add_argument("-f", "--file", help="bind file to load and save, default \"%(default)s\".",
	  dest='bindfile', default=DEFAULT_BINDFILE)
	group = parser.add_mutually_exclusive_group()
	group.add_argument("-u", "--up", help="search this and ancestors for the bindfile.",
	  metavar='DIR', dest='binddir_up')
	group.add_argument("-d", "--down", help="search this and descendants for the bindfile.",
	  metavar='DIR', dest='binddir_down')
	parser.add_argument("-v", "--verbose", help="increase output verbosity",
	  action="count", default=0)
	subps = parser.add_subparsers(help="sub-command to run",
	  metavar="SUBCMD")

	subp = subps.add_parser("init", help="init or reset a bind file", aliases=["n"])
	subp.add_argument("-y", "--yes", help="force-reset bind fb bile even if it already exists",
	  action='store_true', default=False)
	subp.add_argument("args", help="""<source> <target>, or export <target> or import <source>. """
	  "in the latter cases, source/target is implicitly %s, and the second argument "
	  "may be omitted to implicitly be %s." % (os.curdir, ROOT),
	  nargs='+')
	subp.set_defaults(subcmd="init")

	subp = subps.add_parser("insert", help="insert a point", aliases=["i"])
	subp.add_argument("-n", "--dry-run", help="print what would be done, don't actually do",
	  action="store_true", default=False)
	subp.add_argument("points", nargs='+', help="point to insert")
	subp.set_defaults(subcmd="insert")
	subp = subps.add_parser("remove", help="remove a point", aliases=["r"])
	subp.add_argument("-n", "--dry-run", help="print what would be done, don't actually do",
	  action="store_true", default=False)
	subp.add_argument("points", nargs='+', help="point to remove")
	subp.set_defaults(subcmd="remove")

	subp = subps.add_parser("status", help="get the current status", aliases=["s"],
	  description="output meanings: NONE: no points are active; FULL: all points are active; "
	  "PARTIAL: some points are active; INVALID: extraneous points are active.")
	subp.set_defaults(subcmd="status")
	subp = subps.add_parser("list", help="list all points", aliases=["l"])
	subp.set_defaults(subcmd="list")
	subp = subps.add_parser("mountAll", help="mount all points", aliases=["m"])
	subp.set_defaults(subcmd="mountAll")
	subp = subps.add_parser("umountAll", help="umount all points", aliases=["u"])
	subp.set_defaults(subcmd="umountAll")
	subp = subps.add_parser("umountForce", help="in INVALID state, umount all existing points")
	subp.set_defaults(subcmd="umountForce")

	opts = parser.parse_args(args)
	if hasattr(opts, "dry_run") and opts.dry_run and opts.verbose < 2:
		opts.verbose = 2
	if opts.verbose >= 2:
		logging.getLogger().setLevel(logging.DEBUG)
	elif opts.verbose >= 1:
		logging.getLogger().setLevel(logging.INFO)

	fn = os.path.normpath(opts.bindfile)
	if os.path.dirname(fn):
		if opts.binddir_down or opts.binddir_up:
			logging.debug("ignore -u/-d since -f is a path: %s" % fn)
	elif opts.binddir_up:
		fn = d_get_bindfile(opts.binddir_up, fn)
	elif opts.binddir_down:
		cands = [dp for (dp, dns, fns) in os.walk(opts.binddir_down) if fn in fns]
		if not cands:
			raise ValueError("no %s found in %s" % (fn, opts.binddir_down))
		elif len(cands) == 1:
			fn = os.path.normpath(os.path.join(cands[0], fn))
		else:
			raise ValueError("multiple %s found in %s: %s" % (fn, opts.binddir_down, cands))

	b = BindFile(fn)

	if not hasattr(opts, "subcmd"):
		parser.error("must specify a subcommand")

	elif opts.subcmd == "init":
		if os.path.exists(b.fn):
			b.load()
			if b.args == opts.args:
				pass
			elif opts.yes:
				b.init(*opts.args)
			else:
				raise ValueError("bindfile already exists with different args: %s" % b.args)
		else:
			b.init(*opts.args)
	elif not os.path.exists(fn):
		raise ValueError("bindfile does not exist; run 'init' first: %s" % fn)

	if opts.subcmd in ["insert", "remove"]:
		b.load()
		getattr(b, opts.subcmd)(*opts.points, dryrun=opts.dry_run)

	if opts.subcmd == "status":
		b.load()
		if opts.verbose >= 1:
			print("file: %s" % b.fn)
			print("args: %s" % ','.join(b.args))
			print("status: %s" % b.bind.status)
			print("wanted: %r" % b.bind.points.pts)
			print("active: %r" % b.bind.state.pts)
		else:
			print(b.bind.status)

	if opts.subcmd == "list":
		b.load()
		for pp in b.bind.points.pts:
			print(encode_path(pp))

	if opts.subcmd in ["mountAll", "umountAll", "umountForce"]:
		b.load()
		getattr(b.bind, opts.subcmd)()

	return 0


class Test(unittest.TestCase):

	def test_p_to_rel(self):
		self.assertEquals(p_to_rel(ROOT), os.curdir)
		self.assertEquals(p_to_rel(os.curdir), os.curdir)
		self.assertTrue(os.path.join('a', p_to_rel(os.getcwd())).startswith('a'))

	def test_p_key(self):
		self.assertEquals(sorted(['x', 'xa', 'x+', 'x/']), ['x', 'x+', 'x/', 'xa'])
		self.assertEquals(sorted(['x', 'xa', 'x+', 'x/'], key=p_key), ['x', 'x/', 'x+', 'xa'])

	def test_Points_dup(self):
		self.assertRaises(ValueError, Points, ['/a', '/a/'])

	def test_Points_sub(self):
		self.assertRaises(ValueError, Points, ['a/', 'a/b'])
		self.assertRaises(ValueError, Points, ['a', 'a/b'])
		self.assertRaises(ValueError, Points, ['a', 'a+b', 'a/b'])

	def test_Points_get_anc(self):
		p = Points(['/a/', '/b/c'])
		self.assertEquals(p.get_anc('/a'), None)
		self.assertEquals(p.get_anc('/'), None)
		self.assertEquals(p.get_anc('/a/c'), '/a/')

	def test_Points_get_anc_root(self):
		p = Points(['/'])
		self.assertEquals(p.get_anc('/a'), '/')
		self.assertEquals(p.get_anc('/'), None)

	def test_Points_get_self(self):
		p = Points(['/ab/'])
		self.assertEquals(p.get_self('/ab'), '/ab/')
		p = Points(['/ab'])
		self.assertEquals(p.get_self('/ab'), '/ab')

	def test_Points_get_self_root(self):
		p = Points(['/'])
		self.assertEquals(p.get_self('/a'), None)
		self.assertEquals(p.get_self('/'), '/')

	def test_Points_get_desc(self):
		p = Points(['/a/b', '/a/c'])
		self.assertEquals(p.get_desc('/a'), ['/a/b', '/a/c'])
		self.assertEquals(p.get_desc('/'), ['/a/b', '/a/c'])


if __name__ == '__main__':
	sys.exit(main(*sys.argv[1:]))