mem_edit/mem_edit/utils.py
Jan Petykiewicz e56ec88761 modernize some code style
indentation, type annotations, f-strings
2024-03-30 17:41:18 -07:00

106 lines
2.7 KiB
Python

"""
Utility functions and types:
Type definition for buffers:
ctypes_buffer_t
Custom exception type:
MemEditError
Search for a buffer inside another buffer:
search_buffer(needle_buffer, haystack_buffer)
Check if two buffers (ctypes objects) store equal values:
ctypes_equal(a, b)
"""
from typing import Union
import ctypes
ctypes_buffer_t = Union[
ctypes._SimpleCData,
ctypes.Array,
ctypes.Structure,
ctypes.Union,
]
class MemEditError(Exception):
pass
def search_buffer_verbatim(
needle_buffer: ctypes_buffer_t,
haystack_buffer: ctypes_buffer_t,
) -> list[int]:
"""
Search for a buffer inside another buffer, using a direct (bitwise) comparison
Args:
needle_buffer: Buffer to search for.
haystack_buffer: Buffer to search in.
Returns:
List of offsets where the `needle_buffer` was found.
"""
found = []
haystack = bytes(haystack_buffer)
needle = bytes(needle_buffer)
start = 0
result = haystack.find(needle, start)
while start < len(haystack) and result != -1:
found.append(result)
start = result + 1
result = haystack.find(needle, start)
return found
def search_buffer(
needle_buffer: ctypes_buffer_t,
haystack_buffer: ctypes_buffer_t,
) -> list[int]:
"""
Search for a buffer inside another buffer, using `ctypes_equal` for comparison.
Much slower than `search_buffer_verbatim`.
Args:
needle_buffer: Buffer to search for.
haystack_buffer: Buffer to search in.
Returns:
List of offsets where the needle_buffer was found.
"""
found = []
read_type = type(needle_buffer)
for offset in range(0, len(haystack_buffer) - ctypes.sizeof(needle_buffer)):
v = read_type.from_buffer(haystack_buffer, offset)
if ctypes_equal(needle_buffer, v):
found.append(offset)
return found
def ctypes_equal(
a: ctypes_buffer_t,
b: ctypes_buffer_t,
) -> bool:
"""
Check if the values stored inside two ctypes buffers are equal.
"""
if not type(a) == type(b):
return False
if isinstance(a, ctypes.Array):
return a[:] == b[:]
elif isinstance(a, ctypes.Structure) or isinstance(a, ctypes.Union):
for attr_name, attr_type in a._fields_:
a_attr, b_attr = (getattr(x, attr_name) for x in (a, b))
if isinstance(a, (ctypes.Array, ctypes.Structure, ctypes.Union, ctypes._SimpleCData)):
if not ctypes_equal(a_attr, b_attr):
return False
elif not a_attr == b_attr:
return False
return True
else:
return a.value == b.value