#!/usr/bin/env python3
# Copyright (c) 2026 Storm Dragon
# Authors: Storm Dragon, Angus Kola
# SPDX-License-Identifier: GPL-3.0-or-later
"""Runtime isolation and real AT-SPI socket collision regressions."""
import ast
import json
import os
from pathlib import Path
import socket
import stat
import subprocess
import sys
import tempfile
import unittest
from unittest.mock import patch

from native_test_session import run_private_session


class PrivateSessionTests(unittest.TestCase):
    def test_child_runtime_and_bus_are_private(self):
        environment = dict(os.environ, XDG_RUNTIME_DIR="/invalid/desktop-runtime",
                           DBUS_SESSION_BUS_ADDRESS="unix:path=/invalid/desktop-bus",
                           AT_SPI_BUS_ADDRESS="unix:path=/invalid/accessibility-bus",
                           DBUS_STARTER_ADDRESS="unix:path=/invalid/starter-bus",
                           DBUS_STARTER_BUS_TYPE="session", WAYLAND_DISPLAY="desktop-wayland")
        originalEnv = environment.copy()
        script = """import json, os, socket, stat
from pathlib import Path
runtime = Path(os.environ['XDG_RUNTIME_DIR'])
listener = socket.socket(socket.AF_UNIX)
listener.bind(str(runtime / 'probe.sock'))
print(json.dumps(dict(runtime=str(runtime), mode=stat.S_IMODE(runtime.stat().st_mode),
    bus=os.environ['DBUS_SESSION_BUS_ADDRESS'],
    inherited=[name for name in ('AT_SPI_BUS_ADDRESS', 'DBUS_STARTER_ADDRESS',
        'DBUS_STARTER_BUS_TYPE', 'WAYLAND_DISPLAY') if name in os.environ])))
"""
        result = run_private_session(
            ["dbus-run-session", "--", sys.executable, "-c", script], env=environment,
            capture_output=True, text=True, check=True, timeout=15)
        state = json.loads(result.stdout)
        self.assertEqual(state["mode"], 0o700)
        self.assertEqual(state["inherited"], [])
        self.assertNotEqual(state["bus"], originalEnv["DBUS_SESSION_BUS_ADDRESS"])
        self.assertFalse(Path(state["runtime"]).exists(), "child runtime was not cleaned up")
        self.assertEqual(environment, originalEnv)

    def test_cleanup_after_child_failure(self):
        script = "import os, sys; print(os.environ['XDG_RUNTIME_DIR'], flush=True); sys.exit(7)"
        with self.assertRaises(subprocess.CalledProcessError) as caught:
            run_private_session(
                ["dbus-run-session", "--", sys.executable, "-c", script], env=os.environ,
                capture_output=True, text=True, check=True, timeout=15)
        self.assertEqual(caught.exception.returncode, 7)
        self.assertFalse(Path(caught.exception.stdout.strip()).exists())

    def test_rejects_dbus_before_xvfb_without_launching(self):
        with patch("native_test_session.subprocess.run") as launch:
            with self.assertRaisesRegex(ValueError, "Start Xvfb before D-Bus"):
                run_private_session(
                    ["dbus-run-session", "--", "xvfb-run", "-a", "true"], env=os.environ)
        launch.assert_not_called()

    def test_cli_activates_atspi_on_virtual_display(self):
        script = """import ast, json, os, subprocess
result = subprocess.run(['gdbus', 'call', '--session', '--dest', 'org.a11y.Bus',
    '--object-path', '/org/a11y/bus', '--method', 'org.a11y.Bus.GetAddress'],
    capture_output=True, text=True, check=True, timeout=10)
print(json.dumps(dict(display=os.environ['DISPLAY'], runtime=os.environ['XDG_RUNTIME_DIR'],
    address=ast.literal_eval(result.stdout)[0])))
"""
        result = subprocess.run(
            [sys.executable, "-B", str(Path(__file__).with_name("native_test_session.py")),
             sys.executable, "-c", script],
            env=dict(os.environ, DISPLAY=":7777"), capture_output=True, text=True,
            check=True, timeout=20)
        state = json.loads(result.stdout)
        self.assertNotEqual(state["display"], ":7777")
        expectedPath = Path(state["runtime"]) / "at-spi" / ("bus" + state["display"].replace(":", "_"))
        self.assertEqual(state["address"].split(",", 1)[0], "unix:path=" + str(expectedPath))
        self.assertFalse(Path(state["runtime"]).exists())

    def test_atspi_activation_preserves_other_session_socket(self):
        # No real display is contacted. Reproduce D-Bus inheriting one display
        # while its client uses another, with a live canary in the parent runtime.
        with tempfile.TemporaryDirectory(prefix="cth-parent-", dir="/tmp") as parentRuntime:
            socketDir = Path(parentRuntime) / "at-spi"
            socketDir.mkdir(mode=0o700)
            socketPath = socketDir / "bus_7777"
            with socket.socket(socket.AF_UNIX) as listener:
                listener.bind(str(socketPath))
                listener.listen(1)
                originalInode = socketPath.stat().st_ino
                environment = dict(os.environ, XDG_RUNTIME_DIR=parentRuntime, DISPLAY=":7777",
                                   AT_SPI_BUS_ADDRESS="unix:path=/invalid/desktop-accessibility")
                result = run_private_session(
                    ["dbus-run-session", "--", "env", "DISPLAY=:7778", "gdbus", "call",
                     "--session", "--dest", "org.a11y.Bus", "--object-path", "/org/a11y/bus",
                     "--method", "org.a11y.Bus.GetAddress"], env=environment,
                    capture_output=True, text=True, check=True, timeout=15)
                address = ast.literal_eval(result.stdout)[0]
                busPath = Path(address.split(",", 1)[0].removeprefix("unix:path="))
                self.assertEqual(busPath.name, "bus_7777", address)
                self.assertNotEqual(busPath.parent.parent, Path(parentRuntime))
                self.assertFalse(busPath.parent.parent.exists(), "private runtime survived cleanup")
                self.assertEqual(socketPath.stat().st_ino, originalInode)
                self.assertTrue(stat.S_ISSOCK(socketPath.stat().st_mode))
                with socket.socket(socket.AF_UNIX) as client:
                    client.settimeout(2)
                    client.connect(str(socketPath))
                    connection, _ = listener.accept()
                    connection.close()


if __name__ == "__main__":
    unittest.main()