diff --git a/src/netshaper/core/mitm_manager.py b/src/netshaper/core/mitm_manager.py index e45f4e7..f7f14f3 100644 --- a/src/netshaper/core/mitm_manager.py +++ b/src/netshaper/core/mitm_manager.py @@ -240,7 +240,10 @@ def get_state_for_persistence(self) -> dict: if self._mitm_process_identity is None: log.error("Live mitmproxy process has no recoverable identity") return state - command = list(self._mitm_command or getattr(proc, "args", []) or []) + command = list(self._mitm_command or []) + if not command: + log.error("Live mitmproxy process has no recoverable command") + return state state.update( { "service": "mitmproxy", diff --git a/src/netshaper/core/orchestrator.py b/src/netshaper/core/orchestrator.py index 4affc66..fe18041 100644 --- a/src/netshaper/core/orchestrator.py +++ b/src/netshaper/core/orchestrator.py @@ -125,13 +125,17 @@ def _portal_manager(self) -> PortalManager: self.authorized_cidrs, journal=self._journal_state_if_ready, ) + self.portal_manager = manager legacy_proc = getattr(self, "_fake_server_proc", None) - if legacy_proc is not None: - manager.process = legacy_proc legacy_token = getattr(self, "_fake_server_health_token", None) - if legacy_token: - manager._health_token = legacy_token - self.portal_manager = manager + if legacy_proc is not None: + if not manager.attach_owned_process( + legacy_proc, + health_token=legacy_token, + ): + self._fake_server_proc = None + elif legacy_token: + manager.use_health_token(legacy_token) return manager @staticmethod diff --git a/src/netshaper/core/portal_manager.py b/src/netshaper/core/portal_manager.py index 17ac276..1e4e980 100644 --- a/src/netshaper/core/portal_manager.py +++ b/src/netshaper/core/portal_manager.py @@ -69,6 +69,12 @@ def start(self, portal_config: PortalConfig) -> bool: return True if self.process and self.process.poll() is None: + if self._command is None: + self._command = self._command_from_process(self.process) + if not self._command: + log.error("Could not establish recoverable portal process command") + self.stop() + return False if self._process_identity is None and not self._capture_process_identity( self.process, ): @@ -148,6 +154,40 @@ def start(self, portal_config: PortalConfig) -> bool: self.stop() return False + def attach_owned_process( + self, + process: subprocess.Popen[Any], + *, + health_token: Optional[str] = None, + ) -> bool: + """Attach a legacy NetShaper-owned portal child to this manager.""" + if process.poll() is not None: + return False + + command = self._command_from_process(process) + self.process = process + self._command = command + if health_token: + self._health_token = health_token + + if not command: + log.error("Could not establish recoverable portal process command") + self.stop() + return False + if not self._capture_process_identity(process): + log.error("Could not establish recoverable portal process identity") + self.stop() + return False + if not self._journal_state(): + log.error("Refusing to attach netshaper-portal without recovery state") + self.stop() + return False + return True + + def use_health_token(self, token: str) -> None: + """Use an existing portal health token without adopting a process.""" + self._health_token = token + def health_token(self) -> str: if not self._health_token: self._health_token = secrets.token_urlsafe(32) @@ -214,6 +254,15 @@ def _capture_process_identity(self, process: subprocess.Popen[Any]) -> bool: self._process_identity = dict(identity) return True + @staticmethod + def _command_from_process(process: subprocess.Popen[Any]) -> list[str]: + args = getattr(process, "args", None) + if isinstance(args, (list, tuple)): + return [str(item) for item in args] + if isinstance(args, str) and args: + return [args] + return [] + def _journal_state(self) -> bool: if self._journal is None: return True @@ -230,7 +279,10 @@ def get_state_for_persistence(self) -> dict[str, object]: if self._process_identity is None: log.error("Live portal process has no recoverable identity") return {} - command = list(self._command or getattr(process, "args", []) or []) + command = list(self._command or []) + if not command: + log.error("Live portal process has no recoverable command") + return {} executable = command[0] if command else None return { "service": "portal", diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py index 4930b99..0a124e8 100644 --- a/tests/test_orchestrator.py +++ b/tests/test_orchestrator.py @@ -420,6 +420,46 @@ def test_fake_server_health_token_is_stable_for_manual_launch(self): self.assertEqual(ns.fake_server_health_token(), "manual-token") self.assertEqual(ns.fake_server_health_token(), "manual-token") + def test_legacy_fake_server_proc_attaches_recoverable_portal_before_ready(self): + ns = NetShaper.__new__(NetShaper) + self._set_authorized(ns) + ns.own_ip = "192.0.2.1" + ns._fake_server_health_token = "test-health-token" + process = mock.Mock() + process.pid = 1234 + process.poll.return_value = None + process.args = [ + "/usr/bin/python3", + "-m", + "netshaper.portal", + "--health-token", + "test-health-token", + ] + ns._fake_server_proc = process + + with mock.patch( + "netshaper.core.portal_manager.process_owner_metadata", + return_value={ + "pid": 1234, + "process_create_time": 456.0, + "created_at": 789.0, + }, + ) as owner_metadata, mock.patch( + "netshaper.core.portal_manager.PortalManager.health_ready", + return_value=True, + ), mock.patch( + "netshaper.core.portal_manager.subprocess.Popen" + ) as popen: + self.assertTrue(ns.fake_server_ready()) + state = ns._portal_manager().get_state_for_persistence() + + owner_metadata.assert_called_once_with(1234) + popen.assert_not_called() + self.assertEqual(state["pid"], 1234) + self.assertEqual(state["process_create_time"], 456.0) + self.assertEqual(state["argv"], process.args) + self.assertEqual(state["ownership_token"], "test-health-token") + def test_fake_server_launch_refuses_unverified_claimed_listener(self): ns = NetShaper.__new__(NetShaper) ns.own_ip = "192.0.2.1" diff --git a/tests/test_portal_manager.py b/tests/test_portal_manager.py index ed07951..7fa7f3e 100644 --- a/tests/test_portal_manager.py +++ b/tests/test_portal_manager.py @@ -66,11 +66,95 @@ def test_state_includes_owned_portal_process(self): self.assertEqual(state["ownership_token"], "token") self.assertIn("netshaper.portal", state["argv"]) + def test_attach_owned_process_captures_identity_and_command(self): + journal = mock.Mock(return_value=True) + manager = PortalManager( + "192.0.2.10", + ["192.0.2.0/24"], + journal=journal, + ) + process = mock.Mock() + process.pid = 1234 + process.poll.return_value = None + process.args = [ + "/usr/bin/python3", + "-m", + "netshaper.portal", + "--health-token", + "token", + ] + + with mock.patch( + "netshaper.core.portal_manager.process_owner_metadata", + return_value={ + "pid": 1234, + "process_create_time": 456.0, + "created_at": 789.0, + }, + ) as owner_metadata: + self.assertTrue( + manager.attach_owned_process(process, health_token="token") + ) + state = manager.get_state_for_persistence() + + owner_metadata.assert_called_once_with(1234) + journal.assert_called_once() + self.assertEqual(state["pid"], 1234) + self.assertEqual(state["process_create_time"], 456.0) + self.assertEqual(state["argv"], process.args) + self.assertEqual(state["ownership_token"], "token") + + def test_attach_owned_process_stops_when_command_is_missing(self): + manager = PortalManager("192.0.2.10", ["192.0.2.0/24"]) + process = mock.Mock() + process.pid = 1234 + process.poll.side_effect = [None, None, 0] + + self.assertFalse(manager.attach_owned_process(process, health_token="token")) + + process.terminate.assert_called_once() + self.assertIsNone(manager.process) + + def test_attach_owned_process_stops_when_identity_capture_fails(self): + manager = PortalManager("192.0.2.10", ["192.0.2.0/24"]) + process = mock.Mock() + process.pid = 1234 + process.poll.side_effect = [None, None, 0] + process.args = [ + "/usr/bin/python3", + "-m", + "netshaper.portal", + "--health-token", + "token", + ] + + with mock.patch( + "netshaper.core.portal_manager.process_owner_metadata", + return_value={ + "pid": 1234, + "process_create_time": None, + "created_at": 789.0, + }, + ): + self.assertFalse( + manager.attach_owned_process(process, health_token="token") + ) + + process.terminate.assert_called_once() + self.assertIsNone(manager.process) + def test_start_waits_for_existing_child(self): manager = PortalManager("192.0.2.10", ["192.0.2.0/24"]) manager.process = mock.Mock() manager.process.pid = 1234 manager.process.poll.return_value = None + manager.process.args = [ + "/usr/bin/python3", + "-m", + "netshaper.portal", + "--health-token", + "token", + ] with mock.patch.object(manager, "ready", side_effect=[False, True]), \ mock.patch("netshaper.core.portal_manager.check_local_port"