#!/usr/bin/python3
#-------------------------------------------------------------------------------
# # Copyright (C) 20xx-2021 NVIDIA Corporation. ALL RIGHTS RESERVED. #
#-------------------------------------------------------------------------------
# determine-reset-reason
#
# Program designed to run once, soon after system boot which will
# determine the cause of the previous reset and store it to the disk,
#

try:
    import datetime
    import pytz
    import json
    import os
    import pwd 
    import re
    import sys
    import logging
    import logging.handlers
    import io
    from os.path import isfile, join
    from collections import OrderedDict

except ImportError as err:
    raise ImportError("%s - required module not found" % str(err))

VERSION = "1.0"

SYSLOG_IDENTIFIER = "determine-reset-reason"

RESET_REASON_DIR = "/var/reset-reason/"
RESET_REASON_HISTORY_DIR = "/var/reset-reason/history/"
CSMGRD_RESET_REASON_FILE = "/var/cache/csmsgr/csmgrd.state"
PREVIOUS_RESET_REASON_FILE = os.path.join(RESET_REASON_DIR,  "previous-reset-reason.json")
FIRST_BOOT_PLATFORM_FILE = "/tmp/notify_firstboot_to_platform"

REBOOT_TYPE_CSMGRD_PATTERN_WARM = ".*reboot_warm.*"
REBOOT_TYPE_CSMGRD_PATTERN_FAST = ".*reboot_fast.*"

HWMGMT_SYSTEM_ROOT = '/var/run/hw-management/system/'
#reboot cause related definitions
RESET_REASON_ROOT = HWMGMT_SYSTEM_ROOT
RESET_REASON_FILE_LENGTH = 1

# Possible reboot causes
RESET_REASON_POWER_LOSS = "Power Loss"
RESET_REASON_AC_POWER_LOSS = "AC Power Loss"
RESET_REASON_PWR_CONVERTER_FAIL = "Power converter failure"
RESET_REASON_RESET_FROM_ASIC = "Reset from ASIC"
RESET_REASON_RESET_FROM_SOC = "Reset from SOC"
RESET_REASON_RESET_FROM_FW = "Reset from FW"
RESET_REASON_SW_PWR_OFF = "SW asserted power off signal through CPLD"
RESET_REASON_RESET_SYSTEM = "System reset"
RESET_REASON_RESET_PLATFORM = "SW asserted reset through CPLD"
RESET_REASON_SWB_DC_DC_PWR_FAIL = "Switch board power loss"
RESET_REASON_THERMAL_OVERLOAD_CPU = "Thermal Overload: CPU"
RESET_REASON_THERMAL_OVERLOAD_ASIC = "Thermal Overload: ASIC"
RESET_REASON_THERMAL_OVERLOAD_OTHER = "Thermal Overload: Other"
RESET_REASON_INSUFFICIENT_FAN_SPEED = "Insufficient Fan Speed"
RESET_REASON_WATCHDOG = "Watchdog"
RESET_REASON_UNKNOWN = "Unknown"
RESET_REASON_HARDWARE_OTHER = "Hardware - Other"
RESET_REASON_NON_HARDWARE = "Non-Hardware"

# initialize the global logger
global logger
logger = logging.getLogger(SYSLOG_IDENTIFIER)
fmt = logging.Formatter(fmt='%(name)s %(levelname)s: %(message)s')
handler = logging.handlers.SysLogHandler('/dev/log')
handler.setFormatter(fmt)
logger.setLevel(logging.INFO)
#logger.setLevel(logging.DEBUG)
logger.addHandler(handler)


# ============================= Functions =============================

def find_software_reset_reason_from_csmgrd():
    software_reset_reason = RESET_REASON_UNKNOWN
    if os.path.isfile(CSMGRD_RESET_REASON_FILE):
        with open(CSMGRD_RESET_REASON_FILE) as cause_file:
            cause_file_csmgrd = cause_file.readline().rstrip('\n')
        m = re.search(REBOOT_TYPE_CSMGRD_PATTERN_WARM, cause_file_csmgrd)
        if m:
            software_reset_reason = 'warm-reboot'
        m = re.search(REBOOT_TYPE_CSMGRD_PATTERN_FAST, cause_file_csmgrd)
        if m:
            software_reset_reason = 'fast-reboot'

        logger.info("{} indicates the reboot cause: {}".format(CSMGRD_RESET_REASON_FILE, software_reset_reason))
    else:
        logger.info("Reboot cause file {} not found".format(CSMGRD_RESET_REASON_FILE))
    return software_reset_reason


def find_first_boot_version():
    build_version = "unknown"

    if os.path.isfile("/etc/os-release"):
        with open("/etc/os-release") as release_file:
            release_data = release_file.read().strip()
            m = re.search('^VERSION=(\S+)$', release_data, re.MULTILINE)
            build_version = m.group(1)

    return " (First boot of CL version {})".format(build_version)


def find_thermal_shutdown_reason():
    thermal_shutdown_reason = None
    thermal_shutdown_file = os.path.join(RESET_REASON_DIR, 'thermal_shutdown_indicator')
    if os.path.isfile(thermal_shutdown_file):
        try:
            with open(thermal_shutdown_file, 'r') as f:
                content = f.read().strip()
                if content == 'thermal_shutdown':
                    thermal_shutdown_reason = 'Thermal shutdown from NOS'
                    # Remove the indicator file after reading
                    os.remove(thermal_shutdown_file)
        except Exception as e:
            logger.error("Failed to read thermal shutdown indicator: %s", str(e))
    return thermal_shutdown_reason


def find_software_reset_reason():
    software_reset_reason = find_software_reset_reason_from_csmgrd()
    if software_reset_reason == RESET_REASON_UNKNOWN:
        #TODO: Find a way to notify platform about first boot case
        if os.path.isfile(FIRST_BOOT_PLATFORM_FILE):
            software_reset_reason += find_first_boot_version()
            os.remove(FIRST_BOOT_PLATFORM_FILE)
    return software_reset_reason


def _read_generic_file(filename, len):
    """
    Read a generic file, returns the contents of the file
    """
    result = ''
    try:
        fileobj = io.open(filename)
        result = fileobj.read(len)
        fileobj.close()
        return result
    except Exception as e:
        #logger.info("Fail to read file {} due to {}".format(filename, repr(e)))
        return '0'


def _verify_reset_reason(filename):
    '''
    Open and read the reboot cause file in
    /var/run/hwmanagement/system (which is defined as RESET_REASON_ROOT)
    If a reboot cause file doesn't exists, returns '0'.
    '''
    return bool(int(_read_generic_file(join(RESET_REASON_ROOT, filename), RESET_REASON_FILE_LENGTH).rstrip('\n')))


def get_reset_reason_hw_mgmt():
    """
    Retrieves the cause of the previous reboot

    Returns:
        A tuple (string, string) where the first element is a string
        containing the cause of the previous reboot. This string must be
        one of the predefined strings in this class. If the first string
        is "RESET_REASON_HARDWARE_OTHER", the second string can be used
        to pass a description of the reboot cause.
    """
    #read reboot causes files in the following order
    reboot_major_cause_dict = {
        'reset_main_pwr_fail'       :   RESET_REASON_POWER_LOSS,
        'reset_aux_pwr_or_ref'      :   RESET_REASON_POWER_LOSS,
        'reset_aux_pwr_or_fu'       :   RESET_REASON_POWER_LOSS,
        'reset_comex_pwr_fail'      :   RESET_REASON_POWER_LOSS,
        'reset_ac_pwr_fail'         :   RESET_REASON_AC_POWER_LOSS,
        'reset_asic_thermal'        :   RESET_REASON_THERMAL_OVERLOAD_ASIC,
        'reset_comex_thermal'       :   RESET_REASON_THERMAL_OVERLOAD_CPU,
        'reset_hotswap_or_wd'       :   RESET_REASON_WATCHDOG,
        'reset_comex_wd'            :   RESET_REASON_WATCHDOG,
        'reset_swb_wd'              :   RESET_REASON_WATCHDOG,
        'reset_sff_wd'              :   RESET_REASON_WATCHDOG,
        'reset_fw_reset'            :   RESET_REASON_RESET_FROM_FW,
        'reset_from_asic'           :   RESET_REASON_RESET_FROM_ASIC,
        'reset_system'              :   RESET_REASON_RESET_SYSTEM,
        'reset_soc'                 :   RESET_REASON_RESET_FROM_SOC,
        'reset_sw_pwr_off'          :   RESET_REASON_SW_PWR_OFF,
        'reset_platform'            :   RESET_REASON_RESET_PLATFORM,
        'reset_ac_pwr_fail'         :   RESET_REASON_POWER_LOSS,
        'reset_pwr_converter_fail'  :   RESET_REASON_PWR_CONVERTER_FAIL,
        'reset_swb_dc_dc_pwr_fail'  :   RESET_REASON_SWB_DC_DC_PWR_FAIL

    }
    reboot_minor_cause_dict = {
        'reset_sw_reset'            :   "Reset by Software",
        'reset_long_pb'             :   "Long push button",
        'reset_short_pb'            :   "Short push button",
        'reset_reload_bios'         :   "BIOS",
        'reset_hotswap_or_halt'     :   "Reset caused by hotswap or halt",
        'reset_voltmon_upgrade_fail':   "Reset due to voltage monitor devices upgrade failure"
    }

    for reset_file, reset_cause in list(reboot_major_cause_dict.items()):
        if _verify_reset_reason(reset_file):
            return reset_cause, ''

    for reset_file, reset_cause in list(reboot_minor_cause_dict.items()):
        if _verify_reset_reason(reset_file):
            return RESET_REASON_HARDWARE_OTHER, reset_cause

    return RESET_REASON_NON_HARDWARE, ''


def find_hardware_reset_reason():
    hardware_reset_reason_major, hardware_reset_reason_minor = get_reset_reason_hw_mgmt()
    if hardware_reset_reason_major:
        logger.info("Platform api indicates reboot cause {}, {}".format(hardware_reset_reason_major, hardware_reset_reason_minor))
    else:
        logger.info("No reboot cause found from platform api")

    hardware_minor_reason = f" ({hardware_reset_reason_minor})" if hardware_reset_reason_minor else ""
    hardware_reset_reason = "{}{}".format(hardware_reset_reason_major, hardware_minor_reason)

    return hardware_reset_reason


def get_reset_reason_dict(previous_reset_reason, comment, gen_time):
    """Store the key infomation of device reboot into a dictionary by parsing the string in
    previous_reset_reason.

    If user issused a command to reboot device, then user, command and time will be
    stored into a dictionary.

    If device was rebooted due to the kernel panic, then the string `Kernel Panic` 
    and time will be stored into a dictionary.
    """
    reset_reason_dict = OrderedDict()
    reset_reason_dict['gen_time'] = gen_time
    reset_reason_dict['cause'] = previous_reset_reason
    reset_reason_dict['user'] = "system/root"
    reset_reason_dict['time'] = "N/A"
    reset_reason_dict['comment'] = comment if comment is not None else "N/A"

    '''
    if re.search(r'User issued', previous_reset_reason):
        # Match with "User issued '{}' command [User: {}, Time: {}]"
        match = re.search(r'User issued \'(.*)\' command \[User: (.*), Time: (.*)\]', previous_reset_reason)
        if match is not None:
            reset_reason_dict['cause'] = match.group(1)
            reset_reason_dict['user'] = match.group(2)
            reset_reason_dict['time'] = match.group(3)
    elif re.search(r'Kernel Panic', previous_reset_reason):
        match = re.search(r'Kernel Panic \[Time: (.*)\]', previous_reset_reason)
        if match is not None:
            reset_reason_dict['cause'] = "Kernel Panic"
            reset_reason_dict['time'] = match.group(1)
    '''

    return reset_reason_dict


def check_log_file_count(reset_path):
    list_of_files = os.listdir(reset_path)

    if len(list_of_files) >= 10:
        oldest_file = min([reset_path+f for f in list_of_files], key=os.path.getctime)
        #print "Removing the oldest file {}".format(oldest_file)
        os.remove(oldest_file)


def main():

    logger.info("Starting up...")

    if not os.geteuid() == 0:
        logger.error("User {} does not have permission to execute".format(pwd.getpwuid(os.getuid()).pw_name))
        sys.exit("This utility must be run as root")

    # Create RESET_REASON_DIR if it doesn't exist
    if not os.path.exists(RESET_REASON_DIR):
        os.makedirs(RESET_REASON_DIR)

    # Remove stale PREVIOUS_RESET_REASON_FILE if it exists
    if os.path.exists(PREVIOUS_RESET_REASON_FILE):
        os.remove(PREVIOUS_RESET_REASON_FILE)

    # This variable is kept for future-use purpose. When proc_cmd_line/csmgrd/software provides
    # any additional_reboot_info it will be stored as a "comment" in RESET_REASON_HISTORY_FILE
    additional_reboot_info = "N/A"

    # First check for NOS thermal shutdown
    thermal_shutdown_reason = find_thermal_shutdown_reason()
    if thermal_shutdown_reason:
        previous_reset_reason = thermal_shutdown_reason
    else:
        previous_reset_reason = find_hardware_reset_reason()
        if previous_reset_reason.startswith(RESET_REASON_NON_HARDWARE):
            # If the reboot cause is non-hardware, get the reboot cause from CSMGRD_RESET_REASON_FILE
            previous_reset_reason = find_software_reset_reason()

    # Current time
    local_time = datetime.datetime.now(pytz.UTC)
    reset_reason_gen_time = local_time.isoformat()

    # Save the previous cause info into its history file as json format
    reset_reason_dict = get_reset_reason_dict(previous_reset_reason, additional_reboot_info, reset_reason_gen_time)

    # Create reset-reason-#time#.json under history directory
    RESET_REASON_HISTORY_FILE = os.path.join(RESET_REASON_HISTORY_DIR, "reset-reason-{}.json".format(reset_reason_gen_time))

    # Create RESET_REASON_HISTORY_DIR if it doesn't exist
    if not os.path.exists(RESET_REASON_HISTORY_DIR):
        os.makedirs(RESET_REASON_HISTORY_DIR)

    # Ensure that log file count doesnt increase the MAX
    check_log_file_count(RESET_REASON_HISTORY_DIR)

    # Write the previous reboot cause to RESET_REASON_HISTORY_FILE as a JSON format
    with open(RESET_REASON_HISTORY_FILE, "w") as reset_reason_history_file:
        json.dump(reset_reason_dict, reset_reason_history_file, indent=2)
        reset_reason_history_file.write('\n')

    # Create a symbolic link to previous-reset-reason.json file
    os.symlink(RESET_REASON_HISTORY_FILE, PREVIOUS_RESET_REASON_FILE)


if __name__ == "__main__":
    main()
