#!/opt/cloudlinux/venv/bin/python3
"""
lvestats-gov-helper — Bridge to clcommon.lib.MySQLGovernor for DB Governor.

Imports: clcommon.lib (MySQLGovernor, MySQLGovException, MySQLGovernorAbsent)

Returns JSON to stdout.

Usage:
    lvestats-gov-helper limits <username> [<username>...]
    lvestats-gov-helper governor-mode
    lvestats-gov-helper user-ignore <username>
    lvestats-gov-helper chart-info <username>
"""

import json
import os
import pwd
import sys


def _reseller_scope(caller):
    """Extra usernames a non-root *reseller* caller is entitled to: the users
    under the reseller it owns.

    ``limits`` is invoked non-root by ``cloudlinux_top --for-reseller <self>``
    (the CLI refuses ``--for-reseller`` for anyone but the caller), which
    legitimately needs its subaccounts' governor limits — their UIDs differ
    from the caller's, so the own-UID check alone would wrongly blank them.
    Mirrors the cpapi-helper reseller own-scope model.

    ``caller`` is the login name resolved from the caller's *euid* (never from
    argv), and the reseller-users lookup is gated on the caller actually being a
    reseller, so an unprivileged user can only ever widen scope to its own
    reseller's users — not an arbitrary tenant's. Fail-closed: any error, a
    panel without reseller support, or a non-reseller caller yields an empty
    set (the caller then keeps only its own row — no cross-tenant exposure).
    """
    try:
        from clcommon import cpapi
        resellers_fn = getattr(cpapi, "resellers", None)
        reseller_users_fn = getattr(cpapi, "reseller_users", None)
        if resellers_fn is None or reseller_users_fn is None:
            return set()
        if caller not in set(resellers_fn()):
            return set()
        return set(reseller_users_fn(caller))
    except Exception:
        return set()


def _authorized_usernames(usernames):
    """Restrict a non-root caller to the usernames it is entitled to.

    lvestats-gov-helper is installed world-executable, and its ``limits``
    subcommand is invoked by the unprivileged tenant CLIs (``cloudlinux_top``,
    ``lveinfo``). Root callers (the CMCollector daemon) need the full list. For
    a non-root caller keep only the usernames it is entitled to, so an
    unprivileged local user cannot read other tenants' MySQL Governor limits by
    invoking the helper directly (finding CLOS-4630 #30):

    * its **own** account — usernames whose UID matches the caller's euid; and
    * if the caller is a **reseller**, the users under it — the legitimate
      ``cloudlinux_top --for-reseller`` path, whose subaccounts have different
      UIDs (see :func:`_reseller_scope`).

    Everything else is dropped. Root is unrestricted.
    """
    if os.geteuid() == 0:
        return list(usernames)
    euid = os.geteuid()
    try:
        caller = pwd.getpwuid(euid).pw_name
    except KeyError:
        caller = None
    scope = _reseller_scope(caller) if caller is not None else set()
    authorized = []
    for username in usernames:
        try:
            own = pwd.getpwnam(username).pw_uid == euid
        except KeyError:
            own = False
        if own or username in scope:
            authorized.append(username)
    return authorized


def cmd_limits(usernames):
    usernames = _authorized_usernames(usernames)
    results = {}
    if not usernames:
        return results
    try:
        from clcommon.lib import MySQLGovernor, MySQLGovException, MySQLGovernorAbsent
        governor = MySQLGovernor()
    except ImportError as e:
        for u in usernames:
            results[u] = {"error": "clcommon.lib.MySQLGovernor not available: %s" % e, "absent": True}
        return results

    for username in usernames:
        try:
            status = governor.get_governor_status_by_username(username)
            if status != 'watched':
                results[username] = {"cpu": 0, "io": 0}
                continue
            cpu, io = governor.get_limits_by_user(username)
            results[username] = {"cpu": cpu, "io": io}
        except MySQLGovernorAbsent as e:
            results[username] = {"error": str(e), "absent": True}
        except MySQLGovException as e:
            results[username] = {"error": str(e), "absent": False}
        except Exception as e:
            results[username] = {"error": "%s: %s" % (type(e).__name__, e), "absent": False}

    return results


def cmd_governor_mode():
    """Get MySQL Governor mode using clcommon.lib.MySQLGovernor.

    Returns the mode string: "abusers", "all", "single", "on", or "none".
    """
    try:
        from clcommon.lib import MySQLGovernor
        g = MySQLGovernor()
        mode = g.get_governor_mode()
        return mode or "none"
    except Exception:
        return "none"


def cmd_user_ignore(username):
    """Whether the DB Governor is configured to ignore this user.

    Mirrors Python lve-stats `func.get_governor_ignore_for_user`: a user with
    governor status "ignored" is excluded from DB accounting/charting.
    Returns a bool (False when governor is absent, on any error, or when a
    non-root caller is not entitled to this username, so the caller falls back
    to the default "not ignored" behaviour).
    """
    # Gate the per-user governor bit like ``limits`` does, so an unprivileged
    # caller cannot read another tenant's ignored status (CLOS-4630 #30).
    if not _authorized_usernames([username]):
        return False
    try:
        from clcommon.lib import MySQLGovernor
        g = MySQLGovernor()
        return g.get_governor_status_by_username(username) == 'ignored'
    except Exception:
        return False


def cmd_chart_info(username):
    """Combined lookup for the lvechart DB-overlay: governor mode + whether
    this user is ignored. Returned together so the chart needs only one
    helper invocation instead of two.
    """
    return {"mode": cmd_governor_mode(), "ignored": cmd_user_ignore(username)}


def main():
    if len(sys.argv) < 2:
        print("Usage: lvestats-gov-helper <command> [args...]", file=sys.stderr)
        print("Commands: limits, governor-mode, user-ignore, chart-info", file=sys.stderr)
        sys.exit(1)

    command = sys.argv[1]

    try:
        if command == "limits":
            if len(sys.argv) < 3:
                print("Usage: lvestats-gov-helper limits <username> [<username>...]", file=sys.stderr)
                sys.exit(1)
            result = cmd_limits(sys.argv[2:])

        elif command == "governor-mode":
            result = cmd_governor_mode()

        elif command == "user-ignore":
            if len(sys.argv) < 3:
                print("Usage: lvestats-gov-helper user-ignore <username>", file=sys.stderr)
                sys.exit(1)
            result = cmd_user_ignore(sys.argv[2])

        elif command == "chart-info":
            if len(sys.argv) < 3:
                print("Usage: lvestats-gov-helper chart-info <username>", file=sys.stderr)
                sys.exit(1)
            result = cmd_chart_info(sys.argv[2])

        else:
            print("Unknown command: %s" % command, file=sys.stderr)
            sys.exit(1)

        json.dump(result, sys.stdout)
        sys.stdout.write("\n")

    except Exception as e:
        print("lvestats-gov-helper: unexpected error: %s: %s" % (type(e).__name__, e), file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()