diff --git a/keepercommander/commands/tunnel/port_forward/tunnel_helpers.py b/keepercommander/commands/tunnel/port_forward/tunnel_helpers.py index 9d26d78c5..b533d648e 100644 --- a/keepercommander/commands/tunnel/port_forward/tunnel_helpers.py +++ b/keepercommander/commands/tunnel/port_forward/tunnel_helpers.py @@ -167,6 +167,21 @@ def print_above_keeper_prompt(msg): pass +def _print_ai_closed_banner(): + """User-facing notice for a tunnel closed with the AIClosed reason. + + Wording matches `pam_launch.launch`'s `_print_close_reason_notice` for the + same condition — duplicated rather than shared across the two commands, + per this codebase's preference for small per-command copies over a + premature cross-module helper. + """ + print_above_keeper_prompt( + f"\n{bcolors.FAIL}Session terminated by KeeperAI.{bcolors.ENDC}\n" + "Critical activity was detected during this session.\n" + "Contact your administrator to unlock this record." + ) + + # Constants NONCE_LENGTH = 12 MAIN_NONCE_LENGTH = 16 @@ -2028,6 +2043,27 @@ def _signal_from_rust_inner(self, response: dict): elif close_reason.is_retryable(): logging.debug(f"{bcolors.WARNING}Tunnel closed with retryable error - '{tube_id}': {close_reason.name}{bcolors.ENDC}") + elif close_reason.code == CloseConnectionReasons.AIClosed: + # Explicit branch (rather than relying on the generic `else` + # below, where this currently lands) so the user actually + # sees why their session ended — mirrors `pam launch`'s + # handling of the same close reason. + logging.debug(f"{bcolors.OKBLUE}Tunnel '{tube_id}' closed by KeeperAI{bcolors.ENDC}") + _print_ai_closed_banner() + if self.record_uid: + try: + from ...tunnel_and_connections import ( + _AI_CLOSED_EXIT_CODE_BY_RECORD, + _LEASE_SHUTDOWN_EVENTS_BY_RECORD, + EXIT_CODE_AI_TERMINATED, + ) + _AI_CLOSED_EXIT_CODE_BY_RECORD[self.record_uid] = EXIT_CODE_AI_TERMINATED + shutdown_event = _LEASE_SHUTDOWN_EVENTS_BY_RECORD.get(self.record_uid) + if shutdown_event is not None: + shutdown_event.set() + except Exception: + logging.debug("Could not wake --foreground wait for AI-closed tunnel", exc_info=True) + else: logging.debug(f"{bcolors.OKBLUE}Tunnel '{tube_id}' closed with reason: {close_reason.name}{bcolors.ENDC}") diff --git a/keepercommander/commands/tunnel_and_connections.py b/keepercommander/commands/tunnel_and_connections.py index b183a39e1..44ad2090b 100644 --- a/keepercommander/commands/tunnel_and_connections.py +++ b/keepercommander/commands/tunnel_and_connections.py @@ -39,6 +39,7 @@ from .pam_import.keeper_ai_settings import get_resource_jit_settings from .pam_import.nsf_helpers import sync_down_preserving_nsf_keys from .pam_launch.terminal_connection import _extract_user_record_credentials +from .pam_launch.launch import EXIT_CODE_AI_TERMINATED from .tunnel_registry import ( PARENT_GRACE_SECONDS, is_pid_alive, @@ -67,6 +68,11 @@ # read by the lease-expiry callback. Default interactive mode does NOT register # (it has no blocking wait to interrupt; user SSH session continues naturally). _LEASE_SHUTDOWN_EVENTS_BY_RECORD: Dict[str, threading.Event] = {} +# Maps record_uid -> exit code set by tunnel_helpers' channel_closed handler when +# the Gateway/Rust reports an AIClosed close reason. Read by the --foreground +# block below after fg_shutdown wakes, so the process exits with the same code +# `pam launch` uses for the same condition. +_AI_CLOSED_EXIT_CODE_BY_RECORD: Dict[str, int] = {} def _coerce_settings_subdicts(entry, *keys): @@ -1699,11 +1705,13 @@ def _fg_signal_handler(signum, _frame): print(f"{bcolors.WARNING}Note: On Windows, tunnel stop uses hard termination. " f"WebRTC cleanup is best-effort.{bcolors.ENDC}\n") + ai_exit_code = None try: fg_shutdown.wait() except KeyboardInterrupt: pass finally: + ai_exit_code = _AI_CLOSED_EXIT_CODE_BY_RECORD.pop(record_uid, None) _LEASE_SHUTDOWN_EVENTS_BY_RECORD.pop(record_uid, None) unregister_tunnel() signal.signal(signal.SIGTERM, prev_sigterm) @@ -1727,6 +1735,14 @@ def _fg_signal_handler(signum, _frame): os.remove(pid_file) except OSError: pass + # Mirror `pam launch`'s exit code for the same condition, but only + # in batch / scripted mode. In the interactive shell a SystemExit + # would tear down the whole shell (and its other sessions); the + # user should land back at the Keeper> prompt. The banner itself + # is printed from tunnel_helpers' close_reason handling as soon as + # the signal arrives, not here. + if ai_exit_code is not None and getattr(params, 'batch_mode', False): + raise SystemExit(ai_exit_code) else: # Print failure message error_msg = result.get("error", "Unknown error") if result else "Failed to start tunnel" diff --git a/unit-tests/pam/test_pam_tunnel_ai_closed_banner.py b/unit-tests/pam/test_pam_tunnel_ai_closed_banner.py new file mode 100644 index 000000000..9e6837ce6 --- /dev/null +++ b/unit-tests/pam/test_pam_tunnel_ai_closed_banner.py @@ -0,0 +1,93 @@ +""" +Unit tests for `pam tunnel start`'s handling of the AIClosed (code 15) close +reason — the signal already reaches `TunnelSignalHandler` today, but used to +fall through to a silent `logging.debug(...)` with nothing shown to the user. +""" +import io +import json +import threading +import unittest +from contextlib import redirect_stdout + +from keepercommander.commands.pam_launch.launch import EXIT_CODE_AI_TERMINATED +from keepercommander.commands.tunnel.port_forward.tunnel_helpers import ( + TunnelSignalHandler, + _print_ai_closed_banner, +) +from keepercommander.commands.tunnel_and_connections import ( + _AI_CLOSED_EXIT_CODE_BY_RECORD, + _LEASE_SHUTDOWN_EVENTS_BY_RECORD, +) + + +class TestPrintAIClosedBanner(unittest.TestCase): + + def test_banner_text(self): + buf = io.StringIO() + with redirect_stdout(buf): + _print_ai_closed_banner() + out = buf.getvalue() + self.assertIn('Session terminated by KeeperAI.', out) + self.assertIn('Critical activity was detected during this session.', out) + self.assertIn('Contact your administrator to unlock this record.', out) + + +class TestTunnelSignalHandlerAIClosed(unittest.TestCase): + + def setUp(self): + self.record_uid = 'test-ai-closed-record-uid' + self.addCleanup(_AI_CLOSED_EXIT_CODE_BY_RECORD.pop, self.record_uid, None) + self.addCleanup(_LEASE_SHUTDOWN_EVENTS_BY_RECORD.pop, self.record_uid, None) + + def _make_handler(self): + return TunnelSignalHandler( + params=None, + record_uid=self.record_uid, + gateway_uid='gateway-uid', + symmetric_key=b'0' * 32, + base64_nonce='nonce', + conversation_id='conversation-id', + tube_registry=None, + ) + + def _send_ai_closed(self, handler): + payload = { + 'kind': 'channel_closed', + 'tube_id': 'fake-tube-id', + 'conversation_id': 'conversation-id', + 'data': json.dumps({'close_reason': {'code': 15}}), + } + buf = io.StringIO() + with redirect_stdout(buf): + handler.signal_from_rust(payload) + return buf.getvalue() + + def test_prints_banner_and_sets_exit_code(self): + handler = self._make_handler() + out = self._send_ai_closed(handler) + self.assertIn('Session terminated by KeeperAI.', out) + self.assertEqual( + _AI_CLOSED_EXIT_CODE_BY_RECORD.get(self.record_uid), + EXIT_CODE_AI_TERMINATED, + ) + + def test_wakes_registered_foreground_shutdown_event(self): + shutdown_event = threading.Event() + _LEASE_SHUTDOWN_EVENTS_BY_RECORD[self.record_uid] = shutdown_event + handler = self._make_handler() + self._send_ai_closed(handler) + self.assertTrue(shutdown_event.is_set()) + + def test_no_shutdown_event_registered_does_not_raise(self): + handler = self._make_handler() + # No entry in _LEASE_SHUTDOWN_EVENTS_BY_RECORD for this record_uid — + # mirrors interactive-shell mode, which never registers one. + self._send_ai_closed(handler) + self.assertEqual( + _AI_CLOSED_EXIT_CODE_BY_RECORD.get(self.record_uid), + EXIT_CODE_AI_TERMINATED, + ) + + +if __name__ == '__main__': + unittest.main()