"""Opt-in checks against an already-provisioned disposable VM on an isolated host."""

from __future__ import annotations

import json
import os
import unittest
from urllib.parse import quote
from urllib.request import Request, urlopen


@unittest.skipUnless(os.environ.get("UVM_EXEC_INTEGRATION") == "1", "set UVM_EXEC_INTEGRATION=1 on an isolated VM host")
class GuestExecIntegrationTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.api_url = os.environ.get("UVM_API_URL", "http://127.0.0.1:8000").rstrip("/")
        cls.token = os.environ["UVM_API_TOKEN"]
        cls.vm_id = os.environ["UVM_EXEC_VM_ID"]
        cls.path = f"/vms/{quote(cls.vm_id, safe='')}"
        vm = cls.request("GET", cls.path)
        if (vm["id"] != cls.vm_id or vm["observed_status"] != "running"
                or vm["username"] == "root" or vm["ports"] != []):
            raise RuntimeError("provide a running disposable VM with a non-root login, installed agent, and empty TCP allowlist")
        cls.username = vm["username"]

    @classmethod
    def request(cls, method, path, body=None):
        request = Request(
            cls.api_url + path, method=method,
            data=json.dumps(body).encode("utf-8") if body is not None else None,
            headers={"X-UVM-Token": cls.token, "Content-Type": "application/json"},
        )
        with urlopen(request, timeout=30) as response:
            return json.load(response)

    def execute(self, command, *, sudo=False, timeout=10):
        return self.request("POST", self.path + "/exec", {
            "command": command, "sudo": sudo, "timeout_seconds": timeout,
        })

    def test_user_and_root_execution_work_with_all_tcp_ports_closed(self) -> None:
        ordinary = self.execute("id -u; whoami")
        self.assertEqual(ordinary["exit_code"], 0)
        uid, username = ordinary["stdout"].splitlines()
        self.assertNotEqual(uid, "0")
        self.assertEqual(username, self.username)
        elevated = self.execute("id -u", sudo=True)
        self.assertEqual(elevated["exit_code"], 0)
        self.assertEqual(elevated["stdout"], "0\n")

    def test_multiline_script_returns_output_and_nonzero_status(self) -> None:
        result = self.execute("printf 'hello\\n'\nprintf 'problem\\n' >&2\nexit 7")
        self.assertEqual(result, {
            "exit_code": 7, "stdout": "hello\n", "stderr": "problem\n", "timed_out": False, "truncated": False,
        })

    def test_timeout_kills_a_descendant_that_started_a_new_session(self) -> None:
        result = self.execute("""python3 - <<'PY'
import os, time
child = os.fork()
if child == 0:
    os.setsid()
else:
    print(child, flush=True)
time.sleep(30)
PY""", timeout=2)
        self.assertTrue(result["timed_out"])
        self.assertIsNone(result["exit_code"])
        child = int(result["stdout"].strip())
        checked = self.execute(f"""python3 - <<'PY'
from pathlib import Path
try:
    state = Path('/proc/{child}/stat').read_text().rsplit(') ', 1)[1].split()[0]
except FileNotFoundError:
    state = 'gone'
assert state in ('gone', 'Z'), state
PY""")
        self.assertEqual(checked["exit_code"], 0, checked["stderr"])

    def test_output_is_bounded_while_both_pipes_are_drained(self) -> None:
        result = self.execute("""python3 - <<'PY'
import os
for _ in range(32):
    os.write(1, b'o' * 65536)
    os.write(2, b'e' * 65536)
PY""")
        self.assertEqual(result["exit_code"], 0)
        self.assertTrue(result["truncated"])
        self.assertEqual(len(result["stdout"]) + len(result["stderr"]), 1024 * 1024)