diff --git a/tests/test_elf.py b/tests/test_elf.py index 7c42342a..480bd1ff 100644 --- a/tests/test_elf.py +++ b/tests/test_elf.py @@ -9,10 +9,10 @@ sys.path.append("..") from qiling import Qiling from qiling.const import QL_OS, QL_INTERCEPT, QL_VERBOSE from qiling.exception import * +from qiling.extensions import pipe from qiling.os.const import STRING from qiling.os.posix import syscall from qiling.os.mapper import QlFsMappedObject -from qiling.os.posix.stat import Fstat class ELFTest(unittest.TestCase): @@ -823,41 +823,6 @@ class ELFTest(unittest.TestCase): def test_elf_linux_x86_crackme(self): - class MyPipe(): - def __init__(self): - self.buf = b'' - - def write(self, s): - self.buf += s - - def read(self, l): - if l <= len(self.buf): - ret = self.buf[ : l] - self.buf = self.buf[l : ] - else: - ret = self.buf - self.buf = '' - return ret - - def fileno(self): - return 0 - - def fstat(self): - return Fstat(sys.stdin.fileno()) - - def show(self): - pass - - def clear(self): - pass - - def flush(self): - pass - - def close(self): - self.outpipe.close() - - def instruction_count(ql, address, size, user_data): user_data[0] += 1 @@ -865,14 +830,14 @@ class ELFTest(unittest.TestCase): pass def run_one_round(payload): - stdin = MyPipe() - ql = Qiling(["../examples/rootfs/x86_linux/bin/crackme_linux"], "../examples/rootfs/x86_linux", console = False, stdin = stdin) + mock_stdin = pipe.SimpleInStream(sys.stdin.fileno()) + ql = Qiling(["../examples/rootfs/x86_linux/bin/crackme_linux"], "../examples/rootfs/x86_linux", console=False, stdin=mock_stdin) ins_count = [0] ql.hook_code(instruction_count, ins_count) ql.set_syscall("_llseek", my__llseek) - stdin.write(payload) + ql.os.stdin.write(payload) ql.run() - del stdin + del mock_stdin del ql return ins_count[0] @@ -1014,77 +979,20 @@ class ELFTest(unittest.TestCase): del ql def test_x8664_absolute_path(self): - class MyPipe(): - def __init__(self): - self.buf = b'' - - def write(self, s): - self.buf += s - - def read(self, l): - pass - - def fileno(self): - return 0 - - def fstat(self): - return Fstat(sys.stdin.fileno()) - - def show(self): - pass - - def clear(self): - pass - - def flush(self): - pass - - def close(self): - pass - - pipe = MyPipe() - ql = Qiling(["../examples/rootfs/x8664_linux/bin/absolutepath"], "../examples/rootfs/x8664_linux", verbose=QL_VERBOSE.DEBUG, stdout=pipe) + mock_stdout = pipe.SimpleOutStream(sys.stdout.fileno()) + ql = Qiling(["../examples/rootfs/x8664_linux/bin/absolutepath"], "../examples/rootfs/x8664_linux", verbose=QL_VERBOSE.DEBUG, stdout=mock_stdout) ql.run() - - self.assertEqual(pipe.buf, b'test_complete\n\ntest_complete\n\n') + self.assertEqual(ql.os.stdout.read(), b'test_complete\n\ntest_complete\n\n') del ql def test_x8664_getcwd(self): - class MyPipe(): - def __init__(self): - self.buf = b'' - - def write(self, s): - self.buf += s - - def read(self, l): - pass - - def fileno(self): - return 0 - - def fstat(self): - return Fstat(sys.stdin.fileno()) - - def show(self): - pass - - def clear(self): - pass - - def flush(self): - pass - - def close(self): - pass - - pipe = MyPipe() - ql = Qiling(["../examples/rootfs/x8664_linux/bin/testcwd"], "../examples/rootfs/x8664_linux", verbose=QL_VERBOSE.DEBUG, stdout=pipe) + mock_stdout = pipe.SimpleOutStream(sys.stdout.fileno()) + ql = Qiling(["../examples/rootfs/x8664_linux/bin/testcwd"], "../examples/rootfs/x8664_linux", verbose=QL_VERBOSE.DEBUG, stdout=mock_stdout) ql.run() - self.assertEqual(pipe.buf, b'/\n/lib\n/bin\n/\n') + self.assertEqual(ql.os.stdout.read(), b'/\n/lib\n/bin\n/\n') del ql @@ -1138,41 +1046,12 @@ class ELFTest(unittest.TestCase): del ql def test_elf_linux_x86_getdents64(self): - class MyPipe(): - def __init__(self): - self.buf = b'' + mock_stdout = pipe.SimpleOutStream(sys.stdout.fileno()) + ql = Qiling(["../examples/rootfs/x86_linux/bin/x86_getdents64"], "../examples/rootfs/x86_linux", verbose=QL_VERBOSE.DEBUG, stdout=mock_stdout) - def write(self, s): - self.buf += s - - def read(self, l): - pass - - def fileno(self): - return 0 - - def fstat(self): - return Fstat(sys.stdin.fileno()) - - def fstat64(self): - return self.fstat() - - def show(self): - pass - - def clear(self): - pass - - def flush(self): - pass - - def close(self): - pass - - pipe = MyPipe() - ql = Qiling(["../examples/rootfs/x86_linux/bin/x86_getdents64"], "../examples/rootfs/x86_linux", verbose=QL_VERBOSE.DEBUG, stdout=pipe) ql.run() - self.assertTrue("bin\n" in pipe.buf.decode("utf-8")) + self.assertTrue("bin\n" in ql.os.stdout.read().decode("utf-8")) + del ql if __name__ == "__main__": diff --git a/tests/test_pe.py b/tests/test_pe.py index 51947485..860b69c5 100644 --- a/tests/test_pe.py +++ b/tests/test_pe.py @@ -10,6 +10,7 @@ sys.path.append("..") from qiling import Qiling from qiling.const import * from qiling.exception import * +from qiling.extensions import pipe from qiling.loader.pe import QlPeCache from qiling.os.const import * from qiling.os.windows.fncc import * @@ -349,28 +350,6 @@ class PETest(unittest.TestCase): def test_pe_win_x86_crackme(self): - class StringBuffer: - def __init__(self): - self.buffer = b'' - - def read(self, n): - ret = self.buffer[:n] - self.buffer = self.buffer[n:] - return ret - - def readline(self, end=b'\n'): - ret = b'' - while True: - c = self.read(1) - ret += c - if c == end: - break - return ret - - def write(self, string): - self.buffer += string - return len(string) - def force_call_dialog_func(ql): # get DialogFunc address lpDialogFunc = ql.unpack32(ql.mem.read(ql.reg.esp - 0x8, 4)) @@ -390,7 +369,7 @@ class PETest(unittest.TestCase): ql.patch(0x0040110B, b'\x90\x90') ql.patch(0x00401112, b'\x90\x90') - ql.os.stdin = StringBuffer() + ql.os.stdin = pipe.SimpleStringBuffer() ql.os.stdin.write(b"Ea5yR3versing\n") ql.hook_address(force_call_dialog_func, 0x00401016) diff --git a/tests/test_windows_stdio.py b/tests/test_windows_stdio.py index c82b34d5..9050ec57 100644 --- a/tests/test_windows_stdio.py +++ b/tests/test_windows_stdio.py @@ -4,47 +4,31 @@ # import sys - -from unicorn import * -from unicorn.x86_const import * +from typing import Sequence sys.path.append("..") from qiling import * from qiling.const import QL_VERBOSE +from qiling.extensions import pipe -class StringBuffer: - def __init__(self): - self.buffer = b'' - - def read(self, n): - ret = self.buffer[:n] - self.buffer = self.buffer[n:] - return ret - - def read_all(self): - ret = self.buffer - self.buffer = b'' - return ret - - def write(self, string): - self.buffer += string - return len(string) - - -def instruction_count(ql, address, size, user_data): +def instruction_count(ql: Qiling, address: int, size: int, user_data): user_data[0] += 1 -def get_count(flag): +def get_count(flag: Sequence[str]): ql = Qiling(["../examples/rootfs/x86_windows/bin/crackme.exe"], "../examples/rootfs/x86_windows", verbose=QL_VERBOSE.OFF, libcache = True) - ql.os.stdin = StringBuffer() - ql.os.stdout = StringBuffer() - ql.os.stdin.write(bytes("".join(flag) + "\n", 'utf-8')) + ql.os.stdin = pipe.SimpleStringBuffer() + ql.os.stdout = pipe.SimpleStringBuffer() + count = [0] ql.hook_code(instruction_count, count) + + ql.os.stdin.write(bytes("".join(flag) + "\n", 'utf-8')) ql.run() - print(ql.os.stdout.read_all().decode('utf-8'), end='') - print(" ============ count: %d ============ " % count[0]) + + print(ql.os.stdout.read().decode('utf-8'), end='') + print(f' ============ count: {count[0]:d} ============ ') + return count[0]