##
# Utility Functions to support re-use in python scripts.
# Includes functions for running external commands, etc
#
# Copyright (c) Microsoft Corporation
#
# SPDX-License-Identifier: BSD-2-Clause-Patent
##
import re
import os
import stat
import logging
import datetime
import time
import shutil
import threading
import subprocess
import sys
import inspect
import platform
import importlib
from collections import namedtuple
from enum import Enum as StdEnum
import locale


def Enum(*args):
    items = []
    if len(args) == 1:
        item = args[0]
        if isinstance(item, list):
            items = item
        elif isinstance(item, tuple):
            items = list(item)
        else:
            items.append(item)
    else:
        items = args
    calling_frame = inspect.stack()[1]
    calling_mod = inspect.getmodule(calling_frame[0])
    calling_func = calling_frame.function
    calling_line = calling_frame.lineno

    enum_name = calling_func + ":" + str(calling_line)
    logging.warning(f"Enum is deprecated! Please fix it in {calling_mod}:{calling_func}")
    return StdEnum(enum_name, items, module=calling_mod)


####
# Class to support running commands from the shell in a python environment.
# Don't use directly.
#
# PropagatingThread copied from sample here:
# https://stackoverflow.com/questions/2829329/catch-a-threads-exception-in-the-caller-thread-in-python
####
class PropagatingThread(threading.Thread):
    def run(self):
        self.exc = None
        try:
            if hasattr(self, '_Thread__target'):
                # Thread uses name mangling prior to Python 3.
                self.ret = self._Thread__target(*self._Thread__args, **self._Thread__kwargs)
            else:
                self.ret = self._target(*self._args, **self._kwargs)
        except BaseException as e:
            self.exc = e

    def join(self, timeout=None):
        super(PropagatingThread, self).join()
        if self.exc:
            raise self.exc
        return self.ret


####
# Helper functions for running commands from the shell in python environment
# Don't use directly
#
# process output stream and write to log.
# part of the threading pattern.
#
#  http://stackoverflow.com/questions/19423008/logged-subprocess-communicate
####
def reader(filepath, outstream, stream, logging_level=logging.INFO, encodingErrors='strict'):
    f = None
    # open file if caller provided path
    if(filepath):
        f = open(filepath, "w")

    (_, encoding) = locale.getdefaultlocale()
    while True:
        s = stream.readline().decode(encoding, errors=encodingErrors)
        if not s:
            break
        if(f is not None):
            # write to file if caller provided file
            f.write(s)
        if(outstream is not None):
            # write to stream object if caller provided object
            outstream.write(s)
        logging.log(logging_level, s.rstrip())
    stream.close()
    if(f is not None):
        f.close()


####
# Returns a namedtuple containing information about host machine.
#
# @return namedtuple Host(os=OS Type, arch=System Architecture, bit=Highest Order Bit)
####
def GetHostInfo():
    Host = namedtuple('Host', 'os arch bit')
    host_info = platform.uname()
    os = host_info.system
    processor_info = host_info.machine

    arch = None
    bit = None

    if ("x86" in processor_info.lower()) or ("AMD" in processor_info.upper()) or ("INTEL" in processor_info.upper()):
        arch = "x86"
    elif ("ARM" in processor_info.upper()) or ("AARCH" in processor_info.upper()):
        arch = "ARM"

    if "32" in processor_info:
        bit = "32"
    elif "64" in processor_info:
        bit = "64"

    if (arch is None) or (bit is None):
        raise EnvironmentError("Host info could not be parsed: {0}".format(str(host_info)))

    return Host(os=os, arch=arch, bit=bit)


####
# This is a mixing to do timing on a function. Use it like this:
#
# @timing
# def function_i_want_to_time():
#   ...
####
def timing(f):
    def wrap(*args):
        time1 = time.time()
        ret = f(*args)
        time2 = time.time()
        logging.debug('{:s} function took {:.3f} ms'.format(f.__name__, (time2 - time1) * 1000.0))
        return ret
    return wrap


####
# Run a shell command and print the output to the log file
# This is the public function that should be used to run commands from the shell in python environment
# @param cmd - command being run, either quoted or not quoted
# @param parameters - parameters string taken as is
# @param capture - boolean to determine if caller wants the output captured in any format.
# @param workingdir - path to set to the working directory before running the command.
# @param outfile - capture output to file of given path.
# @param outstream - capture output to a stream.
# @param environ - shell environment variables dictionary that replaces the one inherited from the
#                  current process.
# @param logging_level - log level to log output at.  Default is INFO
# @param raise_exception_on_nonzero - Setting to true causes exception to be raised if the cmd
#                                     return code is not zero.
# @param encodingErrors - may be given to set the desired error handling for encoding errors decoding cmd output.
#                         Default is 'strict'.
#
# @return returncode of called cmd
####
def RunCmd(cmd, parameters, capture=True, workingdir=None, outfile=None, outstream=None, environ=None,
           logging_level=logging.INFO, raise_exception_on_nonzero=False, encodingErrors='strict'):
    cmd = cmd.strip('"\'')
    if " " in cmd:
        cmd = '"' + cmd + '"'
    if parameters is not None:
        parameters = parameters.strip()
        cmd += " " + parameters
    starttime = datetime.datetime.now()
    logging.log(logging_level, "Cmd to run is: " + cmd)
    logging.log(logging_level, "------------------------------------------------")
    logging.log(logging_level, "--------------Cmd Output Starting---------------")
    logging.log(logging_level, "------------------------------------------------")
    c = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=workingdir, shell=True, env=environ)
    if(capture):
        thread = PropagatingThread(target=reader, args=(outfile, outstream, c.stdout, logging_level, encodingErrors))
        thread.start()
        c.wait()
        thread.join()
    else:
        c.wait()

    endtime = datetime.datetime.now()
    delta = endtime - starttime
    endtime_str = "{0[0]:02}:{0[1]:02}".format(divmod(delta.seconds, 60))
    returncode_str = "{0:#010x}".format(c.returncode)
    logging.log(logging_level, "------------------------------------------------")
    logging.log(logging_level, "--------------Cmd Output Finished---------------")
    logging.log(logging_level, "--------- Running Time (mm:ss): " + endtime_str + " ----------")
    logging.log(logging_level, "----------- Return Code: " + returncode_str + " ------------")
    logging.log(logging_level, "------------------------------------------------")

    if raise_exception_on_nonzero and c.returncode != 0:
        raise Exception("{0} failed with Return Code: {1}".format(cmd, returncode_str))
    return c.returncode

####
# Run a python script and print the output to the log file
# This is the public function that should be used to execute python scripts from the shell in python environment.
# The python script will be located using the path as if it was an executable.
#
# @param cmd - cmd string to run including parameters
# @param capture - boolean to determine if caller wants the output captured in any format.
# @param workingdir - path to set to the working directory before running the command.
# @param outfile - capture output to file of given path.
# @param outstream - capture output to a stream.
# @param environ - shell environment variables dictionary that replaces the one inherited from the
#                  current process.
#
# @return returncode of called cmd
####


def RunPythonScript(pythonfile, params, capture=True, workingdir=None, outfile=None, outstream=None,
                    environ=None, logging_level=logging.INFO, raise_exception_on_nonzero=False):
    # locate python file on path
    pythonfile.strip('"\'')
    if " " in pythonfile:
        pythonfile = '"' + pythonfile + '"'
    params.strip()
    logging.debug("RunPythonScript: {0} {1}".format(pythonfile, params))
    if(os.path.isabs(pythonfile)):
        logging.debug("Python Script was given as absolute path: %s" % pythonfile)
    elif(os.path.isfile(os.path.join(os.getcwd(), pythonfile))):
        pythonfile = os.path.join(os.getcwd(), pythonfile)
        logging.debug("Python Script was given as relative path: %s" % pythonfile)
    else:
        # loop thru path environment variable
        for a in os.getenv("PATH").split(os.pathsep):
            a = os.path.normpath(a)
            if os.path.isfile(os.path.join(a, pythonfile)):
                pythonfile = os.path.join(a, pythonfile)
                logging.debug("Python Script was found on the path: %s" % pythonfile)
                break
    params = pythonfile + " " + params
    return RunCmd(sys.executable, params, capture=capture, workingdir=workingdir, outfile=outfile,
                  outstream=outstream, environ=environ, logging_level=logging_level,
                  raise_exception_on_nonzero=raise_exception_on_nonzero)
####
# Locally Sign input file using Windows SDK signtool.  This will use a local Pfx file.
# WARNING!!! : This should not be used for production signing as that process should follow stronger
#               security practices (HSM / smart cards / etc)
#
#  Signing is in format specified by UEFI authentacted variables
####


def DetachedSignWithSignTool(SignToolPath, ToSignFilePath, SignatureOutputFile, PfxFilePath,
                             PfxPass=None, Oid="1.2.840.113549.1.7.2", Eku=None):

    # check signtool path
    if not os.path.exists(SignToolPath):
        logging.error("Path to signtool invalid.  %s" % SignToolPath)
        return -1

    # Adjust for spaces in the path (when calling the command).
    if " " in SignToolPath:
        SignToolPath = '"' + SignToolPath + '"'

    OutputDir = os.path.dirname(SignatureOutputFile)
    # Signtool docs https://docs.microsoft.com/en-us/dotnet/framework/tools/signtool-exe
    # Signtool parameters from
    #   https://docs.microsoft.com/en-us/windows-hardware/manufacture/desktop/secure-boot-key-generation-and-signing-using-hsm--example  # noqa: E501
    # Search for "Secure Boot Key Generation and Signing Using HSM"
    params = 'sign /fd sha256 /p7ce DetachedSignedData /p7co ' + Oid + ' /p7 "' + \
             OutputDir + '" /f "' + PfxFilePath + '"'
    if Eku is not None:
        params += ' /u ' + Eku
    if PfxPass is not None:
        # add password if set
        params = params + ' /p ' + PfxPass
    params = params + ' /debug /v "' + ToSignFilePath + '" '
    ret = RunCmd(SignToolPath, params)
    if(ret != 0):
        logging.error("Signtool failed %d" % ret)
        return ret
    signedfile = os.path.join(OutputDir, os.path.basename(ToSignFilePath) + ".p7")
    if(not os.path.isfile(signedfile)):
        raise Exception("Output file doesn't exist %s" % signedfile)

    shutil.move(signedfile, SignatureOutputFile)
    return ret

####
# Locally Sign input file using Windows SDK signtool.  This will use a local Pfx file.
# WARNING!!! : This should not be used for production signing as that process should follow
#              stronger security practices (HSM / smart cards / etc)
#
#  Signing is catalog format which is an attached signature
####


def CatalogSignWithSignTool(SignToolPath, ToSignFilePath, PfxFilePath, PfxPass=None):

    # check signtool path
    if not os.path.exists(SignToolPath):
        logging.error("Path to signtool invalid.  %s" % SignToolPath)
        return -1

    # Adjust for spaces in the path (when calling the command).
    if " " in SignToolPath:
        SignToolPath = '"' + SignToolPath + '"'

    OutputDir = os.path.dirname(ToSignFilePath)
    # Signtool docs https://docs.microsoft.com/en-us/dotnet/framework/tools/signtool-exe
    # todo: link to catalog signing documentation
    params = "sign /a /fd SHA256 /f " + PfxFilePath
    if PfxPass is not None:
        # add password if set
        params = params + ' /p ' + PfxPass
    params = params + ' /debug /v "' + ToSignFilePath + '" '
    ret = RunCmd(SignToolPath, params, workingdir=OutputDir)
    if(ret != 0):
        logging.error("Signtool failed %d" % ret)
    return ret


###
# Function to print a byte list as hex and optionally output ascii as well as
# offset within the buffer
###
def PrintByteList(ByteList, IncludeAscii=True, IncludeOffset=True, IncludeHexSep=True, OffsetStart=0):
    Ascii = ""
    for index in range(len(ByteList)):
        # Start of New Line
        if(index % 16 == 0):
            if(IncludeOffset):
                print("0x%04X -" % (index + OffsetStart), end='')

        # Midpoint of a Line
        if(index % 16 == 8):
            if(IncludeHexSep):
                print(" -", end='')

        # Print As Hex Byte
        print(" 0x%02X" % ByteList[index], end='')

        # Prepare to Print As Ascii
        if(ByteList[index] < 0x20) or (ByteList[index] > 0x7E):
            Ascii += "."
        else:
            Ascii += ("%c" % ByteList[index])

        # End of Line
        if(index % 16 == 15):
            if(IncludeAscii):
                print(" %s" % Ascii, end='')
            Ascii = ""
            print("")

    # Done - Lets check if we have partial
    if(index % 16 != 15):
        # Lets print any partial line of ascii
        if(IncludeAscii) and (Ascii != ""):
            # Pad out to the correct spot

            while(index % 16 != 15):
                print("     ", end='')
                if(index % 16 == 7):  # account for the - symbol in the hex dump
                    if(IncludeOffset):
                        print("  ", end='')
                index += 1
            # print the ascii partial line
            print(" %s" % Ascii, end='')
            # print a single newline so that next print will be on new line
        print("")


# Simplified Comparison Function borrowed from StackOverflow...
# https://stackoverflow.com/questions/1714027/version-number-comparison
# With Python 3.0 help from:
# https://docs.python.org/3.0/whatsnew/3.0.html#ordering-comparisons
def version_compare(version1, version2):
    def normalize(v):
        return [int(x) for x in re.sub(r'(\.0+)*$', '', v).split(".")]
    (a, b) = (normalize(version1), normalize(version2))
    return (a > b) - (a < b)


def import_module_by_file_name(module_file_path):
    ''' Standard method of importing a Python file. Expecting absolute path. '''
    module_name = os.path.basename(module_file_path)
    spec = importlib.util.spec_from_file_location(module_name, module_file_path)

    if spec is None:
        raise RuntimeError(f"Expected module file named {module_file_path}")

    ImportedModule = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(ImportedModule)

    return ImportedModule


def locate_class_in_module(Module, DesiredClass):
    '''
    Given a module and a class, this function will return the subclass of DesiredClass found in Module.
    It gives preference to classes that are defined in the module itself.
    This means that if you have an import that subclasses DesiredClass, it will be picked unless
    there is a class defined in the module that subclasses DesiredClass.

    In this hypothetical class hierarchy, GrandChildClass would be picked
    --------------      ------------      -----------------
    |DesiredClass|  ->  |ChildClass|  ->  |GrandChildClass|
    --------------      ------------      -----------------
    '''

    DesiredClassInstance = None
    # Pull out the contents of the module that was provided
    module_contents = dir(Module)
    # Filter through the Module, we're only looking for classes.
    classList = [getattr(Module, obj) for obj in module_contents
                 if inspect.isclass(getattr(Module, obj))]

    for _class in classList:
        # Classes that the module import show up in this list too so we need
        # to make sure it's an INSTANCE of DesiredClass, not DesiredClass itself!
        # if multiple instances are found in the same class hierarchy, pick the
        # most specific one. If multiple instances are found belonging to different
        # class hierarchies, raise an error.
        if _class is not DesiredClass and issubclass(_class, DesiredClass):
            if (DesiredClassInstance is None) or issubclass(_class, DesiredClassInstance):
                DesiredClassInstance = _class
            elif not issubclass(DesiredClassInstance, _class):
                raise RuntimeError(f"Multiple instances were found:\n\t{DesiredClassInstance}\n\t{_class}")
    return DesiredClassInstance


def RemoveTree(dir_path: str, ignore_errors: bool = False) -> None:
    '''
    Helper for removing a directory.  Over time there have been
    many private implementations of this due to reliability issues in the
    shutil implementations.  To consolidate on a single function this helper is added.

    On error try to change file attributes.  Also add retry logic.

    dir_path: Path to directory to remove.
    ignore_errors: ignore errors during removal
    '''

    def remove_readonly(func, path, _):
        ''' private function to attempt to change permissions on file/folder being deleted'''
        os.chmod(path, stat.S_IWRITE)
        func(path)

    for _ in range(3):  # retry up to 3 times
        try:
            shutil.rmtree(dir_path, ignore_errors=ignore_errors, onerror=remove_readonly)
        except OSError as err:
            logging.warning(f"Failed to fully remove {dir_path}: {err}")
        else:
            break
    else:
        raise RuntimeError(f"Failed to remove {dir_path}")
