Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor tests for Process #33

Merged
merged 4 commits into from
Apr 8, 2024
Merged

Refactor tests for Process #33

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3605dca
Refactor tests for Process, enhance logging and exception handling
jptomoya Apr 7, 2024
6559835
Fix ResourceWarning in tests.pwn.test_fsb.TestFSB
jptomoya Apr 7, 2024
a87f787
Fix logger determination logic
jptomoya Apr 7, 2024
ff6455a
Improve Process selection to indicate actual implementation
jptomoya Apr 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 4 additions & 11 deletions ptrlib/connection/proc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,6 @@
logger = getLogger(__name__)


def Process(*args, **kwargs) -> Tube:
if _is_windows:
return WinProcess(*args, **kwargs)
else:
return UnixProcess(*args, **kwargs)

class UnixProcess(Tube):
def __init__(
self,
Expand Down Expand Up @@ -198,15 +192,14 @@ def close(self):
"""
if self.proc:
os.close(self.slave)
self.proc.stdin.close()
self.proc.stdout.close()
if self.is_alive():
self.proc.stdin.close()
self.proc.stdout.close()
self.proc.kill()
self.proc.wait()
logger.info("'{0}' (PID={1}) killed".format(self.filepath, self.proc.pid))
self.proc = None
else:
self.proc.stdout.close()
logger.info("'{0}' (PID={1}) has already exited".format(self.filepath, self.proc.pid))
self.proc = None

Expand Down Expand Up @@ -242,5 +235,5 @@ def wait(self) -> int:
def __del__(self):
self.close()

# alias
process = Process
Process = WinProcess if _is_windows else UnixProcess
process = Process # alias for the Process
64 changes: 28 additions & 36 deletions tests/connection/test_proc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import unittest
import inspect
import os
import random
import unittest
from logging import FATAL, getLogger

from ptrlib import Process, is_scanf_safe
from logging import getLogger, FATAL

_is_windows = os.name == 'nt'

Expand All @@ -14,12 +16,17 @@ def setUp(self):
self.skipTest("This test is intended for the Linux platform")

def test_basic(self):
module_name = inspect.getmodule(Process).__name__

while True:
msg = os.urandom(16)
if is_scanf_safe(msg):
break

p = Process("./tests/test.bin/test_echo.x64")
with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_echo.x64")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')

# sendline / recvline
p.sendline(b"Message : " + msg)
Expand Down Expand Up @@ -60,56 +67,41 @@ def test_basic(self):
# wait
self.assertEqual(p.wait(), 0)

p.close()
with self.assertLogs(module_name) as cm:
p.close()
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:.+ \(PID=\d+\) has already exited$')

def test_timeout(self):
p = Process("./tests/test.bin/test_echo.x64")
module_name = inspect.getmodule(Process).__name__

with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_echo.x64")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')
data = os.urandom(16).hex()

# recv
try:
with self.assertRaises(TimeoutError) as cm:
p.recv(timeout=0.5)
result = False
except TimeoutError as err:
self.assertEqual(err.args[1], b"")
result = True
except:
result = False
self.assertEqual(result, True)
self.assertEqual(cm.exception.args[1], b"")

# recvonce
p.sendline(data)
try:
with self.assertRaises(TimeoutError) as cm:
p.recvonce(len(data) + 1 + 1, timeout=0.5)
result = False
except TimeoutError as err:
self.assertEqual(err.args[1].decode().strip(), data)
result = True
except:
result = False
self.assertEqual(result, True)
self.assertEqual(cm.exception.args[1].decode().strip(), data)

# recvuntil
p.sendline(data)
try:
with self.assertRaises(TimeoutError) as cm:
p.recvuntil("*** never expected ***", timeout=0.5)
result = False
except TimeoutError as err:
self.assertEqual(err.args[1].decode().strip(), data)
result = True
except:
result = False
self.assertEqual(result, True)
self.assertEqual(cm.exception.args[1].decode().strip(), data)

# sendlineafter
a, b = os.urandom(16).hex(), os.urandom(16).hex()
p.sendline(a)
try:
with self.assertRaises(TimeoutError) as cm:
p.sendlineafter(b"neko", b, timeout=0.5)
except TimeoutError as err:
self.assertEqual(err.args[1].decode().strip(), a)
result = True
except:
result = False
self.assertEqual(result, True)
self.assertEqual(cm.exception.args[1].decode().strip(), a)

33 changes: 19 additions & 14 deletions tests/connection/test_windows_proc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import unittest
import inspect
import os
import random
import subprocess
import unittest
from logging import FATAL, getLogger

from ptrlib import Process, is_scanf_safe
from logging import getLogger, FATAL

_is_windows = os.name == 'nt'

Expand All @@ -15,12 +17,17 @@ def setUp(self):
self.skipTest("This test is for Windows architecture")

def test_basic(self):
module_name = inspect.getmodule(Process).__name__

while True:
msg = os.urandom(16)
if is_scanf_safe(msg) and b'\x1a' not in msg:
break

p = Process("./tests/test.bin/test_echo.pe.exe")
with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_echo.pe.exe")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')
pid = p.pid

# send / recv
Expand All @@ -46,15 +53,13 @@ def test_basic(self):
self.assertFalse(str(pid) in subprocess.getoutput(f'tasklist /FI "PID eq {pid}"').split())

def test_timeout(self):
p = Process("./tests/test.bin/test_echo.pe.exe")
try:
module_name = inspect.getmodule(Process).__name__

with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_echo.pe.exe")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')

with self.assertRaises(TimeoutError):
p.recvuntil("*** never expected ***", timeout=1)
result = False
except TimeoutError:
result = True
except:
result = False
finally:
p.close()

self.assertEqual(result, True)
p.close()
37 changes: 23 additions & 14 deletions tests/pwn/test_fsb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import unittest
import inspect
import os
import unittest
from logging import FATAL, getLogger

from ptrlib import Process, fsb
from logging import getLogger, FATAL

_is_windows = os.name == 'nt'

Expand All @@ -14,10 +16,14 @@ def setUp(self):
self.skipTest("This test has not been implemented for Windows yet")

def test_fsb32(self):
module_name = inspect.getmodule(Process).__name__

# test 1
result = True
for i in range(3):
p = Process("./tests/test.bin/test_fsb.x86")
with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_fsb.x86")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')
p.recvuntil(": ")
target = int(p.recvline(), 16)
payload = fsb(
Expand All @@ -28,14 +34,15 @@ def test_fsb32(self):
)
p.sendline(payload + b'XXXXXXXX')
p.recvuntil("XXXXXXXX\n")
result |= b'OK' in p.recvline()
self.assertTrue(b'OK' in p.recvline())
p.close()
self.assertEqual(result, True)

# test 2
result = True
for i in range(3):
p = Process("./tests/test.bin/test_fsb.x86")
with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_fsb.x86")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')
p.recvuntil(": ")
target = int(p.recvline(), 16)
payload = fsb(
Expand All @@ -47,15 +54,18 @@ def test_fsb32(self):
)
p.sendline(payload + b'XXXXXXXX')
p.recvuntil("XXXXXXXX\n")
result |= b'OK' in p.recvline()
self.assertTrue(b'OK' in p.recvline())
p.close()
self.assertEqual(result, True)

def test_fsb64(self):
module_name = inspect.getmodule(Process).__name__

# test 3
result = True
for i in range(3):
p = Process("./tests/test.bin/test_fsb.x64")
with self.assertLogs(module_name) as cm:
p = Process("./tests/test.bin/test_fsb.x64")
self.assertEqual(len(cm.output), 1)
self.assertRegex(cm.output[0], fr'^INFO:{module_name}:Successfully created new process \(PID=\d+\)$')
p.recvuntil(": ")
target = int(p.recvline(), 16)
payload = fsb(
Expand All @@ -65,6 +75,5 @@ def test_fsb64(self):
bits = 64
)
p.sendline(payload)
result |= b'OK' in p.recvuntil("OK")
self.assertTrue(b'OK' in p.recvuntil("OK"))
p.close()
self.assertEqual(result, True)