# CHIPSEC: Platform Security Assessment Framework
# Copyright (c) 2017, Google
# Copyright (c) 2021, Intel Corporation
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; Version 2.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#

"""
>>> chipsec_util reg read <reg_name> [<field_name>]
>>> chipsec_util reg read_field <reg_name> <field_name>
>>> chipsec_util reg write <reg_name> <value>
>>> chipsec_util reg write_field <reg_name> <field_name> <value>
>>> chipsec_util reg get_control <control_name>
>>> chipsec_util reg set_control <control_name> <value>

Examples:

>>> chipsec_util reg read SMBUS_VID
>>> chipsec_util reg read HSFC FGO
>>> chipsec_util reg read_field HSFC FGO
>>> chipsec_util reg write SMBUS_VID 0x8088
>>> chipsec_util reg write_field BC BLE 0x1
>>> chipsec_util reg get_control BiosWriteEnable
>>> chipsec_util reg set_control BiosLockEnable 0x1
"""

from chipsec.command import BaseCommand, toLoad
from argparse import ArgumentParser


class RegisterCommand(BaseCommand):

    def requirements(self) -> toLoad:
        return toLoad.All

    def parse_arguments(self) -> None:
        parser = ArgumentParser(prog='chipsec_util reg', usage=__doc__)
        subparsers = parser.add_subparsers()

        parser_read = subparsers.add_parser('read')
        parser_read.add_argument('reg_name', type=str, help='Register name')
        parser_read.add_argument('field_name', type=str, nargs='?', default=None, help='Field name')
        parser_read.set_defaults(func=self.reg_read)

        parser_readfield = subparsers.add_parser('read_field')
        parser_readfield.add_argument('reg_name', type=str, help='Register name')
        parser_readfield.add_argument('field_name', type=str, help='Field name')
        parser_readfield.set_defaults(func=self.reg_read_field)

        parser_write = subparsers.add_parser('write')
        parser_write.add_argument('reg_name', type=str, help='Register name')
        parser_write.add_argument('value', type=lambda x: int(x, 16), help='Value (hex)')
        parser_write.set_defaults(func=self.reg_write)

        parser_writefield = subparsers.add_parser('write_field')
        parser_writefield.add_argument('reg_name', type=str, help='Register name')
        parser_writefield.add_argument('field_name', type=str, help='Field name')
        parser_writefield.add_argument('value', type=lambda x: int(x, 16), help='Value (hex)')
        parser_writefield.set_defaults(func=self.reg_write_field)

        parser_getcontrol = subparsers.add_parser('get_control')
        parser_getcontrol.add_argument('control_name', type=str, help='Control name')
        parser_getcontrol.set_defaults(func=self.reg_get_control)

        parser_setcontrol = subparsers.add_parser('set_control')
        parser_setcontrol.add_argument('control_name', type=str, help='Control name')
        parser_setcontrol.add_argument('value', type=lambda x: int(x, 16), help='Value (hex)')
        parser_setcontrol.set_defaults(func=self.reg_set_control)

        parser.parse_args(self.argv, namespace=self)

    def reg_read(self):
        if self.field_name is not None:
            reglist = self.cs.register.get_list_by_name(self.reg_name)
            values = reglist.read_field(self.field_name)
            for value in values:
                self.logger.log(f"[CHIPSEC] {self.reg_name}.{self.field_name}=0x{value:X}")
        else:
            reglist = self.cs.register.get_list_by_name(self.reg_name)
            if len(reglist) == 0:
                self.logger.log(f'No register found with the name {self.reg_name}')
            else:
                reglist.read_and_print()

    def reg_read_field(self):
        if self.cs.register.has_field(self.reg_name, self.field_name):
            reglist = self.cs.register.get_list_by_name(self.reg_name)
            field_list = reglist.read_field(self.field_name)
            for value in field_list:
                self.logger.log(f"[CHIPSEC] {self.reg_name}.{self.field_name} = 0x{value:X}")
        else:
            self.logger.log_error("[CHIPSEC] Register '{}' doesn't have field '{}' defined".format(self.reg_name, self.field_name))

    def reg_write(self):
        self.logger.log(f"[CHIPSEC] Writing {self.reg_name} < 0x{self.value:X}")
        self.cs.register.get_list_by_name(self.reg_name).write(self.value)

    def reg_write_field(self):
        if self.cs.register.has_field(self.reg_name, self.field_name):
            self.logger.log(f"[CHIPSEC] Writing {self.reg_name}.{self.field_name} < 0x{self.value:X}")
            self.cs.register.get_list_by_name(self.reg_name).write_field(self.field_name, self.value)
        else:
            self.logger.log_error("[CHIPSEC] Register '{self.reg_name}' doesn't have field '{self.field_name}' defined")

    def reg_get_control(self):
        if self.cs.control.is_defined(self.control_name):
            ctrl_list = self.cs.control.get_list_by_name(self.control_name)
            for ctrl in ctrl_list:
                ctrl.read()
                self.logger.log(f"[CHIPSEC] {self.control_name} = 0x{ctrl.value:X}")
        else:
            self.logger.log_error("[CHIPSEC] Control '{}' isn't defined".format(self.control_name))

    def reg_set_control(self):
        if self.cs.control.is_defined(self.control_name):
            self.cs.control.get_list_by_name(self.control_name).write(self.value)
            self.logger.log("[CHIPSEC] Setting control {} < 0x{:X}".format(self.control_name, self.value))
        else:
            self.logger.log_error("[CHIPSEC] Control '{}' isn't defined".format(self.control_name))

commands = {'reg': RegisterCommand}