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()