import gdb

import pwndbg
import tests

MEMORY_BINARY = tests.binaries.get('memory.out')
data_addr = '0x400081'


def test_windbg_dX_commands(start_binary):
    """
    Tests windbg compatibility commands that dump memory
    like dq, dw, db, ds etc.
    """
    start_binary(MEMORY_BINARY)

    # Try to fail commands in different way
    for cmd_prefix in ('dq', 'dd', 'dw', 'db'):

        # With a non-existent symbol
        cmd = cmd_prefix + ' nonexistentsymbol'
        assert gdb.execute(cmd, to_string=True) == (
            'usage: XX [-h] address [count]\n'
            "XX: error: argument address: invalid HexOrAddressExpr value: 'nonexistentsymbol'\n"
        ).replace('XX', cmd_prefix)

        # With an invalid/unmapped address
        cmd = cmd_prefix + ' 0'
        assert gdb.execute(cmd, to_string=True) == 'Could not access the provided address\n'

    #################################################
    #### dq command tests
    #################################################
    # Try `dq` with symbol, &symbol, 0x<address> and <address> without 0x prefix (treated as hex!)
    dq1 = gdb.execute('dq data', to_string=True)
    dq2 = gdb.execute('dq &data', to_string=True)
    dq3 = gdb.execute('dq %s' % data_addr, to_string=True)
    dq4 = gdb.execute('dq %s' % data_addr.replace('0x', ''), to_string=True)
    assert dq1 == dq2 == dq3 == dq4 == (
        '0000000000400081     0000000000000000 0000000000000001\n'
        '0000000000400091     0000000100000002 0001000200030004\n'
        '00000000004000a1     0102030405060708 1122334455667788\n'
        '00000000004000b1     0123456789abcdef 0000000000000000\n'
    )
    
    # Try `dq` with different counts
    dq_count1 = gdb.execute('dq data 2', to_string=True)
    dq_count2 = gdb.execute('dq &data 2', to_string=True)
    dq_count3 = gdb.execute('dq %s 2' % data_addr, to_string=True)
    assert dq_count1 == dq_count2 == dq_count3 == '0000000000400081     0000000000000000 0000000000000001\n'
    
    assert gdb.execute('dq data 1', to_string=True) == '0000000000400081     0000000000000000\n'
    assert gdb.execute('dq data 3', to_string=True) == (
        '0000000000400081     0000000000000000 0000000000000001\n'
        '0000000000400091     0000000100000002\n'
    )

    # Try 'dq' with count equal to a register, but lets set it before ;)
    # also note that we use `data2` here
    assert gdb.execute('set $eax=4', to_string=True) == ''  # assert as a sanity check
    assert gdb.execute('dq data2 $eax', to_string=True) == (
        '00000000004000a9     1122334455667788 0123456789abcdef\n'
        '00000000004000b9     0000000000000000 ffffffffffffffff\n'
    )

    # See if we can repeat dq command (use count for shorter data)
    assert gdb.execute('dq data2 2', to_string=True) == (
        '00000000004000a9     1122334455667788 0123456789abcdef\n'
    )

    # TODO/FIXME: Can we test command repeating here? Neither passing `from_tty=True`
    # or setting `pwndbg.commands.windbg.dq.repeat = True` works here
    #assert gdb.execute('dq data2 2', to_string=True, from_tty=True) == (
    #    '00000000004000b9     0000000000000000 ffffffffffffffff\n'
    #)

    #################################################
    #### dd command tests
    #################################################
    dd1 = gdb.execute('dd data', to_string=True)
    dd2 = gdb.execute('dd &data', to_string=True)
    dd3 = gdb.execute('dd %s' % data_addr, to_string=True)
    dd4 = gdb.execute('dd %s' % data_addr.replace('0x', ''), to_string=True)
    assert dd1 == dd2 == dd3 == dd4 == (
        '0000000000400081     00000000 00000000 00000001 00000000\n'
        '0000000000400091     00000002 00000001 00030004 00010002\n'
        '00000000004000a1     05060708 01020304 55667788 11223344\n'
        '00000000004000b1     89abcdef 01234567 00000000 00000000\n'
    )

    # count tests
    assert gdb.execute('dd data 4', to_string=True) == (
        '0000000000400081     00000000 00000000 00000001 00000000\n'
    )
    assert gdb.execute('dd data 3', to_string=True) == (
        '0000000000400081     00000000 00000000 00000001\n'
    )

    #################################################
    #### dw command tests
    #################################################
    dw1 = gdb.execute('dw data', to_string=True)
    dw2 = gdb.execute('dw &data', to_string=True)
    dw3 = gdb.execute('dw %s' % data_addr, to_string=True)
    dw4 = gdb.execute('dw %s' % data_addr.replace('0x', ''), to_string=True)
    assert dw1 == dw2 == dw3 == dw4 == (
        '0000000000400081     0000 0000 0000 0000 0001 0000 0000 0000\n'
        '0000000000400091     0002 0000 0001 0000 0004 0003 0002 0001\n'
        '00000000004000a1     0708 0506 0304 0102 7788 5566 3344 1122\n'
        '00000000004000b1     cdef 89ab 4567 0123 0000 0000 0000 0000\n'
    )

    # count tests
    assert gdb.execute('dw data 8', to_string=True) == (
        '0000000000400081     0000 0000 0000 0000 0001 0000 0000 0000\n'
    )

    assert gdb.execute('dw data 8/2', to_string=True) == (
        '0000000000400081     0000 0000 0000 0000\n'
    )

    assert gdb.execute('dw data $eax', to_string=True) == (
        '0000000000400081     0000 0000 0000 0000\n'
    )

    #################################################
    #### db command tests
    #################################################
    db1 = gdb.execute('db data', to_string=True)
    db2 = gdb.execute('db &data', to_string=True)
    db3 = gdb.execute('db %s' % data_addr, to_string=True)
    db4 = gdb.execute('db %s' % data_addr.replace('0x', ''), to_string=True)
    assert db1 == db2 == db3 == db4 == (
        '0000000000400081     00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00\n'
        '0000000000400091     02 00 00 00 01 00 00 00 04 00 03 00 02 00 01 00\n'
        '00000000004000a1     08 07 06 05 04 03 02 01 88 77 66 55 44 33 22 11\n'
        '00000000004000b1     ef cd ab 89 67 45 23 01 00 00 00 00 00 00 00 00\n'
    )

    # count tests
    assert gdb.execute('db data 31', to_string=True) == (
        '0000000000400081     00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00\n'
        '0000000000400091     02 00 00 00 01 00 00 00 04 00 03 00 02 00 01\n'
    )
    assert gdb.execute('db data $ax', to_string=True) == (
        '0000000000400081     00 00 00 00\n'
    )

    #################################################
    #### dc command tests
    #################################################
    dc1 = gdb.execute('dc data', to_string=True)
    dc2 = gdb.execute('dc &data', to_string=True)
    dc3 = gdb.execute('dc %s' % data_addr, to_string=True)
    dc4 = gdb.execute('dc %s' % data_addr.replace('0x', ''), to_string=True)
    assert dc1 == dc2 == dc3 == dc4 == (
        '+0000 0x400081  00 00 00 00  00 00 00 00                            '
        '│....│....│    │    │\n'
    )

    assert gdb.execute('dc data 3', to_string=True) == (
        '+0000 0x400081  00 00 00                                            │... '
        '│    │    │    │\n'
    )

    #################################################
    #### ds command tests
    #################################################
    ds1 = gdb.execute('ds short_str', to_string=True)
    ds2 = gdb.execute('ds &short_str', to_string=True)
    ds3 = gdb.execute('ds 0x4000d9', to_string=True)
    ds4 = gdb.execute('ds 4000d9', to_string=True)
    assert ds1 == ds2 == ds3 == ds4 == "4000d9 'some cstring here'\n"

    # Check too low maxlen
    assert gdb.execute('ds short_str 5', to_string=True) == (
        'Max str len of 5 too low, changing to 256\n'
        "4000d9 'some cstring here'\n"
    )

    # Check output for a string longer than (the default) maxlen of 256
    assert gdb.execute('ds long_str', to_string=True) == (
        "4000eb 'long string: "
        "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA...'\n"
    )

    # Check impossible address
    assert gdb.execute('ds 0', to_string=True) == (
        "Data at address can't be dereferenced or is not a printable null-terminated "
        'string or is too short.\n'
        'Perhaps try: db <address> <count> or hexdump <address>\n'
    )


def test_windbg_eX_commands(start_binary):
    """
    Tests windbg compatibility commands that write to memory
    like eq, ed, ew, eb etc.
    """
    start_binary(MEMORY_BINARY)

    # Try to fail commands in different way
    for cmd_prefix in ('eq', 'ed', 'ew', 'eb'):
        # With a non-existent symbol
        cmd = cmd_prefix + ' nonexistentsymbol'

        # Seems there is some mismatch between Python 3.x argparse output
        expected_in = (
            # This version occurred locally when tested on Python 3.9.5
            (
                'usage: XX [-h] address [data ...]\n'
                "XX: error: argument address: invalid HexOrAddressExpr value: 'nonexistentsymbol'\n"
            ).replace('XX', cmd_prefix),
            # This version occurs on CI on Python 3.8.10
            (
                'usage: XX [-h] address [data [data ...]]\n'
                "XX: error: argument address: invalid HexOrAddressExpr value: 'nonexistentsymbol'\n"
            ).replace('XX', cmd_prefix),
        )

        assert gdb.execute(cmd, to_string=True) in expected_in
        assert gdb.execute(cmd, to_string=True) in expected_in

        # With no data arguments provided
        cmd = cmd_prefix + ' 0'
        assert gdb.execute(cmd, to_string=True) == 'Cannot write empty data into memory.\n'

        # With invalid/unmapped address 0
        cmd = cmd_prefix + ' 0 1122'
        assert gdb.execute(cmd, to_string=True) == (
            'Cannot access memory at address 0x0\n'
        )

        # With invalid data which can't be parsed as hex
        cmd = cmd_prefix + ' 0 x'
        assert gdb.execute(cmd, to_string=True) == (
            'Incorrect data format: it must all be a hex value (0x1234 or 1234, both '
            'interpreted as 0x1234)\n'
        )
    #########################################
    ### Test eq write
    #########################################
    assert gdb.execute('eq $sp 0xcafebabe', to_string=True) == ''
    assert '0x00000000cafebabe' in gdb.execute('x/xg $sp', to_string=True)

    assert gdb.execute('eq $sp 0xbabe 0xcafe', to_string=True) == ''
    assert '0x000000000000babe\t0x000000000000cafe' in gdb.execute('x/2xg $sp', to_string=True)

    assert gdb.execute('eq $sp cafe000000000000 babe000000000000', to_string=True) == ''
    assert '0xcafe000000000000\t0xbabe000000000000' in gdb.execute('x/2xg $sp', to_string=True)

    # TODO/FIXME: implement tests for others (ed, ew, eb etc)

    #########################################
    ### Test write & output on partial write
    #########################################
    # e.g. when we make a write to the last stack address
    stack_ea = pwndbg.regs[pwndbg.regs.stack]
    stack_page = pwndbg.vmmap.find(stack_ea)

    # Last possible address on stack where we can perform an 8-byte write
    stack_last_qword_ea = stack_page.end - 8

    gdb_result = gdb.execute('eq %#x 0xCAFEBABEdeadbeef 0xABCD' % stack_last_qword_ea, to_string=True).split("\n")
    assert "Cannot access memory at address" in gdb_result[0]
    assert gdb_result[1] == '(Made 1 writes to memory; skipping further writes)'

    # Check if the write actually occurred
    assert pwndbg.memory.read(stack_last_qword_ea, 8) == b'\xef\xbe\xad\xde\xbe\xba\xfe\xca'