# CHIPSEC: Platform Security Assessment Framework
#
# 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.
#
#

import unittest

from chipsec.hal.acpi_tables import APIC, BGRT, DMAR, XSDT, UEFI_TABLE, WSMT


class TestACPITables(unittest.TestCase):
    """Test the ACPI Table Structures and Parsing."""

    def test_verify_dmar_table_format(self):
        test_dmar = DMAR()
        value_re = r'^=?[BHQsI\d]*'
        key_re = r'^\w+_FORMAT'
        self.assertEqual(type(test_dmar.DMAR_TABLE_FORMAT), dict)
        for key, value in test_dmar.DMAR_TABLE_FORMAT.items():
            self.assertRegex(key, key_re)
            self.assertRegex(value, value_re)

    def test_dmar_parse(self):
        test_table_content = b'&\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00\x00\xd9\xfe\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00\x00\x00 \x00\x01\x00\x00\x00\x00\x10\xd9\xfe\x00\x00\x00\x00\x03\x08\x00\x00\x02\xf0\x1f\x00\x04\x08\x00\x00\x00\x00\x1f\x00\x01\x00 \x00\x00\x00\x00\x00\x00\xa0\xb9z\x00\x00\x00\x00\xff?\xdez\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x14\x00\x01\x00 \x00\x00\x00\x00\x00\x00\x00\x80{\x00\x00\x00\x00\xff\xff\xff\x7f\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00'
        test_dmar = DMAR()
        test_dmar.parse(test_table_content)
        self.assertEqual(len(test_dmar.dmar_structures), 4)

    def test__str__(self):
        test_table_content = b'&\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00\x00\xd9\xfe\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00\x00\x00 \x00\x01\x00\x00\x00\x00\x10\xd9\xfe\x00\x00\x00\x00\x03\x08\x00\x00\x02\xf0\x1f\x00\x04\x08\x00\x00\x00\x00\x1f\x00\x01\x00 \x00\x00\x00\x00\x00\x00\xa0\xb9z\x00\x00\x00\x00\xff?\xdez\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x14\x00\x01\x00 \x00\x00\x00\x00\x00\x00\x00\x80{\x00\x00\x00\x00\xff\xff\xff\x7f\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00'
        # test_table_content = b'&\x05\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00\x00\xd9\xfe\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00`\xd8\xfe\x00\x00\x00\x00\x02\x08\x00\x00\x00\x00\x07\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00p\xd8\xfe\x00\x00\x00\x00\x02\x08\x00\x00\x00\x00\x07\x03\x00\x00 \x00\x01\x00\x00\x00\x00\x10\xd9\xfe\x00\x00\x00\x00\x03\x08\x00\x00\x02\x00\x1e\x07\x04\x08\x00\x00\x00\x00\x1e\x06\x01\x00 \x00\x00\x00\x00\x00\x00\x00\x00/\x00\x00\x00\x00\xff\xff\x7fO\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00'
        test_dmar = DMAR()
        test_dmar.parse(test_table_content)
        str_header = 'DMAR Table Contents'
        str_contents = test_dmar.__str__()
        self.assertIn(str_header, str_contents)
        self.assertTrue(1500 <= len(str_contents) <= 2000)

    def test_get_structure_dmar_type_7(self):
        test_dmar = DMAR()
        type = 0x07
        structure_DMAR_ret = test_dmar._get_structure_DMAR(type, b'\x00\x00\x00\x00')
        self.assertIn(f"\n  Unknown DMAR structure 0x{type:02X}\n", structure_DMAR_ret)

    def test_verify_apic_table_format(self):
        test_apic = APIC()
        value_re = r'^[<=]?[BHQsI\d]*'
        key_re = r'^\w+'
        self.assertEqual(type(test_apic.APIC_TABLE_FORMAT), dict)
        for key, value in test_apic.APIC_TABLE_FORMAT.items():
            self.assertRegex(key, key_re)
            self.assertRegex(value, value_re)

    def test_apic_parse(self):
        test_table_content = b'\x00\x00\xe0\xfe\x01\x00\x00\x00\x00\x08\x01\x00\x01\x00\x00\x00\x04\x06\x01\x05\x00\x01\x00\x08\x02\x02\x01\x00\x00\x00\x04\x06\x02\x05\x00\x01\x00\x08\x03\x01\x01\x00\x00\x00\x04\x06\x03\x05\x00\x01\x00\x08\x04\x03\x01\x00\x00\x00\x04\x06\x04\x05\x00\x01\x01\x0c\x02\x00\x00\x00\xc0\xfe\x00\x00\x00\x00\x02\n\x00\x00\x02\x00\x00\x00\x00\x00\x02\n\x00\t\t\x00\x00\x00\r\x00'
        test_apic = APIC()
        test_apic.parse(test_table_content)
        self.assertEqual(len(test_apic.apic_structs), 11)

    def test_xsdt_parse(self):
        test_table_content = b'h\x9daz\x00\x00\x00\x00\x80\x9eaz\x00\x00\x00\x00\x08\x9faz\x00\x00\x00\x00P\x9faz\x00\x00\x00\x00\xf0\x9faz\x00\x00\x00\x000\xa0az\x00\x00\x00\x00\x90\xa3az\x00\x00\x00\x00\xf8\xd4az\x00\x00\x00\x000\xd5az\x00\x00\x00\x00\x18\xdfaz\x00\x00\x00\x00`\xdfaz\x00\x00\x00\x00\xc0\xe9az\x00\x00\x00\x00p\x01bz\x00\x00\x00\x00\x08\x02bz\x00\x00\x00\x00P\x03bz\x00\x00\x00\x00\xf0\x05bz\x00\x00\x00\x00\xf85bz\x00\x00\x00\x00\xd86bz\x00\x00\x00\x00\xc89bz\x00\x00\x00\x00\x00:bz\x00\x00\x00\x00X:bz\x00\x00\x00\x00@Mbz\x00\x00\x00\x00\xe8Mbz\x00\x00\x00\x00\x18Nbz\x00\x00\x00\x00PNbz\x00\x00\x00\x00\x88Nbz\x00\x00\x00\x00\xe0Nbz\x00\x00\x00\x00'
        test_xsdt = XSDT()
        test_xsdt.parse(test_table_content)
        self.assertEqual(len(test_xsdt.Entries), 27)

    def test_bgrt_parse(self):
        test_table_content = b"\x01\x00\x01\x00\x18P\x9fv\x00\x00\x00\x00\xf3\x02\x00\x00'\x01\x00\x00"
        test_bgrt = BGRT()
        test_bgrt.parse(test_table_content)
        self.assertEqual(test_bgrt.Version, 1)
        self.assertEqual(test_bgrt.Status, 1)
        self.assertEqual(test_bgrt.ImageOffsetY, 295)

    def test_uefi_parse(self):
        test_table_content = b'\xe2\xd8\x8e\xc6\xc6\x9d\xbdL\x9d\x94\xdbe\xac\xc5\xc328\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\xb0\x9fz\x00\x00\x00\x00'
        test_uefi = UEFI_TABLE()
        test_uefi.parse(test_table_content)
        self.assertEqual(test_uefi.smi, 1)
        self.assertEqual(test_uefi.buf_addr, 0x7a9fb00000000000)

    def test_wsmt_parse(self):
        test_table_content = b'\x00\x00\x00\x00'
        test_wsmt = WSMT()
        test_wsmt.parse(test_table_content)
        self.assertEqual(test_wsmt.fixed_comm_buffers, 0)
        self.assertEqual(test_wsmt.comm_buffer_nested_ptr_protection, 0)
        self.assertEqual(test_wsmt.system_resource_protection, 0)
    # def test_get_structure_dmar(self):
    #     test_table_content = b'&\x05\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00\x00\xd9\xfe\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00`\xd8\xfe\x00\x00\x00\x00\x02\x08\x00\x00\x00\x00\x07\x00\x00\x00\x18\x00\x00\x00\x00\x00\x00p\xd8\xfe\x00\x00\x00\x00\x02\x08\x00\x00\x00\x00\x07\x03\x00\x00 \x00\x01\x00\x00\x00\x00\x10\xd9\xfe\x00\x00\x00\x00\x03\x08\x00\x00\x02\x00\x1e\x07\x04\x08\x00\x00\x00\x00\x1e\x06\x01\x00 \x00\x00\x00\x00\x00\x00\x00\x00/\x00\x00\x00\x00\xff\xff\x7fO\x00\x00\x00\x00\x01\x08\x00\x00\x00\x00\x02\x00'
    #     test_dmar = DMAR()
    #     i = 1
    #     for key, value in test_dmar.DMAR_TABLE_FORMAT.items():
    #         size = struct.calcsize(value)
    #         test_dmar._get_structure_DMAR(i, test_table_content[])
    #         i += 1
    #     print(test_dmar._get_DMAR_structure_SATC(b'\x00'))


if __name__ == '__main__':
    unittest.main()
