#!/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

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

MAX_DIE_NUM = 2
MAX_UVB_CNA_STR_LEN = 22
MAX_URMA_EID_LENGTH = 39
MAX_NODES = 32
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 write_proc_file(proc_dir, proc_name, proc_value):
    """
    Don't use 'shell=True' for subprocess.run/subprocess.Popen, it's not safe. However, if 'shell=true'
    is not set, it is difficult to modify the proc file for subprocess.run/subprocess.Popen.
    """
    exit_code = 0
    try:
        with open("/proc/%s/%s" % (proc_dir, proc_name), mode="w") as f:
            f.write(str(proc_value) + "\n")
    except PermissionError as e:
        exit_code = -e.errno
        print("sentryctl: error: set %s failed for %s, the user does not have the permission!" % (proc_dir, proc_name))
    except FileNotFoundError as e:
        exit_code = -e.errno
        print("sentryctl: error: set %s failed for %s, the proc file does not exist!" % (proc_dir, proc_name))
    except Exception as e:
        exit_code = getattr(e, 'errno', -1)
        exit_code = -exit_code if exit_code > 0 else -1
        print("sentryctl: error: set %s failed for %s" % (proc_dir, proc_name))
    finally:
        return exit_code

def set_sentry_reporter_proc(proc_name, proc_value):
    return write_proc_file("sentry_reporter", proc_name, proc_value)

def set_remote_reporter_proc(proc_name, proc_value):
    return write_proc_file("sentry_remote_reporter", proc_name, proc_value)

def set_urma_heartbeat(proc_value):
    return write_proc_file("sentry_urma_comm", "heartbeat", proc_value)

def set_uvb_proc(server_cna_str):
    if len(server_cna_str.strip()) == 0:
        print("Invalid args for server_cna")
        sys.exit(-1)
    server_cna_list = server_cna_str.strip().split(";")
    if len(server_cna_list) > MAX_NODES:
        print("Exceeded the maximum number (%d) of nodes supported." % MAX_NODES)
        sys.exit(-1)
    for cna in server_cna_list:
        if len(cna.strip()) == 0:
            print("Find invalid cna (%s) for server_cna" % cna)
            sys.exit(-1)
        if len(cna.strip()) > MAX_UVB_CNA_STR_LEN:
            print("Invalid cna (%s) for server_cna" % cna)
            sys.exit(-1)
    server_cna_string =";".join(server_cna_list)
    return write_proc_file("sentry_uvb_comm", "server_cna", server_cna_string)

def set_urma_proc(server_eid, client_jetty_id):
    if len(server_eid.strip()) == 0:
        print("Invalid args for server_eid, server_eid is empty string")
        sys.exit(-1)
    if client_jetty_id < MIN_CLIENT_JETTY_ID or client_jetty_id > MAX_CLIENT_JETTY_ID:
        print("Invalid args for client_jetty_id")
        sys.exit(-1)

    server_eid_list = server_eid.strip().split(";")
    if len(server_eid_list) > MAX_DIE_NUM:
        print("Invalid args for server_eid, server_eid contains an extra semicolon.")
        sys.exit(-1)
    for server_eid_i in server_eid_list:
        if len(server_eid_i) == 0:
            print("Invalid args for server_eid, server_eid is empty string")
            sys.exit(-1)
        server_eid_list_i = server_eid_i.split(",")
        for eid_i in server_eid_list_i:
            if len(eid_i.strip()) != MAX_URMA_EID_LENGTH:
                print("Invalid args for server_eid, the length of the eid ({}) does not equal {}.".format(eid_i, MAX_URMA_EID_LENGTH))
                sys.exit(-1)
    write_urma_info = " ".join((server_eid, str(client_jetty_id)))
    return write_proc_file("sentry_urma_comm", "client_info", write_urma_info)

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'])


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())
        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":
        print("res msg format error")
        return None

    try:
        res_data_len = int(res_data[3:])
        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"},
        ],
    }
    parser_set = subparsers.add_parser('set', help='set plugins params')
    parsers_set_plugin_param = parser_set.add_subparsers(dest="set_task")
    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':
        ret_code = 0
        if client_args.set_task == "sentry_remote_reporter":
            if client_args.cna != None:
                if client_args.cna < 0:
                    print("cna should be a number not less than 0.")
                    ret_code = -1
                else:
                    ret_code += set_remote_reporter_proc("cna", client_args.cna)
            if client_args.eid != None:
                num_of_semicolons = client_args.eid.count(";")
                if num_of_semicolons > 1:
                    print("invalid value for extraneous semicolon.")
                    sys.exit(-1)
                eid_list = client_args.eid.split(";")
                for eid_i in eid_list:
                    if len(eid_i) != MAX_URMA_EID_LENGTH:
                        print("The length of eid must be {}, but the detected input length is {}.".format(MAX_URMA_EID_LENGTH, len(eid_i)))
                        sys.exit(-1)
                ret_code += set_remote_reporter_proc("eid", client_args.eid)
            if client_args.uvb_comm != None or client_args.urma_comm != None:
                if client_args.urma_comm == "on":
                    ret_code += set_remote_reporter_proc("urma_comm", client_args.urma_comm)
                if client_args.uvb_comm == "on":
                    ret_code += set_remote_reporter_proc("uvb_comm", client_args.uvb_comm)
                if client_args.urma_comm == "off":
                    ret_code += set_remote_reporter_proc("urma_comm", client_args.urma_comm)
                if client_args.uvb_comm == "off":
                    ret_code += set_remote_reporter_proc("uvb_comm", client_args.uvb_comm)
            if client_args.panic != None:
                ret_code += set_remote_reporter_proc("panic", client_args.panic)
            if client_args.panic_timeout_ms != None:
                ret_code += set_remote_reporter_proc("panic_timeout", client_args.panic_timeout_ms)
            if client_args.kernel_reboot != None:
                ret_code += set_remote_reporter_proc("kernel_reboot", client_args.kernel_reboot)
            if client_args.kernel_reboot_timeout_ms != None:
                ret_code += set_remote_reporter_proc("kernel_reboot_timeout", client_args.kernel_reboot_timeout_ms)
        elif client_args.set_task == "sentry_urma_comm":
            if client_args.server_eid != None and client_args.client_jetty_id:
                ret_code += set_urma_proc(client_args.server_eid, client_args.client_jetty_id)
            elif client_args.server_eid != None or client_args.client_jetty_id:
                print("Options --server_eid and --client_jetty_id need to be used together")
                ret_code += -1
            if client_args.heartbeat != None:
                ret_code += set_urma_heartbeat(client_args.heartbeat)
        elif client_args.set_task == "sentry_uvb_comm":
            ret_code += set_uvb_proc(client_args.server_cna)
        elif client_args.set_task == "sentry_reporter":
            if client_args.ub_mem_fault_with_kill != None:
                ret_code += set_sentry_reporter_proc("ub_mem_fault_with_kill", client_args.ub_mem_fault_with_kill)
            if client_args.ub_mem_fault != None:
                ret_code += set_sentry_reporter_proc("ub_mem_fault", client_args.ub_mem_fault);
            if client_args.power_off != None:
                ret_code += set_sentry_reporter_proc("power_off", client_args.power_off);
            if client_args.oom != None:
                ret_code += set_sentry_reporter_proc("oom", client_args.oom);
        sys.exit(ret_code)
    else:
        parser.print_help()
        sys.exit(-1)

    if client_args.cmd_type != 'list' 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'])
