84 lines
3.6 KiB
Python
84 lines
3.6 KiB
Python
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
@unittest.skipUnless(sys.platform == "win32", "Windows Job Object integration")
|
|
class WindowsContainmentTests(unittest.TestCase):
|
|
def run_child(self, source):
|
|
source = f"import sys; sys.path.insert(0, {str(Path(__file__).parent)!r});\n" + source
|
|
return subprocess.run([sys.executable, "-c", source], cwd=Path(__file__).parent,
|
|
capture_output=True, text=True, timeout=20,
|
|
creationflags=subprocess.CREATE_NO_WINDOW)
|
|
|
|
def test_process_limit_and_no_breakaway(self):
|
|
result = self.run_child("""
|
|
import json, subprocess, sys
|
|
from windows_containment import contain_inspector
|
|
status=contain_inspector()
|
|
children=[]
|
|
try:
|
|
for i in range(3):
|
|
children.append(subprocess.Popen([sys.executable,'-c','import time; time.sleep(15)'],creationflags=subprocess.CREATE_NO_WINDOW))
|
|
denied=False
|
|
try:
|
|
extra=subprocess.Popen([sys.executable,'-c','pass'],creationflags=subprocess.CREATE_NO_WINDOW)
|
|
extra.wait()
|
|
except OSError:
|
|
denied=True
|
|
breakaway=False
|
|
try:
|
|
escaped=subprocess.Popen([sys.executable,'-c','pass'],creationflags=subprocess.CREATE_NO_WINDOW | subprocess.CREATE_BREAKAWAY_FROM_JOB)
|
|
escaped.wait()
|
|
except OSError:
|
|
breakaway=True
|
|
print(json.dumps({'status':status,'denied':denied,'breakawayDenied':breakaway}),flush=True)
|
|
finally:
|
|
for child in children:
|
|
child.terminate();child.wait()
|
|
""")
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
proof = json.loads(result.stdout)
|
|
self.assertTrue(proof["status"]["active"])
|
|
self.assertEqual(proof["status"]["memoryBytes"], 4 * 1024 ** 3)
|
|
self.assertTrue(proof["denied"])
|
|
self.assertTrue(proof["breakawayDenied"])
|
|
|
|
def test_descendant_is_killed_when_owner_is_terminated(self):
|
|
source = """
|
|
import subprocess,sys,time
|
|
from windows_containment import contain_inspector
|
|
contain_inspector()
|
|
child=subprocess.Popen([sys.executable,'-c','import time; time.sleep(60)'],creationflags=subprocess.CREATE_NO_WINDOW)
|
|
print(child.pid,flush=True)
|
|
time.sleep(60)
|
|
"""
|
|
source = f"import sys; sys.path.insert(0, {str(Path(__file__).parent)!r});\n" + source
|
|
owner = subprocess.Popen([sys.executable, "-c", source], cwd=Path(__file__).parent,
|
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
|
|
creationflags=subprocess.CREATE_NO_WINDOW)
|
|
handle = None
|
|
try:
|
|
# Opening the descendant before termination permits a reliable wait,
|
|
# avoiding PID reuse or process enumeration heuristics.
|
|
child_pid = int(owner.stdout.readline())
|
|
import ctypes
|
|
kernel = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
kernel.OpenProcess.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.c_uint32]
|
|
kernel.OpenProcess.restype = ctypes.c_void_p
|
|
kernel.WaitForSingleObject.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
|
|
kernel.CloseHandle.argtypes = [ctypes.c_void_p]
|
|
handle = kernel.OpenProcess(0x00100000, False, child_pid)
|
|
self.assertTrue(handle)
|
|
owner.terminate();owner.wait(timeout=10)
|
|
self.assertEqual(kernel.WaitForSingleObject(handle, 10000), 0)
|
|
finally:
|
|
if owner.poll() is None:
|
|
owner.terminate();owner.wait(timeout=10)
|
|
if handle:
|
|
kernel.CloseHandle(handle)
|
|
owner.stdout.close();owner.stderr.close()
|