From 96329fae4f6d3f642d243d3c39125190696f883c Mon Sep 17 00:00:00 2001 From: Ingmar Steen Date: Sat, 21 Mar 2015 14:03:24 +0100 Subject: [PATCH] Add nosetests based test suite. Covers everything except for flow, which is due for refactoring. --- tests/__init__.py | 5 + tests/test_asm.py | 50 +++++++ tests/test_codec.py | 38 +++++ tests/test_elf.py | 323 ++++++++++++++++++++++++++++++++++++++++++ tests/test_packing.py | 131 +++++++++++++++++ tests/test_target.py | 58 ++++++++ tests/test_util.py | 21 +++ 7 files changed, 626 insertions(+) create mode 100644 tests/__init__.py create mode 100644 tests/test_asm.py create mode 100644 tests/test_codec.py create mode 100644 tests/test_elf.py create mode 100644 tests/test_packing.py create mode 100644 tests/test_target.py create mode 100644 tests/test_util.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..27369d0 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,5 @@ +import pwny + + +def setup(): + pwny.target.assume(pwny.Target(arch=pwny.Architecture.x86)) diff --git a/tests/test_asm.py b/tests/test_asm.py new file mode 100644 index 0000000..ea0da19 --- /dev/null +++ b/tests/test_asm.py @@ -0,0 +1,50 @@ +from nose.tools import raises +import pwny + + +SOURCE = 'mov al, [0xced]' +RESULT_BIN_32 = b'\xa0\xed\x0c\x00\x00' +RESULT_BIN_64 = b'\x8a\x04%\xed\x0c\x00\x00' +RESULT_ITH_32 = b':05000000A0ED0C000062\n:00000001FF\n' +RESULT_ITH_64 = b':070000008A0425ED0C00004D\n:00000001FF\n' + + +def test_asm(): + if pwny.target.bits == 32: + assert pwny.asm(SOURCE) == RESULT_BIN_32 + else: + assert pwny.asm(SOURCE) == RESULT_BIN_64 + + +def test_asm_with_bits(): + assert pwny.asm(SOURCE, bits=32) == RESULT_BIN_32 + assert pwny.asm(SOURCE, bits=64) == RESULT_BIN_64 + + +def test_asm_with_target(): + target_32 = pwny.Target(arch=pwny.Architecture.x86) + target_64 = pwny.Target(arch=pwny.Architecture.x86_64) + assert pwny.asm(SOURCE, target=target_32) == RESULT_BIN_32 + assert pwny.asm(SOURCE, target=target_64) == RESULT_BIN_64 + + +def test_asm_with_bits_and_target(): + target_32 = pwny.Target(arch=pwny.Architecture.x86) + target_64 = pwny.Target(arch=pwny.Architecture.x86_64) + assert pwny.asm(SOURCE, bits=32, target=target_64) == RESULT_BIN_32 + assert pwny.asm(SOURCE, bits=64, target=target_32) == RESULT_BIN_64 + + +def test_asm_with_format(): + assert pwny.asm(SOURCE, fmt=pwny.asm.Format.ith, bits=32) == RESULT_ITH_32 + assert pwny.asm(SOURCE, fmt=pwny.asm.Format.ith, bits=64) == RESULT_ITH_64 + + +@raises(ValueError) +def test_asm_invalid_format(): + pwny.asm(SOURCE, fmt='ced') + + +@raises(SyntaxError) +def test_asm_syntax_error(): + pwny.asm('mov ced, 3') diff --git a/tests/test_codec.py b/tests/test_codec.py new file mode 100644 index 0000000..ef17541 --- /dev/null +++ b/tests/test_codec.py @@ -0,0 +1,38 @@ +import pwny + + +def test_xor_int(): + assert pwny.xor(61, b'fooo') == b'[RRR' + + +def test_xor_str(): + assert pwny.xor(b'abcd', b'fooo') == b'\x07\r\x0c\x0b' + assert pwny.xor(b'abcd', b'fooofooo') == b'\x07\r\x0c\x0b\x07\r\x0c\x0b' + + +def test_rot13(): + assert pwny.rot13('whax') == 'junk' + + +def test_caesar(): + assert pwny.caesar(1, 'abcXYZ') == 'bcdYZA' + + +def test_enhex(): + assert pwny.enhex(b'ABCD') == '41424344' + + +def test_dehex(): + assert pwny.dehex('41424344') == b'ABCD' + + +def test_enb64(): + assert pwny.enb64(b'ABCD') == 'QUJDRA==' + + +def test_deb64(): + assert pwny.deb64('QUJDRA==') == b'ABCD' + + +def test_frequency(): + assert pwny.frequency('ABCD') == {'A': 1, 'B': 1, 'C': 1, 'D': 1} diff --git a/tests/test_elf.py b/tests/test_elf.py new file mode 100644 index 0000000..79f42a5 --- /dev/null +++ b/tests/test_elf.py @@ -0,0 +1,323 @@ +import six +import mock +import pwny + + +headers = [ + { + 'data': b'\x7fELF\x02\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00>\x00\x01\x00\x00\x00@\x04@\x00\x00' + b'\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00p\x11\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00@\x008\x00\t' + b'\x00@\x00\x1e\x00\x1b\x00', + 'header': { + 'arch': pwny.Architecture.x86_64, + 'bits': 64, + 'endian': pwny.Endianness.little, + 'abi_version': 0, + 'entry': 4195392, + 'flags': 0, + 'hsize': 64, + 'osabi': pwny.ELF.OSABI.system_v, + 'phentsize': 56, + 'phnum': 9, + 'phoff': 64, + 'shentsize': 64, + 'shnum': 30, + 'shoff': 4464, + 'shstrndx': 27, + 'type': pwny.ELF.Type.executable + }, + }, + { + 'data': b'\x7fELF\x02\x01\x01\x03\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00>\x00\x01\x00\x00\x00N\x0f@\x00\x00' + b'\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00\xa8\x1d\x0c\x00\x00\x00\x00\x00\x00\x00\x00\x00@\x008\x00' + b'\x06\x00@\x00\x1f\x00\x1c\x00', + 'header': { + 'arch': pwny.Architecture.x86_64, + 'bits': 64, + 'endian': pwny.Endianness.little, + 'abi_version': 0, + 'entry': 4198222, + 'flags': 0, + 'hsize': 64, + 'osabi': pwny.ELF.OSABI.linux, + 'phentsize': 56, + 'phnum': 6, + 'phoff': 64, + 'shentsize': 64, + 'shnum': 31, + 'shoff': 794024, + 'shstrndx': 28, + 'type': pwny.ELF.Type.executable, + }, + }, + { + 'data': b'\x7fELF\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x03\x00\x01\x00\x00\x00 \x83\x04\x08' + b'4\x00\x00\x00L\x11\x00\x00\x00\x00\x00\x004\x00 \x00\t\x00(\x00\x1e\x00\x1b\x00\x06\x00\x00\x004\x00' + b'\x00\x004\x80\x04\x08', + 'header': { + 'arch': pwny.Architecture.x86, + 'bits': 32, + 'endian': pwny.Endianness.little, + 'abi_version': 0, + 'entry': 134513440, + 'flags': 0, + 'hsize': 52, + 'osabi': pwny.ELF.OSABI.system_v, + 'phentsize': 32, + 'phnum': 9, + 'phoff': 52, + 'shentsize': 40, + 'shnum': 30, + 'shoff': 4428, + 'shstrndx': 27, + 'type': pwny.ELF.Type.executable, + }, + }, + { + 'data': b'\x7fELF\x01\x01\x01\x03\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x03\x00\x01\x00\x00\x00\n\x8d\x04\x08' + b'4\x00\x00\x00\xf0 \n\x00\x00\x00\x00\x004\x00 \x00\x06\x00(\x00\x1f\x00\x1c\x00\x01\x00\x00\x00\x00' + b'\x00\x00\x00\x00\x80\x04\x08', + 'header': { + 'arch': pwny.Architecture.x86, + 'bits': 32, + 'endian': pwny.Endianness.little, + 'abi_version': 0, + 'entry': 134515978, + 'flags': 0, + 'hsize': 52, + 'osabi': pwny.ELF.OSABI.linux, + 'phentsize': 32, + 'phnum': 6, + 'phoff': 52, + 'sections': [], + 'shentsize': 40, + 'shnum': 31, + 'shoff': 663792, + 'shstrndx': 28, + 'strings': None, + 'type': pwny.ELF.Type.executable, + }, + } +] + + +def test_elf_header_parse(): + for header in headers: + yield check_elf_header_parse, header['data'], header['header'] + + +def check_elf_header_parse(data, values): + elf = pwny.ELF(data) + for key, value in values.items(): + assert getattr(elf, key, value) == value + + +def test_elf_parse_section_invalid_type(): + section_fmt = 'IIIIIIIIII' + section_fmt_size = pwny.pack_size(section_fmt) + section = pwny.pack( + section_fmt, + 1, + 0x5fffffff, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + section_fmt_size, + ) + elf = pwny.ELF() + elf.bits = 32 + section = elf.parse_section_header(section) + assert section['type'] == pwny.ELF.SectionType.UNKNOWN + + +def test_elf_parse_file_32(): + b = six.BytesIO() + + header_fmt = '4sBBBBB7sHHIIIIIHHHHHH' + header_fmt_size = pwny.pack_size(header_fmt) + + section_fmt = 'IIIIIIIIII' + section_fmt_size = pwny.pack_size(section_fmt) + + strings_section = b'\x00strings\x00' + + b.write(pwny.pack( + header_fmt, + b'\x7fELF', + 1, + 1, + 1, + pwny.ELF.OSABI.linux.value, + 0, + b'\x00' * 7, + pwny.ELF.Type.executable.value, + pwny.Architecture.x86.value, + 1, + 0, + 0, + header_fmt_size, + 0, + header_fmt_size, + 0, + 0, + section_fmt_size, + 1, + 0, + )) + + b.write(pwny.pack( + section_fmt, + 1, + pwny.ELF.SectionType.NULL.value, + 0, + 0, + header_fmt_size + section_fmt_size, + len(strings_section), + 0, + 0, + 0, + section_fmt_size, + )) + + b.write(strings_section) + b.seek(0) + + elf = pwny.ELF.parse(b) + + for key, value in { + 'arch': pwny.Architecture.x86, + 'bits': 32, + 'endian': pwny.Endianness.little, + 'abi_version': 0, + 'entry': 0, + 'flags': 0, + 'hsize': header_fmt_size, + 'osabi': pwny.ELF.OSABI.linux, + 'phentsize': 0, + 'phnum': 0, + 'phoff': 0, + 'shentsize': section_fmt_size, + 'shnum': 1, + 'shoff': header_fmt_size, + 'shstrndx': 0, + 'type': pwny.ELF.Type.executable, + }.items(): + assert getattr(elf, key) == value, '%s != %r' % (key, value) + + assert len(elf.sections) == 1 + + +def test_elf_parse_file_64(): + b = six.BytesIO() + + header_fmt = '4sBBBBB7sHHIQQQIHHHHHH' + header_fmt_size = pwny.pack_size(header_fmt) + + section_fmt = 'IIQQQQIIQQ' + section_fmt_size = pwny.pack_size(section_fmt) + + strings_section = b'\x00strings\x00' + + b.write(pwny.pack( + header_fmt, + b'\x7fELF', + 2, + 1, + 1, + pwny.ELF.OSABI.linux.value, + 0, + b'\x00' * 7, + pwny.ELF.Type.executable.value, + pwny.Architecture.x86_64.value, + 1, + 0, + 0, + header_fmt_size, + 0, + header_fmt_size, + 0, + 0, + section_fmt_size, + 1, + 0, + )) + + b.write(pwny.pack( + section_fmt, + 1, + pwny.ELF.SectionType.NULL.value, + 0, + 0, + header_fmt_size + section_fmt_size, + len(strings_section), + 0, + 0, + 0, + section_fmt_size, + )) + + b.write(strings_section) + b.seek(0) + elf = pwny.ELF.parse(b) + + for key, value in { + 'arch': pwny.Architecture.x86_64, + 'bits': 64, + 'endian': pwny.Endianness.little, + 'abi_version': 0, + 'entry': 0, + 'flags': 0, + 'hsize': header_fmt_size, + 'osabi': pwny.ELF.OSABI.linux, + 'phentsize': 0, + 'phnum': 0, + 'phoff': 0, + 'shentsize': section_fmt_size, + 'shnum': 1, + 'shoff': header_fmt_size, + 'shstrndx': 0, + 'type': pwny.ELF.Type.executable, + }.items(): + assert getattr(elf, key) == value, '%s != %r' % (key, value) + + assert len(elf.sections) == 1 + +def test_elf_parse_file_open(): + b = six.BytesIO() + + header_fmt = '4sBBBBB7sHHIIIIIHHHHHH' + header_fmt_size = pwny.pack_size(header_fmt) + + b.write(pwny.pack( + header_fmt, + b'\x7fELF', + 1, + 1, + 1, + pwny.ELF.OSABI.linux.value, + 0, + b'\x00' * 7, + pwny.ELF.Type.executable.value, + pwny.Architecture.x86.value, + 1, + 0, + 0, + header_fmt_size, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + )) + + b.seek(0) + + with mock.patch('pwnypack.elf.open', create=True) as mock_open: + mock_open.return_value = b + elf = pwny.ELF.parse('test.elf') diff --git a/tests/test_packing.py b/tests/test_packing.py new file mode 100644 index 0000000..dab719d --- /dev/null +++ b/tests/test_packing.py @@ -0,0 +1,131 @@ +from nose.tools import raises +import pwny + + +def test_pack(): + assert pwny.pack('I', 0x41424344) == b'DCBA' + + +def test_pack_format_with_endianness(): + assert pwny.pack('>I', 0x41424344) == b'ABCD' + + +def test_pack_explicit_endianness(): + assert pwny.pack('I', 0x41424344, endian=pwny.Endianness.big) == b'ABCD' + + +def test_pack_explicit_target(): + target = pwny.Target(pwny.Architecture.x86, endian=pwny.Endianness.big) + assert pwny.pack('I', 0x41424344, target=target) == b'ABCD' + + +@raises(NotImplementedError) +def test_pack_invalid_endianness(): + pwny.pack('I', 1, endian='invalid') + + +def test_unpack(): + assert pwny.unpack('I', b'DCBA') == (0x41424344,) + + +def test_unpack_format_with_endianness(): + assert pwny.unpack('>I', b'ABCD') == (0x41424344,) + + +def test_unpack_explicit_endianness(): + assert pwny.unpack('I', b'ABCD', endian=pwny.Endianness.big) == (0x41424344,) + + +def test_unpack_explicit_target(): + target = pwny.Target(pwny.Architecture.x86, endian=pwny.Endianness.big) + assert pwny.unpack('I', b'ABCD', target=target) == (0x41424344,) + + +@raises(NotImplementedError) +def test_unpack_invalid_endianness(): + pwny.unpack('I', 'AAAA', endian='invalid') + + +def test_pack_size(): + # This tests both pack_size in general as well as not padding the byte. + assert pwny.pack_size('bq') == 9 + + +short_signed_data = [ + [8, -0x7f, b'\x81'], + [16, -0x7fff, b'\x80\x01'], + [32, -0x7fffffff, b'\x80\x00\x00\x01'], + [64, -0x7fffffffffffffff, b'\x80\x00\x00\x00\x00\x00\x00\x01'], +] + + +short_unsigned_data = [ + [8, 0x61, b'a'], + [16, 0x6162, b'ab'], + [32, 0x61626364, b'abcd'], + [64, 0x6162636465666768, b'abcdefgh'], +] + + +def test_short_form_pack(): + for width, num, bytestr in short_signed_data: + f = 'p%d' % width + yield check_short_form_pack, f, num, bytestr[::-1] + yield check_short_form_pack_endian, f, num, bytestr[::-1], pwny.Endianness.little + yield check_short_form_pack_endian, f, num, bytestr, pwny.Endianness.big + + for width, num, bytestr in short_unsigned_data: + f = 'P%d' % width + yield check_short_form_pack, f, num, bytestr[::-1] + yield check_short_form_pack_endian, f, num, bytestr[::-1], pwny.Endianness.little + yield check_short_form_pack_endian, f, num, bytestr, pwny.Endianness.big + + +def test_short_form_unpack(): + for width, num, bytestr in short_signed_data: + f = 'u%d' % width + yield check_short_form_unpack, f, num, bytestr[::-1] + yield check_short_form_unpack_endian, f, num, bytestr[::-1], pwny.Endianness.little + yield check_short_form_unpack_endian, f, num, bytestr, pwny.Endianness.big + + for width, num, bytestr in short_unsigned_data: + f = 'U%d' % width + yield check_short_form_unpack, f, num, bytestr[::-1] + yield check_short_form_unpack_endian, f, num, bytestr[::-1], pwny.Endianness.little + yield check_short_form_unpack_endian, f, num, bytestr, pwny.Endianness.big + + +def test_pointer_pack(): + yield check_short_form_pack, 'p', -66052, b'\xfc\xfd\xfe\xff' + yield check_short_form_pack_endian, 'p', -66052, b'\xfc\xfd\xfe\xff', pwny.Endianness.little + yield check_short_form_pack_endian, 'p', -66052, b'\xff\xfe\xfd\xfc', pwny.Endianness.big + + yield check_short_form_pack, 'P', 4294901244, b'\xfc\xfd\xfe\xff' + yield check_short_form_pack_endian, 'P', 4294901244, b'\xfc\xfd\xfe\xff', pwny.Endianness.little + yield check_short_form_pack_endian, 'P', 4294901244, b'\xff\xfe\xfd\xfc', pwny.Endianness.big + + +def test_pointer_unpack(): + yield check_short_form_unpack, 'u', -66052, b'\xfc\xfd\xfe\xff' + yield check_short_form_unpack_endian, 'u', -66052, b'\xfc\xfd\xfe\xff', pwny.Endianness.little + yield check_short_form_unpack_endian, 'u', -66052, b'\xff\xfe\xfd\xfc', pwny.Endianness.big + + yield check_short_form_unpack, 'U', 4294901244, b'\xfc\xfd\xfe\xff' + yield check_short_form_unpack_endian, 'U', 4294901244, b'\xfc\xfd\xfe\xff', pwny.Endianness.little + yield check_short_form_unpack_endian, 'U', 4294901244, b'\xff\xfe\xfd\xfc', pwny.Endianness.big + + +def check_short_form_pack(f, num, bytestr): + assert getattr(pwny, f)(num) == bytestr + + +def check_short_form_pack_endian(f, num, bytestr, endian): + assert getattr(pwny, f)(num, endian=endian) == bytestr + + +def check_short_form_unpack(f, num, bytestr): + assert getattr(pwny, f)(bytestr) == num + + +def check_short_form_unpack_endian(f, num, bytestr, endian): + assert getattr(pwny, f)(bytestr, endian=endian) == num diff --git a/tests/test_target.py b/tests/test_target.py new file mode 100644 index 0000000..a310147 --- /dev/null +++ b/tests/test_target.py @@ -0,0 +1,58 @@ +import mock +from nose.tools import raises +import pwny + + +def test_default_arch_32bit(): + with mock.patch('platform.architecture') as platform_mock: + platform_mock.return_value = ('32bit',) + assert pwny.Target().arch is pwny.Architecture.x86 + + +def test_default_arch_64bit(): + with mock.patch('platform.architecture') as platform_mock: + platform_mock.return_value = ('64bit',) + assert pwny.Target().arch is pwny.Architecture.x86_64 + + +def test_set_arch(): + with mock.patch('platform.architecture') as platform_mock: + platform_mock.return_value = ('64bit',) + target = pwny.Target(arch=pwny.Architecture.x86) + assert target.arch is pwny.Architecture.x86 + + +def test_default_endianness(): + assert pwny.Target().endian is pwny.Endianness.little + + +def test_set_endianness(): + target = pwny.Target(endian=pwny.Endianness.big) + assert target.endian is pwny.Endianness.big + + +def test_default_bits_x86(): + target = pwny.Target(arch=pwny.Architecture.x86) + assert target.bits == 32 + + +def test_default_bits_x86_64(): + target = pwny.Target(arch=pwny.Architecture.x86_64) + assert target.bits == 64 + + +@raises(NotImplementedError) +def test_default_bits_unsupported(): + target = pwny.Target(arch=pwny.Architecture.aarch64) + _ = target.bits + + +def test_set_bits(): + target = pwny.Target(bits=33) + assert target.bits == 33 + + +def test_target_assume(): + target = pwny.Target() + target.assume(pwny.Target(arch=pwny.Architecture.aarch64, bits=33, endian=pwny.Endianness.little)) + assert target.arch is pwny.Architecture.aarch64 and target.bits == 33 and target.endian == pwny.Endianness.little diff --git a/tests/test_util.py b/tests/test_util.py new file mode 100644 index 0000000..1ce7c0a --- /dev/null +++ b/tests/test_util.py @@ -0,0 +1,21 @@ +import pwny + + +def test_cycle(): + assert pwny.cycle(64) == 'AAAABAAACAAADAAAEAAAFAAAGAAAHAAAIAAAJAAAKAAALAAAMAAANAAAOAAAPAAA' + + +def test_cycle_width(): + assert pwny.cycle(64, width=2) == 'AABACADAEAFAGAHAIAJAKALAMANAOAPAQARASATAUAVAWAXAYAZBBCBDBEBFBGBH' + + +def test_cycle_find(): + assert pwny.cycle.find('PAAA') == 60 + + +def test_cycle_find_start(): + assert pwny.cycle.find('AAAA') == 0 + + +def test_cycle_find_not_found(): + assert pwny.cycle.find('\x00', width=1) == -1