test_generic.py

import logging

from ctypes import (
    c_uint16, c_uint32, c_int64, c_uint64, c_float, c_double,
    create_string_buffer, c_void_p, addressof,
    c_longlong, c_ulonglong
)

from time import sleep
from os import getpid
from sys import exit, platform, version_info
from multiprocessing import Process as MProcess

from memorpy import (
    MemWorker, Process,
    PROT_READ, PROT_WRITE, PROT_EXEC
)

from memorpy.utils import expected_type_size

if version_info.major > 2:
    MATH_TYPES = (float, int)
else:
    MATH_TYPES = (float, int, long)

VALUE0_DATA = 0xBEEF, 'ushort'
VALUE1_DATA = 0xDEADBEEF, 'uint'
VALUE2_DATA = -0xB15B00B5, 'longlong'
VALUE3_DATA = 0xB15B00B5DEADBEEF, 'ulonglong'
VALUE4_DATA = -31337.0, 'float'
VALUE5_DATA = -1337.899999912341, 'double'

STRING1_TERM = b'Hello bytes world!'
STRING2_TERM = 'Hello unicode world!'

STRING1_DATA = STRING1_TERM * 32, 'bytes'
STRING2_DATA = STRING2_TERM * 32, 'unicode'

VALUES = (
    VALUE0_DATA, VALUE1_DATA, VALUE2_DATA, VALUE3_DATA, VALUE4_DATA,
    VALUE5_DATA, STRING1_DATA, STRING2_DATA
)


def cow():
    STRING1 = create_string_buffer(STRING1_DATA[0])

    STRING2_DATA_ENCODED_UNICODE = STRING2_DATA[0].encode('utf-16le')

    STRING2 = create_string_buffer(
        STRING2_DATA_ENCODED_UNICODE, len(STRING2_DATA_ENCODED_UNICODE))

    VALUE0 = c_uint16(VALUE0_DATA[0])
    VALUE1 = c_uint32(VALUE1_DATA[0])
    VALUE2 = c_longlong(VALUE2_DATA[0])
    VALUE3 = c_ulonglong(VALUE3_DATA[0])
    VALUE4 = c_float(VALUE4_DATA[0])
    VALUE5 = c_double(VALUE5_DATA[0])
    VALUE6 = c_void_p(addressof(VALUE2))

    print('Current pid={} String1: 0x{:X} String2: 0x{:X}'.format(
        getpid(), addressof(STRING1), addressof(STRING2)))

    while True:
        sleep(10)

    STORAGES = (
        VALUE0, VALUE1, VALUE2, VALUE3, VALUE4, VALUE5, VALUE6,
        STRING1, STRING2
    )

    assert(STORAGES)


def setup():
    process = MProcess(target=cow)
    process.start()

    assert(process.pid)
    assert(process.pid != getpid())

    logging.info('Wait a bit, as we are not using locks')
    sleep(5)
    logging.info('Lets go')
    return process


def search_value(worker, value, search_type):
    logging.info('Search constant %s as type %s', value, search_type)
    for item in worker.mem_search(value, search_type, protec=PROT_WRITE):
        found_value = worker[item]

        if found_value != value:
            # Something changed
            logging.info(
                'Found candidate: at %s (read => %s (type=%s))',
                hex(item.address), repr(found_value), type(found_value)
            )
            continue

        logging.info(
            'Found: at %s (read => %s (type=%s))',
            hex(item.address), repr(found_value), type(found_value)
        )

        return item

    logging.error('Failed to find %s type %s', value, search_type)
    return None


def test(process, prefer_lock):
    worker = MemWorker(process.pid, prefer_lock=prefer_lock)

    logging.info('Pid: %d', process.pid)
    logging.info('Iterating exectuable regions')
    with worker:
        for address, size in worker.process.iter_region(protec=PROT_EXEC):
            logging.info(
                'Region: 0x%X size=%d', address, size)

    logging.info('Iterating RW regions')
    with worker:
        for address, size in worker.process.iter_region(
                protec=PROT_WRITE | PROT_READ):
            logging.info(
                'Region: 0x%X size=%d', address, size)

    logging.info('Search for values')

    values_table = {}

    with worker:
        for (value, value_type) in VALUES:
            address = search_value(worker, value, value_type)
            if address is None:
                return False

            if search_value(worker, value, None) is None:
                return False

            values_table[value] = address, value_type

    logging.info('Test address dereference')
    with worker:
        found_2x_deref = False
        for address in worker.mem_search(*VALUE2_DATA):
            for address_of_address in worker.mem_search(address):
                try:
                    deref1 = worker[address_of_address]
                except (IOError, OSError):
                    continue

                deref1.type = VALUE2_DATA[1]
                deref1.size = expected_type_size(deref1.type)

                try:
                    value = worker[deref1]
                except (IOError, OSError):
                    continue

                if value == VALUE2_DATA[0]:
                    found_2x_deref = True
                    break

        if not found_2x_deref:
            return False

    logging.info('Search by regex')
    with worker:
        found_bytes_re = False
        for _, address in worker.mem_search(
                br'Hello\s+([^\s]+)\s+world!', 're'):
            try:
                data = worker[address]
                if data == STRING1_TERM:
                    found_bytes_re = True
                    break
                else:
                    logging.info(
                        'Candidate: %s => %s (!= %s)',
                        address, data, STRING1_TERM
                    )
            except (IOError, OSError):
                continue

        if not found_bytes_re:
            logging.info('Failed to find regex')
            return False

    logging.info('Change values')
    with worker:
        for value, (address, _) in values_table.items():
            if isinstance(value, bytes):
                new_value = b'\xAB' * len(value)
            elif isinstance(value, MATH_TYPES):
                new_value = value + 1

            if address.type == 'unicode':
                continue

            worker[address] = new_value
            current_value = worker[address]

            if current_value != new_value:
                if isinstance(value, bytes):
                    logging.error(
                        'Failed to change value %s to %s: '
                        'Read: %s (address %s)',
                        repr(value), repr(new_value),
                        repr(current_value), address
                    )
                else:
                    logging.error(
                        'Failed to change value 0x%X to 0x%X: '
                        'Read: 0x%X (address %s)',
                        value, new_value, current_value, address
                    )

                return False

            # Set old value back
            worker[address] = value

    return True


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    logging.info('Testing at platform %s', repr(platform))

    current_process_name = None

    for process in Process.list():
        if process.pid == getpid():
            current_process_name = process.name
            logging.info(
                'Current process: pid=%d name=%s',
                process['pid'], process['name']
            )
            break

    assert(current_process_name is not None)

    try:
        current_process = Process(getpid())
    except ValueError:
        current_process = None

    assert(current_process is None)

    py_process = Process.pid_from_name(current_process_name)
    logging.info('Search for any python proces: %s', py_process)

    assert(py_process is not None)

    logging.info('Start child process')

    process = setup()

    try:
        if not test(process, False):
            exit('Failed to pass tests with prefer_lock=False')

        if not test(process, True):
            exit('Failed to pass tests with prefer_lock=True')

        if platform.startswith('linux'):
            from memorpy import LinProcess
            if LinProcess.PROCESS_VM_SUPPORT:
                logging.info('Disabling VM Support')

                LinProcess.PROCESS_VM_SUPPORT = False

                if not test(process, False):
                    exit(
                        'Failed to pass tests with prefer_lock=False '
                        'without process vm'
                    )

                if not test(process, True):
                    exit(
                        'Failed to pass tests with prefer_lock=True'
                        'without process vm'
                    )
            else:
                logging.info('Unable to test with process vm support')

            if LinProcess.LARGE_FILE_SUPPORT:
                logging.info('Disabling Large file Support')

                LinProcess.LARGE_FILE_SUPPORT = False

                if not test(process, False):
                    exit(
                        'Failed to pass tests with prefer_lock=False '
                        'without large file support'
                    )

                if not test(process, True):
                    exit(
                        'Failed to pass tests with prefer_lock=True '
                        'without large file support'
                    )
            else:
                logging.info('Unable to test with large file vm support')

    finally:
        process.terminate()

    logging.info('Completed successfully')