# -*- coding: utf-8 -*-
from binascii import unhexlify
import unittest

import pefile

PE_32 = data = bytearray(
    unhexlify(
        (
            "4D5A90000300000004000000FFFF0000B800000000000000400000000000000000000000000000000000000000000000000000000000000000000000"
            "B00000000E1FBA0E00B409CD21B8014CCD21546869732070726F6772616D2063616E6E6F742062652072756E20696E20444F53206D6F64652E0D0D0A"
            "2400000000000000AD723BDFE913558CE913558CE913558C1B64558DE813558C1B64578DE813558C52696368E913558C000000000000000050450000"
            "4C0101009DE473630000000000000000E00002210B010E1E00000000800200000000000000000000D0010000D0010000000000101000000010000000"
            "0600000000000000060000000000000050040000D00100000000000002004005000010000010000000001000001000000000000010000000F0010000"
            "E801000000000000000000000000000000000000000000000000000000000000000000000000000000000000D00100001C0000000000000000000000"
            "000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
            "000000002E7264617461000074020000D001000080020000D001000000000000000000000000000040000040000000009DE47363000000000D000000"
            "54000000F0030000F00300000000000000000000FFFFFFFF000000003A030000010000001D0000001D000000180200008C0200000003000000000000"
            "000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
            "00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000450300004A030000"
            "4F03000054030000590300005E03000063030000680300006D03000072030000770300007C03000081030000860300008B0300009003000095030000"
            "9A0300009F030000A4030000A9030000AE030000B3030000B8030000BD030000C2030000C7030000CC030000D1030000000001000200030004000500"
            "06000700080009000A000B000C000D000E000F0010001100120013001400150016001700180019001A001B001C007465737433322E646C6C00545354"
            "2100545354220054535423005453542400545354250054535426005453542700545354280054535429005453542A005453542B005453542C00545354"
            "2D005453542E005453542F005453543A005453543C005453543E005453543F005453545B005453545C005453545D005453545E005453545F00545354"
            "60005453547B005453547C005453547D005453547E00000018000000008000800000000000000000000000000000000000000000D001000020000000"
            "2E72646174610000F0010000E80100002E65646174610000D8030000180000002E726461746124766F6C746D64000000F0030000540000002E726461"
            "7461247A7A7A646267000000000000000000000000000000"
        ).encode()
    )
)

PE_64 = data = bytearray(
    unhexlify(
        (
            "4D5A90000300000004000000FFFF0000B800000000000000400000000000000000000000000000000000000000000000000000000000000000000000"
            "B00000000E1FBA0E00B409CD21B8014CCD21546869732070726F6772616D2063616E6E6F742062652072756E20696E20444F53206D6F64652E0D0D0A"
            "2400000000000000AD723BDFE913558CE913558CE913558C1B64558DE813558C1B64578DE813558C52696368E913558C000000000000000050450000"
            "648601009DE473630000000000000000F00022200B020E1E00000000800200000000000000000000E001000000000080010000001000000010000000"
            "0600000000000000060000000000000060040000E0010000000000000200600100001000000000000010000000000000000010000000000000100000"
            "00000000000000001000000000020000E801000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
            "E00100001C00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
            "00000000000000000000000000000000000000002E7264617461000074020000E001000080020000E001000000000000000000000000000040000040"
            "000000009DE47363000000000D0000005400000000040000000400000000000000000000FFFFFFFF000000004A030000010000001D0000001D000000"
            "280200009C02000010030000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
            "000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
            "0000000000000000550300005A0300005F03000064030000690300006E03000073030000780300007D03000082030000870300008C03000091030000"
            "960300009B030000A0030000A5030000AA030000AF030000B4030000B9030000BE030000C3030000C8030000CD030000D2030000D7030000DC030000"
            "E103000000000100020003000400050006000700080009000A000B000C000D000E000F0010001100120013001400150016001700180019001A001B00"
            "1C007465737436342E646C6C005453542100545354220054535423005453542400545354250054535426005453542700545354280054535429005453"
            "542A005453542B005453542C005453542D005453542E005453542F005453543A005453543C005453543E005453543F005453545B005453545C005453"
            "545D005453545E005453545F0054535460005453547B005453547C005453547D005453547E0000001800000000800080000000000000000000000000"
            "0000000000000000E0010000200000002E7264617461000000020000E80100002E65646174610000E8030000180000002E726461746124766F6C746D"
            "6400000000040000540000002E7264617461247A7A7A646267000000000000000000000000000000"
        ).encode()
    )
)


EXPECT_NAMES = [
    (1, "TST!"),
    (2, 'TST"'),
    (3, "TST#"),
    (4, "TST$"),
    (5, "TST%"),
    (6, "TST&"),
    (7, "TST'"),
    (8, "TST("),
    (9, "TST)"),
    (10, "TST*"),
    (11, "TST+"),
    (12, "TST,"),
    (13, "TST-"),
    (14, "TST."),
    (15, "TST/"),
    (16, "TST:"),
    (17, "TST<"),
    (18, "TST>"),
    (19, "TST?"),
    (20, "TST["),
    (21, "TST\\"),
    (22, "TST]"),
    (23, "TST^"),
    (24, "TST_"),
    (25, "TST`"),
    (26, "TST{"),
    (27, "TST|"),
    (28, "TST}"),
    (29, "TST~"),
]


class Test_exports(unittest.TestCase):
    def test_exports32(self):
        """Iterate all exports of the 32 bit binary, and validate that they are all accepted"""
        # Update all the symbol addresses to be __ImageBase+0x01
        for idx, _ in enumerate(EXPECT_NAMES):
            PE_32[0x218 + 4 * idx] = 0x01
        pe = pefile.PE(data=PE_32)

        exports = []
        for exp in pe.DIRECTORY_ENTRY_EXPORT.symbols:
            name = exp.name.decode("utf-8") if exp.name else "-INVALID-"
            exports.append((exp.ordinal, name))

        self.assertEqual(EXPECT_NAMES, exports)

    def test_exports64(self):
        """Iterate all exports of the 64 bit binary, and validate that they are all accepted"""
        # Update all the symbol addresses to be __ImageBase+0x01
        for idx, _ in enumerate(EXPECT_NAMES):
            PE_64[0x228 + 4 * idx] = 0x01
        pe = pefile.PE(data=PE_64)

        exports = []
        for exp in pe.DIRECTORY_ENTRY_EXPORT.symbols:
            name = exp.name.decode("utf-8") if exp.name else "-INVALID-"
            exports.append((exp.ordinal, name))

        self.assertEqual(EXPECT_NAMES, exports)


"""
To generate the embedded files, run from a VS command prompt:
link /NOLOGO /SUBSYSTEM:WINDOWS /NOIMPLIB /NOEXP /ALIGN:0x10 /DEF:test.def /DLL /NOENTRY /MACHINE:X86 /OUT:test32.dll
link /NOLOGO /SUBSYSTEM:WINDOWS /NOIMPLIB /NOEXP /ALIGN:0x10 /DEF:test.def /DLL /NOENTRY /MACHINE:X64 /OUT:test64.dll

test.def:
EXPORTS
  TST!=__ImageBase
  TST"=__ImageBase
  TST#=__ImageBase
  TST$=__ImageBase
  TST%=__ImageBase
  TST&=__ImageBase
  TST'=__ImageBase
  TST(=__ImageBase
  TST)=__ImageBase
  TST*=__ImageBase
  TST+=__ImageBase
  TST,=__ImageBase
  TST-=__ImageBase
  TST.=__ImageBase
  TST/=__ImageBase
  TST:=__ImageBase
  TST<=__ImageBase
  TST>=__ImageBase
  TST?=__ImageBase
  TST[=__ImageBase
  TST\=__ImageBase
  TST]=__ImageBase
  TST^=__ImageBase
  TST_=__ImageBase
  TST`=__ImageBase
  TST{=__ImageBase
  TST|=__ImageBase
  TST}=__ImageBase
  TST~=__ImageBase

"""