2 files changed, 212 insertions, 0 deletions

diff --git a/hosts/vidhar/borg/borgsnap/borgsnap/__main__.py b/hosts/vidhar/borg/borgsnap/borgsnap/__main__.py
new file mode 100644
index 00000000..e93e6a60
--- /dev/null
+++ b/hosts/vidhar/borg/borgsnap/borgsnap/__main__.py
@@ -0,0 +1,202 @@
+import argparse
+import os, sys, signal
+from pyprctl import cap_permitted, cap_inheritable, cap_effective, cap_ambient, Cap
+from pwd import getpwnam
+
+from datetime import datetime, timezone
+from dateutil.parser import isoparse
+
+from xdg import xdg_runtime_dir
+import unshare
+from tempfile import TemporaryDirectory
+
+import logging
+
+import json
+import subprocess
+
+import pathlib
+from pathlib import Path
+
+from atomicwrites import atomic_write
+
+from traceback import format_exc
+
+
+borg_pwd = getpwnam('borg')
+
+def as_borg(caps=set(), cwd=None):
+    if caps:
+        cap_permitted.add(*caps)
+        cap_inheritable.add(*caps)
+        cap_effective.add(*caps)
+        cap_ambient.add(*caps)
+
+    os.setgid(borg_pwd.pw_gid)
+    os.setuid(borg_pwd.pw_uid)
+
+    if cwd is not None:
+        os.chdir(cwd)
+
+
+def _archive_name(snapshot, target, archive_prefix):
+    _, _, ts = snapshot.rpartition('@')
+    creation_time = isoparse(ts).astimezone(timezone.utc)
+    archive_name = _archive_basename(snapshot, archive_prefix)
+    return f'{target}::{archive_name}-{creation_time.strftime("%Y-%m-%dT%H:%M:%S")}'
+
+def _archive_basename(snapshot, archive_prefix):
+    base_name, _, _ = snapshot.rpartition('@')
+    return archive_prefix + base_name.replace('-', '--').replace('/', '-')
+
+def check(*, snapshot, target, archive_prefix, cache_file):
+    archives = None
+    if cache_file:
+        logger.debug('Trying cache...')
+        try:
+            with open(cache_file, mode='r', encoding='utf-8') as fp:
+                archives = set(json.load(fp))
+            logger.info('Loaded archive list from cache')
+        except FileNotFoundError:
+            pass
+
+    if not archives:
+        logger.info('Loading archive list from remote...')
+        with subprocess.Popen(['borg', 'list', '--info', '--lock-wait=600', '--json', target], stdout=subprocess.PIPE, preexec_fn=lambda: as_borg()) as proc:
+            archives = set([archive['barchive'] for archive in json.load(proc.stdout)['archives']])
+        if cache_file:
+            logger.debug('Saving archive list to cache...')
+            with atomic_write(cache_file, mode='w', encoding='utf-8', overwrite=True) as fp:
+                json.dump(list(archives), fp)
+
+    # logger.debug(f'archives: {archives}')
+    _, _, archive_name = _archive_name(snapshot, target, archive_prefix).partition('::')
+    if archive_name in archives:
+        logger.info(f'{archive_name} found')
+        return 0
+    else:
+        logger.info(f'{archive_name} not found')
+        return 126
+
+def create(*, snapshot, target, archive_prefix, dry_run):
+    basename = _archive_basename(snapshot, archive_prefix)
+
+    with TemporaryDirectory(prefix=f'borg-mount_{basename}_', dir=os.environ.get('RUNTIME_DIRECTORY')) as tmpdir:
+        child = os.fork()
+        if child == 0:
+            unshare.unshare(unshare.CLONE_NEWNS)
+            subprocess.run(['mount', '--make-rprivate', '/'], check=True)
+            chroot = pathlib.Path(tmpdir) / 'chroot'
+            upper = pathlib.Path(tmpdir) / 'upper'
+            work = pathlib.Path(tmpdir) / 'work'
+            for path in [chroot,upper,work]:
+                path.mkdir()
+            subprocess.run(['mount', '-t', 'overlay', 'overlay', '-o', f'lowerdir=/,upperdir={upper},workdir={work}', chroot], check=True)
+            bindMounts = ['nix', 'run', 'run/secrets.d', 'run/wrappers', 'proc', 'dev', 'sys', pathlib.Path(os.path.expanduser('~')).relative_to('/')]
+            if os.environ.get('BORG_BASE_DIR'):
+                bindMounts.append(pathlib.Path(os.environ['BORG_BASE_DIR']).relative_to('/'))
+            if 'SSH_AUTH_SOCK' in os.environ:
+                bindMounts.append(pathlib.Path(os.environ['SSH_AUTH_SOCK']).parent.relative_to('/'))
+            for bindMount in bindMounts:
+                (chroot / bindMount).mkdir(parents=True,exist_ok=True)
+                # print(*['mount', '--bind', pathlib.Path('/') / bindMount, chroot / bindMount], file=stderr)
+                subprocess.run(['mount', '--bind', pathlib.Path('/') / bindMount, chroot / bindMount], check=True)
+            os.chroot(chroot)
+            os.chdir('/')
+            dir = pathlib.Path('/borg')
+            dir.mkdir(parents=True,exist_ok=True,mode=0o0750)
+            os.chown(dir, borg_pwd.pw_uid, borg_pwd.pw_gid)
+            try:
+                subprocess.run(['mount', '-t', 'zfs', '-o', 'ro', snapshot, dir], check=True)
+                env = os.environ.copy()
+                create_args = ['borg',
+                               'create',
+                               '--lock-wait=600',
+                               '--one-file-system',
+                               '--compression=auto,zstd,10',
+                               '--chunker-params=10,23,16,4095',
+                               '--files-cache=ctime,size',
+                               '--show-rc',
+                               # '--remote-ratelimit=20480',
+                               '--progress',
+                               '--list',
+                               '--filter=AMEi-x?',
+                               '--stats' if not dry_run else '--dry-run'
+                               ]
+                _, _, ts = snapshot.rpartition('@')
+                creation_time = isoparse(ts).astimezone(timezone.utc)
+                create_args += [f'--timestamp={creation_time.strftime("%Y-%m-%dT%H:%M:%S")}']
+                env['BORG_FILES_CACHE_SUFFIX'] = basename
+                create_args += [_archive_name(snapshot, target, archive_prefix), '.']
+                print({'create_args': create_args, 'cwd': dir, 'env': env}, file=sys.stderr)
+                subprocess.run(create_args, stdin=subprocess.DEVNULL, env=env, preexec_fn=lambda: as_borg(caps={CAP.DAC_READ_SEARCH}, cwd=dir), check=True)
+                #  subprocess.run(create_args, stdin=subprocess.DEVNULL, env=env, preexec_fn=lambda: None, cwd=dir, check=True)
+            finally:
+                subprocess.run(['umount', dir], check=True)
+            os._exit(0)
+        else:
+            while True:
+                waitpid, waitret = os.wait()
+                if waitret != 0:
+                    sys.exit(waitret)
+                if waitpid == child:
+                    break
+    return 0
+
+def sigterm(signum, frame):
+    raise SystemExit(128 + signum)
+
+def main():
+    signal.signal(signal.SIGTERM, sigterm)
+
+    global logger
+    logger = logging.getLogger(__name__)
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter( logging.Formatter('[%(levelname)s](%(name)s): %(message)s') )
+    if sys.stderr.isatty():
+        console_handler.setFormatter( logging.Formatter('%(asctime)s [%(levelname)s](%(name)s): %(message)s') )
+    logger.addHandler(console_handler)
+
+    # log uncaught exceptions
+    def log_exceptions(type, value, tb):
+        global logger
+
+        logger.error(value)
+        sys.__excepthook__(type, value, tb) # calls default excepthook
+
+    sys.excepthook = log_exceptions
+
+    parser = argparse.ArgumentParser(prog='borgsnap')
+    parser.add_argument('--verbose', '-v', action='count', default=0)
+    parser.add_argument('--target', metavar='REPO', default='yggdrasil.borgbase:repo')
+    parser.add_argument('--archive-prefix', metavar='REPO', default='yggdrasil.vidhar.')
+    subparsers = parser.add_subparsers()
+    subparsers.required = True
+    parser.set_defaults(cmd=None)
+    check_parser = subparsers.add_parser('check')
+    check_parser.add_argument('--cache-file', type=lambda p: Path(p).absolute(), default=None)
+    check_parser.add_argument('snapshot')
+    check_parser.set_defaults(cmd=check)
+    create_parser = subparsers.add_parser('create')
+    create_parser.add_argument('--dry-run', '-n', action='store_true', default=False)
+    create_parser.add_argument('snapshot')
+    create_parser.set_defaults(cmd=create)
+    args = parser.parse_args()
+
+    if args.verbose <= 0:
+        logger.setLevel(logging.WARNING)
+    elif args.verbose <= 1:
+        logger.setLevel(logging.INFO)
+    else:
+        logger.setLevel(logging.DEBUG)
+
+    cmdArgs = {}
+    for copy in {'target', 'archive_prefix', 'snapshot', 'cache_file', 'dry_run'}:
+        if copy in vars(args):
+            cmdArgs[copy] = vars(args)[copy]
+
+    return args.cmd(**cmdArgs)
+
+
+if __name__ == '__main__':
+    sys.exit(main())
diff --git a/hosts/vidhar/borg/borgsnap/setup.py b/hosts/vidhar/borg/borgsnap/setup.py
new file mode 100644
index 00000000..76356bfc
--- /dev/null
+++ b/hosts/vidhar/borg/borgsnap/setup.py
@@ -0,0 +1,10 @@
+from setuptools import setup
+
+setup(name='borgsnap',
+      packages=['borgsnap'],
+      entry_points={
+          'console_scripts': [
+              'borgsnap=borgsnap.__main__:main',
+          ],
+      }
+)