enable-rootfs-protection for TrueNAS Scale

enable-rootfs-protection.py

#!/usr/bin/python3

import argparse
import json
import os
import stat
import sys
from truenas_api_client import Client
from pathlib import Path
from subprocess import run

from middlewared.utils.mount import getmntinfo
from middlewared.utils.filesystem.stat_x import statx


ZFS_CMD = '/usr/sbin/zfs'
TO_CHMOD = ['apt', 'dpkg']
EXECUTE_BITS = stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
PKG_MGMT_DISABLED_PATH = '/usr/local/bin/pkg_mgmt_disabled'


def set_readonly(entry):
    if 'RO' not in entry['fhs_entry']['options']:
        return

    # There shouldn't be a legitimate reason to edit files in /conf
    if entry['fhs_entry']['name'] == 'conf':
        return

    print(f'Setting readonly=on on dataset {entry["ds"]}')
    run([ZFS_CMD, 'set', 'readonly=on', entry['ds']])

def usr_fs_check():
    mntid = statx('/usr').stx_mnt_id
    mntinfo = getmntinfo(mnt_id=mntid)[mntid]
    match mntinfo['fs_type']:
        case 'zfs':
            return

        case 'overlay':
            if mntinfo['mount_source'] == 'sysext':
                print((
                    '/usr is currently provided by a readonly systemd system extension. '
                    'This may occur if nvidia module support is enabled. System extensions '
                    'must be disabled prior to disabling rootfs protection.'
                ))
            else:
                print(f'/usr is currently provided by an unexpected overlayfs filesystem: {mntinfo}.')
        case _:
            print((
                f'{mntinfo["fs_type"]}: /usr is currently provided by an unexpected filesystem type. '
                'Unable to disable rootfs protection.'
            ))

    sys.exit(1)


def chmod_files():
    with os.scandir('/usr/bin') as it:
        for entry in it:
            do_chmod = False
            if not entry.is_file():
                continue

            for prefix in TO_CHMOD:
                if not entry.name.startswith(prefix):
                    continue

                if (stat.S_IMODE(entry.stat().st_mode) & EXECUTE_BITS) == EXECUTE_BITS:
                    do_chmod = True
                    break

            if do_chmod:
                new_mode = stat.S_IMODE(entry.stat().st_mode & ~EXECUTE_BITS)
                print(f'{entry.path}: setting {oct(new_mode)} on file.')
                os.chmod(entry.path, new_mode)

    # Also turn ON execute bits for pkg_mgmt_disabled
    p = Path(PKG_MGMT_DISABLED_PATH)
    if p.exists():
        old_mode = p.stat().st_mode
        if old_mode & ~EXECUTE_BITS:
            new_mode = stat.S_IMODE(old_mode | EXECUTE_BITS)
            print(f'{PKG_MGMT_DISABLED_PATH}: setting {oct(new_mode)} on file.')
            p.chmod(new_mode)


if __name__ == '__main__':
    datasets = []

    if os.getuid() != 0:
        print((
            'Enabling filesystem protections must be done as the root user '
            'or with sudo.'
        ))
        sys.exit(1)

    usr_fs_check()

    try:
        # The following file is created during TrueNAS installation
        # and contains dataset configuration and guid details
        with open('/conf/truenas_root_ds.json', 'r') as f:
            datasets = json.load(f)
    except FileNotFoundError:
        pass

    print('Un-flagging root dataset as developer mode')
    rv = run([ZFS_CMD, 'get', '-o', 'name', '-H', 'name', '/'], capture_output=True)
    root = rv.stdout.decode().strip()
    run([ZFS_CMD, 'set', 'truenas:developer=off', root])

    chmod_files()

    for entry in datasets:
        set_readonly(entry)

    sys.exit(0)

Code adapted from disable-rootfs-protection, all credit to TrueNAS developers.

Added the .py extension for code highlighting; like with disable, it is not needed.