#!/opt/cloudlinux/venv/bin/python3
# coding:utf-8

# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2019 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENSE.TXT
#

import os
import re
import sys
import subprocess
import json


DBCTL_BIN = '/usr/share/lve/dbgovernor/utils/dbctl_orig'

# Mirror of is_safe_panel_name()/SAFE_NAME_RE in governor_package_limitting.py:
# a safe account name is [A-Za-z0-9._-] and must not start with '-' (so it can
# never be parsed as a dbctl_orig option by getopt_long). F-02 option-injection
# guard for the 'dbctl set' entrypoint, which forwards the name to dbctl_orig.
_SAFE_NAME_RE = re.compile(r'^[A-Za-z0-9._][A-Za-z0-9._-]*$')


def _add_flags_column(orig_out, flags_dict):
    dbctl_lines = orig_out.split('\n')[:-1]
    max_col_len = 19
    for i, row in enumerate(dbctl_lines):
        row_data = row.split()
        if i == 0:
            row_data.append('marks')
            row_data[0] = ' ' + row_data[0]
        else:
            user = row_data[0]
            if user in flags_dict:
                row_data.append(flags_dict[user]['cpu']+flags_dict[user]['io'])
            else:
                # suppose that user doesn't have individual limits
                # if its name not present in governor_package_limitting output
                row_data.append('--')
            if len(row_data[2]) > max_col_len:
                max_col_len = len(row_data[2])
        dbctl_lines[i] = '\t'.join(row_data)
    return '\n'.join(dbctl_lines).expandtabs(max_col_len+1)


def _process_call_error(call_err):
    print(call_err.stdout)
    if call_err.stderr is not None:
        print(call_err.stderr)
    print(f'Error: command \'{" ".join(call_err.cmd)}\' exit code = {call_err.returncode}')
    return call_err.returncode


def list_marked():
    """
    Adds column with marks denoting individual/package limits for db cpu and io usage. Example of output:
     user               cpu(%)              read(MB/s)          write(MB/s)         marks
    default             400/380/350/300     953/791/724/562     953/791/724/562     --
    limited             40/380/350/300      953/791/724/562     953/791/724/5       ++

    first '+' sign means at least one of the cpu limits for this user is different from its package limits
    second '+' sign means at least one of the read or write limits for this user is different from its package limits
    """
    try:
        get_flags_cmd = ['/usr/share/lve/dbgovernor/governor_package_limitting.py', 'get_individual']
        flags_json = subprocess.check_output(get_flags_cmd, text=True)
        flags_json = json.loads(flags_json)
        dbctl_out = subprocess.check_output([DBCTL_BIN, 'list'] + sys.argv[2:], text=True)
    except subprocess.CalledProcessError as call_err:
        return _process_call_error(call_err)
    except json.JSONDecodeError as json_err:
        print(f'Error: command \'{" ".join(get_flags_cmd)}\' returned not valid json:\n{json_err.doc}')
        return 1
    usernames_flags = {'default': {'cpu': '-', 'io': '-'}}
    for username in flags_json.keys():
        usernames_flags[username] = {}
        usernames_flags[username]['cpu'] = '+' if any(flags_json[username]['cpu']) else '-'
        usernames_flags[username]['io'] = '+' if any(flags_json[username]['read'] + flags_json[username]['write']) else '-'
    dbctl_out_processed = _add_flags_column(dbctl_out, usernames_flags)
    print(dbctl_out_processed)
    return 0


def _align_flags_with_dbctl_arg(values, flags):
    for i, lim in enumerate(values.split(',')):
        if lim != '0':
            flags[i] = True
        else:
            flags[i] = False


def _bool_to_str_list(input_list):
    return list(map(lambda v: str(v).lower(), input_list))


def _set_individual_cmd(username, cpu_flags, read_flags, write_flags):
    """
    Build a 'governor_package_limitting.py set_individual' command for the given
    user from the supplied cpu/read/write flag lists. Used both to persist the new
    flags and to roll them back to a prior snapshot on partial failure (F-21).
    """
    return ['/usr/share/lve/dbgovernor/governor_package_limitting.py', 'set_individual',
            f'--user={username}', f'--cpu={",".join(_bool_to_str_list(cpu_flags))}',
            f'--read={",".join(_bool_to_str_list(read_flags))}',
            f'--write={",".join(_bool_to_str_list(write_flags))}']


def set_with_packages():
    """
    1. Calls 'governor_package_limitting.py set_individual' to update governor_package_limit.json with
    info about what limits are individual (i.e. not from package) for given user
    2. Then calls 'dbctl_orig set' to actually change limits for the user
    """
    # delegate handling of incorrect arguments and default limits to dbctl_orig
    if len(sys.argv) < 3 or sys.argv[2] == 'default':
        try:
            subprocess.check_call([DBCTL_BIN] + sys.argv[1:])
            subprocess.check_call(['/usr/share/lve/dbgovernor/governor_package_limitting.py', 'sync'])
        except subprocess.CalledProcessError as call_err:
            return _process_call_error(call_err)
        return 0
    else:
        username = sys.argv[2]
        # F-02 option-injection guard on the primary 'dbctl set' entrypoint: this
        # username (sys.argv[2]) is forwarded to dbctl_orig via sys.argv[1:] - where
        # getopt_long would parse a '-'-leading value as an option - and into
        # set_individual. Refuse any name that isn't [A-Za-z0-9._-] / starts with
        # '-' before it reaches either sink, mirroring is_safe_panel_name().
        if not _SAFE_NAME_RE.match(username):
            print(f"Error: refusing unsafe account name {username!r} "
                  "(allowed: letters, digits, '.', '_', '-'; must not start with '-')")
            return 1
        try:
            # Read the whole individual_limits section (no --user) so we can tell
            # whether this user already had a persisted record. get_individual
            # --user synthesises an all-False vector for an unknown user, which
            # would hide that distinction and prevent a clean rollback.
            get_flags_cmd = ['/usr/share/lve/dbgovernor/governor_package_limitting.py', 'get_individual']
            all_flags_json = subprocess.check_output(get_flags_cmd, text=True)
            all_flags_json = json.loads(all_flags_json)
        except subprocess.CalledProcessError as call_err:
            return _process_call_error(call_err)
        except json.JSONDecodeError as json_err:
            print(f'Error: command \'{" ".join(get_flags_cmd)}\' returned not valid json:\n{json_err.doc}')
            return 1

        # Did the user have a real individual record before this call? This decides
        # how to roll the metadata back if the live update fails below.
        had_prior_record = username in all_flags_json
        # expecting that get_individual returns json for any existing user, even if
        # governor_package_limit.json doesn't have records about it
        prior_flags = all_flags_json.get(username, {'cpu': [False] * 4,
                                                    'read': [False] * 4,
                                                    'write': [False] * 4})
        cpu_flags = prior_flags['cpu']
        read_flags = prior_flags['read']
        write_flags = prior_flags['write']
        # Snapshot the prior individual-vs-package flags before mutating them, so we
        # can roll the persistent metadata back if applying the live limit fails
        # (F-21: the two privileged updates must not leave governor_package_limit.json
        # and dbctl_orig state diverged on partial failure).
        prior_individual_cmd = _set_individual_cmd(username, list(cpu_flags), list(read_flags), list(write_flags))
        for arg in sys.argv[3:]:
            if arg.startswith('--cpu'):
                _align_flags_with_dbctl_arg(arg[6:], cpu_flags)
            if arg.startswith('--read'):
                _align_flags_with_dbctl_arg(arg[7:], read_flags)
            if arg.startswith('--write'):
                _align_flags_with_dbctl_arg(arg[8:], write_flags)
        set_individual_cmd = _set_individual_cmd(username, cpu_flags, read_flags, write_flags)
        # F-21: persist the per-user metadata FIRST, then apply the live limit. The
        # metadata write is pure-config and cheaply reversible, whereas the live
        # dbctl_orig change is not. If metadata persistence fails we abort BEFORE
        # touching the live limits, so nothing is applied and there is nothing to
        # undo. If the live update then fails, we roll the metadata back to its
        # captured prior value (or remove a record we just created), so the two
        # stores can never end up diverged on a partial failure - in either failure
        # direction.
        try:
            subprocess.check_call(set_individual_cmd)
        except subprocess.CalledProcessError as call_err:
            # Metadata never persisted and no live change was made yet: just report.
            return _process_call_error(call_err)
        try:
            subprocess.check_call([DBCTL_BIN] + sys.argv[1:])
        except subprocess.CalledProcessError as call_err:
            # The metadata persisted but applying the live limit failed; undo the
            # metadata so the two stores do not diverge. The failed live update
            # never changed dbctl_orig state, so this rollback MUST stay
            # metadata-only and touch governor_package_limit.json alone - it must
            # not run any live dbctl_orig/dbctl command. If the user already had a
            # record, restore its prior snapshot via set_individual (JSON-only).
            # Otherwise remove the record we just created via reset_individual
            # --metadata-only (drops the JSON key without a live op), restoring the
            # true pre-attempt state - writing an all-False vector instead would
            # leave a phantom 'individual_limits' key for a user who had none. If
            # the rollback itself fails, surface that loudly rather than silently
            # swallowing it.
            if had_prior_record:
                rollback_cmd = prior_individual_cmd
            else:
                rollback_cmd = ['/usr/share/lve/dbgovernor/governor_package_limitting.py',
                                'reset_individual', f'--user={username}',
                                '--limits=all', '--metadata-only']
            try:
                subprocess.check_call(rollback_cmd)
            except subprocess.CalledProcessError:
                print('Error: applied individual-limit metadata but could not apply the '
                      'live limit, and rolling the metadata back failed; '
                      'governor_package_limit.json may be out of sync with dbctl for '
                      f'user {username!r} - re-run \'dbctl set\' to reconcile.')
            return _process_call_error(call_err)
        return 0


if __name__ == '__main__':
    if len(sys.argv) == 1:
        print('usage: dbctl command [parameter] [options]')
        sys.exit(1)
    if not os.path.exists(DBCTL_BIN):
        print(f'Error: {DBCTL_BIN}: No such file or directory\nPackage governor-mysql seems to be corrupted')
        sys.exit(1)

    exit_code = 0
    if sys.argv[1] == '--help':
        dbctl_help = subprocess.check_output([DBCTL_BIN, '--help'], text=True).split('\n')
        dbctl_help.insert(5, 'list-marked              list users, their limits and custom limit marks (cpu/io, + if at least one of the limits in group differ from package limits)')
        print('\n'.join(dbctl_help))
    elif sys.argv[1] == 'set':
        exit_code = set_with_packages()
    elif sys.argv[1] == 'list-marked':
        exit_code = list_marked()
    else:
        process = subprocess.Popen([DBCTL_BIN] + sys.argv[1:])
        exit_code = process.wait()
    sys.exit(exit_code)
