add log2mem functionality and refracture

pwnagotchi/fs/__init__.py

@@ -0,0 +1,176 @@
+import os
+import re
+import tempfile
+import contextlib
+import shutil
+import _thread
+import logging
+
+from time import sleep
+from distutils.dir_util import copy_tree
+
+mounts = list()
+
+@contextlib.contextmanager
+def ensure_write(filename, mode='w'):
+    path = os.path.dirname(filename)
+    fd, tmp = tempfile.mkstemp(dir=path)
+
+    with os.fdopen(fd, mode) as f:
+        yield f
+        f.flush()
+        os.fsync(f.fileno())
+
+    os.replace(tmp, filename)
+
+
+def size_of(path):
+    total = 0
+    for root, dirs, files in os.walk(path):
+        for f in files:
+            total += os.path.getsize(os.path.join(root, f))
+    return total
+
+
+def is_mountpoint(path):
+    return os.system(f"mountpoint -q {path}") == 0
+
+
+def setup_mounts(config):
+    global mounts
+    fs_cfg = config['fs']['memory']
+    if not fs_cfg['enabled']:
+        return
+
+    for name, options in fs_cfg['mounts'].items():
+        if not options['enabled']:
+            continue
+        logging.debug("[FS] Trying to setup mount %s (%s)", name, options['mount'])
+        size,unit = re.match(r"(\d+)([a-zA-Z]+)", options['size']).groups()
+        target = os.path.join('/run/pwnagotchi/disk/', os.path.basename(options['mount']))
+
+        is_mounted = is_mountpoint(target)
+        logging.debug("[FS] %s is %s mounted", options['mount'],
+                      "already" if is_mounted else "not yet")
+
+        m = MemoryFS(
+            options['mount'],
+            target,
+            size=options['size'],
+            zram=options['zram'],
+            zram_disk_size=f"{int(size)*2}{unit}",
+            rsync=options['rsync'])
+
+        if not is_mounted:
+            if not m.mount():
+                logging.debug(f"Error while mounting {m.mountpoint}")
+                continue
+
+            if not m.sync(to_ram=True):
+                logging.debug(f"Error while syncing to {m.mountpoint}")
+                m.umount()
+                continue
+
+        interval = int(options['sync'])
+        if interval:
+            logging.debug("[FS] Starting thread to sync %s (interval: %d)",
+                        options['mount'], interval)
+            _thread.start_new_thread(m.daemonize, (interval,))
+        else:
+            logging.debug("[FS] Not syncing %s, because interval is 0",
+            options['mount'])
+
+        mounts.append(m)
+
+class MemoryFS:
+    @staticmethod
+    def zram_install():
+        if not os.path.exists("/sys/class/zram-control"):
+            logging.debug("[FS] Installing zram")
+            return os.system("modprobe zram") == 0
+        return True
+
+    @staticmethod
+    def zram_dev():
+        logging.debug("[FS] Adding zram device")
+        return open("/sys/class/zram-control/hot_add", "rt").read().strip("\n")
+
+    def __init__(self, mount, disk, size="40M",
+                 zram=True, zram_alg="lz4", zram_disk_size="100M",
+                 zram_fs_type="ext4", rsync=True):
+        self.mountpoint = mount
+        self.disk = disk
+        self.size = size
+        self.zram = zram
+        self.zram_alg = zram_alg
+        self.zram_disk_size = zram_disk_size
+        self.zram_fs_type = zram_fs_type
+        self.zdev = None
+        self.rsync = True
+        self._setup()
+
+    def _setup(self):
+        if self.zram and MemoryFS.zram_install():
+            # setup zram
+            self.zdev = MemoryFS.zram_dev()
+            open(f"/sys/block/zram{self.zdev}/comp_algorithm", "wt").write(self.zram_alg)
+            open(f"/sys/block/zram{self.zdev}/disksize", "wt").write(self.zram_disk_size)
+            open(f"/sys/block/zram{self.zdev}/mem_limit", "wt").write(self.size)
+            logging.debug("[FS] Creating fs (type: %s)", self.zram_fs_type)
+            os.system(f"mke2fs -t {self.zram_fs_type} /dev/zram{self.zdev} >/dev/null 2>&1")
+
+        # ensure mountpoints exist
+        if not os.path.exists(self.disk):
+            logging.debug("[FS] Creating %s", self.disk)
+            os.makedirs(self.disk)
+
+        if not os.path.exists(self.mountpoint):
+            logging.debug("[FS] Creating %s", self.mountpoint)
+            os.makedirs(self.mountpoint)
+
+
+    def daemonize(self, interval=60):
+        logging.debug("[FS] Daemonized...")
+        while True:
+            self.sync()
+            sleep(interval)
+
+
+    def sync(self, to_ram=False):
+        source, dest = (self.disk, self.mountpoint) if to_ram else (self.mountpoint, self.disk)
+        needed, actually_free = size_of(source), shutil.disk_usage(dest)[2]
+        if actually_free >= needed:
+            logging.debug("[FS] Syning %s -> %s", source,dest)
+            if self.rsync:
+                os.system(f"rsync -aXv --inplace --no-whole-file --delete-after {source}/ {dest}/ >/dev/null 2>&1")
+            else:
+                copy_tree(source, dest, preserve_symlinks=True)
+            os.system("sync")
+            return True
+        return False
+
+
+    def mount(self):
+        if os.system(f"mount --bind {self.mountpoint} {self.disk}"):
+            return False
+
+        if os.system(f"mount --make-private {self.disk}"):
+            return False
+
+        if self.zram and self.zdev is not None:
+            if os.system(f"mount -t {self.zram_fs_type} -o nosuid,noexec,nodev,user=pwnagotchi /dev/zram{self.zdev} {self.mountpoint}/"):
+                return False
+        else:
+            if os.system(f"mount -t tmpfs -o nosuid,noexec,nodev,mode=0755,size={self.size} pwnagotchi {self.mountpoint}/"):
+                return False
+
+        return True
+
+
+    def umount(self):
+        if os.system(f"umount -l {self.mountpoint}"):
+            return False
+
+        if os.system(f"umount -l {self.disk}"):
+            return False
+        return True