#! python3
## @file
# GenFlashmap.py
#
# Copyright (c) 2020 - 2021, Intel Corporation. All rights reserved.
# This software and associated documentation (if any) is furnished
# under a license and may only be used or copied in accordance
# with the terms of the license. Except as permitted by such
# license, no part of this software or documentation may be
# reproduced, stored in a retrieval system, or transmitted in any
# form or by any means without the express written consent of
# Intel Corporation.
#
#

##
# This script supports new flashmap FDF generation by using a template FDF file and a base flashmap FDF file.
#

"""
usage: GenFlashmap.py [optional arguments] [positional arguments]

Flash Map FDF file generation script

positional arguments:
  {extended}
    extended            Command to generate the flashmap supporting Extended BIOS Region

optional arguments:
  -h, --help            show this help message and exit
  -b BASEFLASHMAP, --baseflashmap BASEFLASHMAP
                        Base FlashMapInclude FDF file path
  -t TEMPLATE, --template TEMPLATE
                        Template FlashMap FDF for Extended BIOS Region file path
  -o OUT_FILE, --out_file OUT_FILE
                        Output FDF file path
"""

import os
import re
import sys
import argparse
from pprint import pprint

__prog__          = sys.argv[0]
__description__   = "Flash Map FDF file generation script"

SIZE_1MB  = 0x100000
SIZE_16MB = 0x1000000
SIZE_4GB  = 0x100000000

##
# This class has base flashmap resource.
#
# param[in]  Base flashmap file path
#
class Basemap ():
    def __init__ (self, in_file):
        self.file_path = in_file
        self.desc = ["#\n#   Generated by the base flashmap {}\n#\n\n".format (self.file_path)]
        with open (self.file_path, "r") as fdf:
            self.__full = fdf.readlines ()

        # This list stores whole lines settings PCDs.
        self.map_list = self.__full [[self.__full.index (i) for i in self.__full if re.match (r"(?m)^\s*(SET|!)\s*", i)][0]:]
        self.map_list = self.desc + self.map_list

        # This list stores all flashmap offset PCD names except NVS region.
        # Do NOT include size PCD names.
        self.fv_offset_list = [
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvPreMemoryOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvSecurityOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvFspTOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvFspMOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvFspSOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvPostMemoryOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvMicrocodeOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashFvFirmwareBinariesOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashFvRsvdOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashIbbOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashIbbROffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvUefiBootOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvOsBootOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashFvOptionalOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashFvAdvancedOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashFvTestMenuOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashFvValidationOffset",
        "gBoardModuleTokenSpaceGuid.PcdFlashObbOffset"
        ]

        # This list stores flashmap offset PCD names for NVS region.
        self.nvs_offset_list = [
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashNvStorageFtwSpareOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashNvStorageFtwWorkingOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashNvStorageVariableOffset",
        "gMinPlatformPkgTokenSpaceGuid.PcdFlashNvStorageOffset"
        ]

        # The following variables store settings which are defined by macro.
        self.flash_base       = "FLASH_BASE"
        self.flash_size       = "FLASH_SIZE"
        self.flash_block_size = "FLASH_BLOCK_SIZE"
        self.flash_num_blocks = "FLASH_NUM_BLOCKS"

    # This method returns empty region size out of 16MB.
    def get_empty_size (self):
        if SIZE_4GB - CommonFunc.read (self.flash_base, self.__full) < SIZE_16MB:
            size = SIZE_16MB - (SIZE_4GB - CommonFunc.read (self.flash_base, self.__full))
        else:
            size = 0
        return size

    #
    # This method self-diagnoses the resource in the class.
    # The script can not generate a correct flashmap file if this class has some missing settings,
    # ending up a build failure around the end of build process. Running the self-diagnosis logic
    # on a file generation helps for capturing missing settings at the begining of build process.
    # It self-diagnoses by the following steps,
    #  1.  parsing the list which has all the lines for setting PCDs in a flashmap file,
    #      and finding the PCD names which suffix "Offset",
    #  2.  looking for the matching "Offset" PCDs in all PCD list of this class (this_class_list),
    #  3.  capturing the lines which contain the missing "Offset" PCD, and showing the PCD names.
    # If missing PCDs are captured, self-diagnosis method returns False. Otherwise returns True.
    # No failure should occur even if all the PCD list in this class support more than the PCDs
    # which are actually supported by a flashmap.
    #
    def self_diagnosis (self):

        this_class_list = self.fv_offset_list + self.nvs_offset_list   # this list is diagnosed by the logic

        result = True
        pattern_offset = "(?P<pcd>g.*Guid\.Pcd.*Offset)"
        pattern = r"(?m)^\s*SET\s*" + pattern_offset + "\s*=\s*0x(?P<value>[0-9a-fA-F]*).*"
        missing_pcd = []
        for full in self.map_list:
            m_offset_pcd_in_full = re.match (pattern, full)
            if m_offset_pcd_in_full:
                match = False
                for this_class_pcd_list in this_class_list:
                    if this_class_pcd_list in full:
                        match = True
                        break
                if not match:
                    missing_pcd.append (m_offset_pcd_in_full.group ("pcd"))

        if missing_pcd:
            result = False
            error_message =  "  Found the PCD(s) are not supported by the script : {}\n".format (missing_pcd)
            error_message += "  Please add the PCD(s) to Basemap class in the script ({})\n".format (__prog__)
            error_message += "  so the script can help for generating a new flashmap file\n"
            error_message += "  based on the flashmap ({})".format (self.file_path)

        if not result:
            print (" Self-diagnosis : {}".format (result))
            print (error_message)
        else:
            result = True

        return result

##
# This class has Extended BIOS Region template flashmap resource.
#
# param[in]  Extended BIOS Region template flashmap file path
#
class ExtendedBiosRegionmap ():
    def __init__ (self, in_file):
        self.file_path = in_file
        self.desc = ["#   Template file {}\n#\n".format (self.file_path)]
        with open (self.file_path, "r") as fdf:
            self.__full = fdf.readlines ()

        # This list stores whole lines which set PCDs.
#        self.map_list = self.__full [[self.__full.index (i) for i in self.__full if re.match (r"(?m)^\s*SET\s*", i)][0]:]
        self.map_list = self.__full [[self.__full.index (i) for i in self.__full if re.match (r"(?m)^\s*(^DEFINE|#)\s*", i)][0]:]

        # The following variables store settings which are defined by macro.
        self.flash_base        = "FLASH_BASE"
        self.flash_size        = "FLASH_SIZE"
        self.flash_block_size  = "FLASH_BLOCK_SIZE"
        self.flash_num_blocks  = "FLASH_NUM_BLOCKS"
        self.extended_region_memmap_address = "EXTENDED_REGION_MEMMAP_ADDRESS"
        self.extended_region_size   = "EXTENDED_REGION_SIZE"
        self.extended_region_in_use = "EXTENDED_REGION_IN_USE"

    # This method inserts comments to the empty region out of 16MB region in a flashmap file.
    # This purposes just to explicitly show the empty region presence in the auto generated flashmap file.
    def insert_comment_to_fill_16mb_region (self, empty_size, in_list):
        comment = []
        comment.append ("\n")
        comment.append ("#===================================================================================#\n")
        comment.append ("#                                                                                   #\n")
        comment.append ("#   Empty {0:.0f} MB to fill up 16MB region                                   0x{1:08X}  #\n".format (empty_size/SIZE_1MB, empty_size))
        comment.append ("#                                                                                   #\n")
        comment.append ("#===================================================================================#\n")
        comment.append("\n")
        pattern_nvs = "g.*\.PcdFlashNvStorage.*"
        pattern_pcd = r"(?m)^\s*SET\s*" + pattern_nvs + "\s*=\s*0x(?P<value>[0-9a-fA-F]*).*"
        end_of_nvs = [in_list.index (i) for i in in_list if re.match (pattern_pcd, i)][-1] + 1
        out_list = in_list [:end_of_nvs]
        out_list.extend (comment)
        out_list.extend (in_list [end_of_nvs + 1:])
        return out_list

    # This method returns extended region size.
    def get_extended_region_size (self):
        return CommonFunc.read (self.extended_region_size, self.__full)

##
# Common Functions
#
class CommonFunc ():
    def __init__ (self):
        pass
    ##
    # This method parses an list which stores whole lines of a flashmap file
    #  and returns a value found by a setting name.
    #
    # param[in]  setting               Setting name. PCD name or macro name
    # param[in]  in_list               List which stores flashmap FDF file lines
    # reval      return_value          Value found by a setting name
    #
    # Note : This method assumes the flashmap fdf has only one entry of the setting.
    #
    @classmethod
    def read (cls, setting, in_list):
        setting_list = []
        pattern     = r"(?m)^\s*DEFINE\s*" + setting + "\s*=\s*0x(?P<value>[0-9a-fA-F]*).*"
        pattern_pcd = r"(?m)^\s*SET\s*" + setting + "\s*=\s*0x(?P<value>[0-9a-fA-F]*).*"
        pattern_acm_option = r"(?m)^\s*DEFINE\s*" + setting + "\s*=\s*(?P<value>[0-9]*K).*"
        for line in in_list:
            m = re.match (pattern, line)
            if m:
                setting_list.append (int (m.group ("value"), base = 16))
            m = re.match (pattern_pcd, line)
            if m:
                setting_list.append (int (m.group ("value"), base = 16))
            m = re.match (pattern_acm_option, line)
            if m:
                setting_list.append (m.group ("value"))
        try:
            read_value = setting_list [-1]
        except IndexError as err:
            print ("\n\n Error!!! {} is not found in {} \n\n".format (setting, self.in_file))
            raise
        finally:
            return read_value

    ##
    # This method writes a list to a file.
    #
    # param[in]  in_list               List which is written to an output file
    # param[in]  template_file         Template file path which has contents to be written
    #                                  to the top of an output file
    #                                  Input None if no template.
    # param[in]  out_file              Output file path
    # reval      return_value          0 if output file is successfully written
    #
    @classmethod
    def write_fdf (cls, in_list, template_file, out_file):
        if os.path.isfile (out_file):
            print ("  {} has already existed so overwriting ...".format (out_file))
            os.remove (out_file)

        with open (out_file, "w") as g:

            if template_file is not None:
                with open (template_file, "r") as t:
                    template = t.read ()
                for rd in template:
                    g.write (rd)

            for wt in in_list:
                g.write (wt)

        return_code = 0 if os.path.isfile (out_file) else 1
        return return_code

    ##
    # This method updates a list.
    #
    # param[in]  original_map_list     List to be updated
    # param[in]  target_setting_list   List of setting names which values are to be updated
    # param[in]  offset_add_value      Value which are added to the value of offset field
    # param[in]  addr_add_value        Value which are added to the value of addr field
    # reval      new_map_list          Updated list
    #
    @classmethod
    def update_offset (cls, original_map_list, target_setting_list, offset_add_value, addr_add_value):
        new_map_list = []
        for line in original_map_list:
            for pcd in target_setting_list:
                pattern_pcd = r"(?m)^\s*SET\s*" + pcd + "\s*=\s*0x(?P<value>[0-9a-fA-F]*).*\(0x(?P<addr>[0-9a-fA-F]{8})\).*"
                m = re.match (pattern_pcd, line)
                if m:
                    new_offset_str  = format ((int (m.group ("value"), base = 16) + offset_add_value), "08X")
                    line            = re.sub (r"=\s*0x[0-9a-fA-F]{8}", "= 0x" + new_offset_str, line)
                    new_addr_str    = format ((int (m.group ("addr"), base = 16) + addr_add_value), "08X")
                    line            = re.sub (r"\(*0x[0-9a-fA-F]{8}\)", "(0x" + new_addr_str + ")", line)
            new_map_list.append (line)
        return new_map_list


def main ():
    def is_valid_file_path (Argument):
        if os.path.isfile (Argument):
            return Argument
        else:
            message = "[" + Argument + "] is not found"
            raise argparse.ArgumentTypeError (message)

    def is_valid_32bit_all (Argument):
        try:
            value = int (Argument, base = 16)
        except:
            if Argument == "all":
                return Argument
            else:
                message = "[" + Argument + "] is not a valid 32bit number"
                raise argparse.ArgumentTypeError (message)
        if value < 0 or value > 0xffffffff:
            message = "[" + Argument + "] is not a valid 32bit number"
            raise argparse.ArgumentTypeError (message)
        return Argument

    ##
    #  Command to generates Extended BIOS Region flashmap file
    #
    #  This method reads the offset addresses from a default flashmap FDF file,
    #  adjusts the offset addresses based on the macro settings defined by the
    #  template FDF for Extended BIOS Region, and writes the adjusted flashmap
    #  settings to a new Extended BIOS Region flashmap FDF file.
    #
    def gen_extended_bios_region_flashmap (Argument):
        print ("\nGenerating FlashMap FDF supporting Extended BIOS Region ... \n")

        if args.baseflashmap is None:
            print ("Error!!! baseflashmap option is not set")
            return 1
        if args.template is None:
            print ("Error!!! template option is not set")
            return 1
        if args.out_file is None:
            print ("Error!!! out_file option is not set")
            return 1

        # initialize
        bm = Basemap (args.baseflashmap)
        if not bm.self_diagnosis (): return 1
        em = ExtendedBiosRegionmap (args.template)

        # update
        in_list   = bm.map_list
        if bm.get_empty_size () != 0:
            in_list = em.insert_comment_to_fill_16mb_region (bm.get_empty_size (), in_list)
        out_list1 = CommonFunc.update_offset (in_list,   bm.nvs_offset_list, 0, (-1) * bm.get_empty_size ())
        pprint (out_list1)
        out_list2 = CommonFunc.update_offset (out_list1, bm.nvs_offset_list, em.get_extended_region_size (), 0)
        out_list3 = CommonFunc.update_offset (out_list2, bm.fv_offset_list,  bm.get_empty_size (), 0)
        out_list4 = CommonFunc.update_offset (out_list3, bm.fv_offset_list,  em.get_extended_region_size (), 0)
        out_listn = out_list4

        # write to a new flashmap
        return_value = CommonFunc.write_fdf (out_listn, args.template, args.out_file)
        if return_value == 0:
            print ("  Flashmap fdf generation has been completed successfully".format (args.out_file))
            print ("    Base flashmap fdf      : {}".format (args.baseflashmap))
            print ("    Generated flashmap fdf : {}".format (args.out_file))

        return return_value

#
# Command options
#
    parser = argparse.ArgumentParser (prog = __prog__, description =__description__, usage = "%(prog)s [optional arguments] [positional arguments]")
    parser.add_argument ("-b", "--baseflashmap", type = is_valid_file_path, help = "Base FlashMapInclude FDF file path")
    parser.add_argument ("-t", "--template", type = is_valid_file_path, help = "Template FlashMap FDF for Extended BIOS Region file path")
    parser.add_argument ("-o", "--out_file", help = "Output FDF file path")

    subparsers = parser.add_subparsers ()
    parser_test = subparsers.add_parser ("extended", help = "Command to generate the flashmap supporting Extended BIOS Region")
    parser_test.set_defaults (handler = gen_extended_bios_region_flashmap)

    args = parser.parse_args ()

    returncode = args.handler (args) if hasattr (args, "handler") else 1
    return returncode

if __name__ == "__main__":
    sys.exit (main())