#!/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 import ProductType
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'
FORCE_COMMENT = (
    'Force disable rootfs protection on enterprise-licensed hardware. This is '
    'a developer option that will result in an unsupported configuration.'
)


def set_readwrite(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=off on dataset {entry["ds"]}')
    run([ZFS_CMD, 'set', 'readonly=off', 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 OFF 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)


def process_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--force',
        action=argparse.BooleanOptionalAction,
        help=FORCE_COMMENT,
    )
    return parser.parse_args()


if __name__ == '__main__':
    datasets = []
    args = process_args()

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

    # We currently use this script inside some jenkins pipelines that are not
    # normal truenas install and so we need to skip the fs check on force.
    if not args.force:
        usr_fs_check()

    rv = run([ZFS_CMD, 'get', '-o', 'value', '-H', 'truenas:developer', '/'], capture_output=True)

    # If we're already in developer-mode, skip license check
    # This is to allow workflow for developer working on HA platform to
    # run this script then run install-dev-tools to get full development
    # environment.
    if rv.stdout.decode().strip() != 'on' and not args.force:
        with Client() as c:
            if c.call('system.product_type') == ProductType.ENTERPRISE:
                print((
                    'Root filesystem protections may not be administratively disabled '
                    'on Enterprise-licensed TrueNAS products. Circumventing this '
                    'restriction is considered an unsupported configuration.'
                ))
                sys.exit(1)
    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('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=on', root])

    for entry in datasets:
        set_readwrite(entry)

    chmod_files()

    sys.exit(0)
