#!/usr/bin/python3
# coding: utf-8
# Copyright (c) 2023 Huawei Technologies Co., Ltd.
# sysSentry is licensed under the Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#     http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
# See the Mulan PSL v2 for more details.

"""
sentryctl: control command for sysSentry.
"""
import socket
import argparse
import json
import sys
import logging
import json

from syssentry.utils import MAX_MSG_LEN

CTL_SOCKET_PATH = "/var/run/sysSentry/control.sock"
MAX_PARAM_LENGTH = 256


MAX_URMA_EID_LENGTH = 39
MIN_PANIC_TIMEOUT_MS = 0
MAX_PANIC_TIMEOUT_MS = 3600000
MIN_KERNEL_REBOOT_TIMEOUT_MS = 0
MAX_KERNEL_REBOOT_TIMEOUT_MS = 3600000
MIN_CLIENT_JETTY_ID = 3
MAX_CLIENT_JETTY_ID = 1023

RESULT_MSG_DATA_LEN = 4
CTL_MSG_LEN_LEN = 3
ALARM_MSG_DATA_LEN = 6
DEFAULT_ALARM_TIME_RANGE = 10

def status_output_format(res_data):
    """format output"""
    print(f"status: {res_data}")

def result_output_format(res_data):
    try:
        print(json.dumps(res_data, indent=4))
    except json.decoder.JSONDecodeError:
        logging.warning("result_output_format: result is \n%s\n, but json.dumps failed!")
        print(res_data)

def mod_list_output_format(res_data):
    """format output list"""
    mod_list = json.loads(res_data)
    print("oneshot:")
    for mod in mod_list['ONESHOT']:
        print(f"\t {mod}")
    print("period:")
    for mod in mod_list['PERIOD']:
        print(f"\t {mod}")
    print("")


def res_output_handle(res_struct, req_type):
    """handle output"""
    if req_type == 'mod_list':
        mod_list_output_format(res_struct['data'])
    elif req_type == 'get_status':
        status_output_format(res_struct['data'])
    elif req_type == 'get_result':
         result_output_format(res_struct['data'])
    elif req_type == 'get_alarm':
         result_output_format(res_struct['data'])
    elif res_struct['ret'] == "failed":
        print(res_struct['data'])
        if req_type == 'set':
            sys.exit(-1)


def res_msg_serial(res_msg):
    """serial res_msg"""
    res_struct = json.loads(res_msg)
    return res_struct


def client_send_and_recv(request_data, data_str_len):
    """client socket send and recv message"""
    try:
        client_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    except socket.error:
        print("sentryctl: client creat socket error")
        return None

    # connect to syssentry
    try:
        client_socket.connect(CTL_SOCKET_PATH)
    except OSError:
        client_socket.close()
        print("sentryctl: client connect error")
        return None

    # msg: CTL{len}{data}
    req_data_len = len(request_data)
    request_msg = "CTL" + str(req_data_len).zfill(3) + request_data

    try:
        client_socket.send(request_msg.encode())
        if len("CTL") + data_str_len > MAX_MSG_LEN:
            client_socket.close()
            logging.error("socket recv data is illegal:%d", len("CTL") + data_str_len)
            return None
        res_data = client_socket.recv(len("CTL") + data_str_len)
        res_data = res_data.decode()
    except (OSError, UnicodeError):
        client_socket.close()
        print("sentryctl: client communicate error")
        return None

    # res: RES{len}{data}
    res_magic = res_data[:3]
    if res_magic != "RES":
        client_socket.close()
        print("res msg format error")
        return None

    try:
        res_data_len = int(res_data[3:])
        if res_data_len > MAX_MSG_LEN:
            client_socket.close()
            logging.error("socket recv data is illegal:%d", res_data_len)
            return
        res_msg_data = client_socket.recv(res_data_len)
        res_msg_data = res_msg_data.decode()
        return res_msg_data
    except (OSError, ValueError, UnicodeError):
        print("sentryctl: client recv res msg error")
    finally:
        client_socket.close()

    return None


if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='sentryctl', description='System inspect tools',
                                     allow_abbrev=False)
    subparsers = parser.add_subparsers(help="sentryctl: control cmdline for sysSentry",
                                       dest='cmd_type')
    parser_start = subparsers.add_parser('start', help='start specified task')
    parser_start.add_argument('task_name')
    parser_stop = subparsers.add_parser('stop', help='stop specified task')
    parser_stop.add_argument('task_name')
    parser_reload = subparsers.add_parser('reload', help='reload specified mod')
    parser_reload.add_argument('task_name')
    parser_status = subparsers.add_parser('status', help='get task status')
    parser_status.add_argument('task_name')
    parser_get_result = subparsers.add_parser('get_result', help='get task result')
    parser_get_result.add_argument('task_name')
    parser_get_alarm = subparsers.add_parser('get_alarm', help='get task alarm')
    parser_get_alarm.add_argument('task_name')
    parser_get_alarm.add_argument('-s', '--time_range', type=int, default=DEFAULT_ALARM_TIME_RANGE, help='Specified time range')
    parser_get_alarm.add_argument('-d', '--detailed', action='store_true', help='Print Detailed Information')
    parser_list = subparsers.add_parser('list', help='show all loaded task mod')

    set_cmd_args_plugins_info = {
        "sentry_remote_reporter": [
            {"name" : "cna", "type": int, "choices": None, "required" : False, "help": "set cna info, it is uint32_t integer"},
            {"name" : "eid", "type": str, "choices": None, "required" : False, "help": "set local eid info"},
            {"name" : "uvb_comm", "type": str, "choices":  ["on", "off"], "required" : False, "help": "Select the UVB communication mode"},
            {"name" : "urma_comm", "type": str, "choices":  ["on", "off"], "required" : False, "help": "Select the URMA communication mode"},
            {"name" : "panic", "type": str, "choices": ["on", "off"], "required" : False, "help": "Panic event control switch"},
            {"name" : "kernel_reboot", "type": str, "choices": ["on", "off"], "required" : False, "help": "Kernel reboot event control switch"},
            {"name" : "panic_timeout_ms", "type": int, "choices": None, "required" : False, "help": "set panic timeout, value range is [%d, %d]" %(MIN_PANIC_TIMEOUT_MS, MAX_PANIC_TIMEOUT_MS)},
            {"name" : "kernel_reboot_timeout_ms", "type": int, "choices": None, "required" : False, "help": "set kernel reboot timeout, value range is [%d, %d]" %(MIN_KERNEL_REBOOT_TIMEOUT_MS, MAX_KERNEL_REBOOT_TIMEOUT_MS)},
        ],
        "sentry_urma_comm" : [
            {"name" : "server_eid", "type": str, "choices": None, "required" : False, "help": "Info about other nodes to be connected to the current node"},
            {"name" : "client_jetty_id", "type": int, "choices": None, "required" : False, "help": "jetty id of the current node, value range is [%d, %d]" %(MIN_CLIENT_JETTY_ID, MAX_CLIENT_JETTY_ID)},
            {"name" : "heartbeat", "type": str, "choices":  ["on", "off"], "required" : False, "help": "Heartbeat detection function switch, default off"},
        ],
        "sentry_uvb_comm" : [
            {"name": "server_cna", "type": str, "choices": None, "required" : True, "help": "server cna array"},
        ],
        "sentry_reporter" : [
            {"name" : "ub_mem_fault_with_kill", "type" : str, "choices":  ["on", "off"], "required" : False, "help" : "Enable/Disable sending SIGBUS signal with UB mem event"},
            {"name" : "ub_mem_fault", "type" : str, "choices": ["on", "off"], "required" : False, "help": "Enable/Disable UB mem event"},
            {"name" : "power_off", "type": str, "choices": ["on", "off"], "required" : False, "help": "Enable/Disable power off event"},
            {"name" : "oom", "type" : str, "choices": ["on", "off"], "required" : False, "help": "Enable/Disable oom event"},
            {"name" : "oom_policy", "type": str, "choices": None, "required" : False, "help": "Set OOM rate limit policy in JSON format, e.g. '{\"RR_KSWAPD\": 5, \"RR_DIRECT_RECLAIM\": 3, \"RR_HUGEPAGE_RECLAIM\": 2}'"},
        ],
    }
    parser_set = subparsers.add_parser('set', help='set plugins params')
    parsers_set_plugin_param = parser_set.add_subparsers(dest="set_ko_name")
    for plugin_name, args_info in set_cmd_args_plugins_info.items():
        parser_set_plugin_param = parsers_set_plugin_param.add_parser(plugin_name, help="set task args for %s" % plugin_name)
        for arg in args_info:
            parser_set_plugin_param.add_argument("--" + arg["name"], type = arg["type"], choices = arg["choices"], required = arg["required"], help = arg["help"])

    client_args = parser.parse_args()
    if client_args.cmd_type == 'list':
        req_msg_struct = {"type": "mod_list", "data":""}
    elif client_args.cmd_type == 'start':
        req_msg_struct = {"type": "start", "data": client_args.task_name}
    elif client_args.cmd_type == 'stop':
        req_msg_struct = {"type": "stop", "data": client_args.task_name}
    elif client_args.cmd_type == 'status':
        req_msg_struct = {"type": "get_status", "data": client_args.task_name}
    elif client_args.cmd_type == 'get_result':
        req_msg_struct = {"type": "get_result", "data": client_args.task_name}
    elif client_args.cmd_type == 'get_alarm':
        if not isinstance(client_args.time_range, int) or client_args.time_range <= 0:
            print(f"time_range is not a positive integer: {client_args.time_range}")
        req_msg_struct = {
            "type": "get_alarm", 
            "data": {
                'task_name': client_args.task_name, 
                'time_range': client_args.time_range,
                'detailed': client_args.detailed,
            }
        }
    elif client_args.cmd_type == 'reload':
        req_msg_struct = {"type": "reload", "data": client_args.task_name}
    elif client_args.cmd_type == 'set':
        params = vars(client_args)
        req_msg_struct = {
            "type": "set",
            "data": params
        }
    else:
        parser.print_help()
        sys.exit(-1)

    if client_args.cmd_type != 'list' and client_args.cmd_type != 'set' and len(client_args.task_name) > MAX_PARAM_LENGTH:
        print(f"sentryctl: task name is longer than {MAX_PARAM_LENGTH}")
        sys.exit(-1)

    request_message = json.dumps(req_msg_struct)
    if client_args.cmd_type == 'get_result':
        result_message = client_send_and_recv(request_message, RESULT_MSG_DATA_LEN)
    elif client_args.cmd_type == 'get_alarm':
        result_message = client_send_and_recv(request_message, ALARM_MSG_DATA_LEN)
    else:
        result_message = client_send_and_recv(request_message, CTL_MSG_LEN_LEN)
    if not result_message:
        print("sentryctl: client_send_and_recv failed")
        sys.exit(-1)

    result_struct = res_msg_serial(result_message)

    if result_struct['ret'] != "success":
        print(result_struct['ret'])
    res_output_handle(result_struct, req_msg_struct['type'])
