diff --git a/riocli/config/config.py b/riocli/config/config.py index 6d05ce43..87649d62 100644 --- a/riocli/config/config.py +++ b/riocli/config/config.py @@ -57,6 +57,7 @@ class Configuration: OIDC_SERVER = "https://oidc.rapyuta.io" DIFF_TOOL = "diff" MERGE_TOOL = "vimdiff" + AUTO_DISCONNECT_VPN = True DEVICE_FLOW_CLIENT_ID = "rio-cli" def __init__(self, filepath: str | None = None): @@ -216,6 +217,15 @@ def device_flow_client_id(self: Configuration) -> str: def oidc_server(self: Configuration) -> str: return self.data.get("oidc_host", self.OIDC_SERVER) + @property + def auto_disconnect_vpn(self: Configuration) -> bool: + return bool(self.data.get("auto_disconnect_vpn", self.AUTO_DISCONNECT_VPN)) + + @property + def current_project_id(self: Configuration) -> str | None: + """Current project guid, or None if no project is selected. Never raises.""" + return self.data.get("project_id") or None + @property def machine_id(self: Configuration): if "machine_id" not in self.data: diff --git a/riocli/organization/select.py b/riocli/organization/select.py index 34c525b5..5dd59bb6 100644 --- a/riocli/organization/select.py +++ b/riocli/organization/select.py @@ -17,10 +17,10 @@ from click_help_colors import HelpColorsCommand from riocli.auth.util import select_project -from riocli.constants import Colors, Symbols +from riocli.constants import Colors from riocli.organization.util import name_to_guid from riocli.utils.context import get_root_context -from riocli.vpn.util import cleanup_hosts_file +from riocli.vpn.util import disconnect_vpn_for_switch @click.command( @@ -30,6 +30,13 @@ help_options_color=Colors.GREEN, ) @click.argument("organization-name", type=str) +@click.option( + "--keep-vpn", + is_flag=True, + default=False, + help="Keep the VPN connected after switching organizations. Skips both " + "VPN disconnect and hosts file cleanup.", +) @click.option( "--interactive/--no-interactive", is_flag=True, @@ -51,6 +58,7 @@ def select_organization( organization_name: str, organization_guid: str, organization_short_id: str, + keep_vpn: bool, interactive: bool, silent: bool, ) -> None: @@ -66,6 +74,13 @@ def select_organization( If your organization name has spaces, use quotes around the name. + By default, if a VPN is active it will be disconnected and the + hosts file will be cleaned up on switch. Use --keep-vpn to suppress + this. You can also set ``auto_disconnect_vpn: false`` in the CLI + config file (``~/.config/rio-cli/config.json`` on Linux, + ``~/Library/Application Support/rio-cli/config.json`` on macOS) + to permanently suppress auto-disconnect. + Usage Examples: Set the current organization to 'Platform JP Staging' @@ -103,13 +118,7 @@ def select_organization( ctx.obj.save() - try: - cleanup_hosts_file() - except Exception as e: - click.secho( - f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}", - fg=Colors.YELLOW, - ) + disconnect_vpn_for_switch(ctx.obj, keep_vpn) if ctx.obj.data.get("project_id"): from riocli.ssh import refresh_ssh_cert diff --git a/riocli/project/select.py b/riocli/project/select.py index e69e319a..e707f6a7 100644 --- a/riocli/project/select.py +++ b/riocli/project/select.py @@ -17,7 +17,7 @@ from riocli.constants import Colors, Symbols from riocli.project.util import name_to_guid from riocli.utils.context import get_root_context -from riocli.vpn.util import cleanup_hosts_file +from riocli.vpn.util import disconnect_vpn_for_switch @click.command( @@ -27,31 +27,43 @@ help_options_color=Colors.GREEN, ) @click.argument("project-name", type=str) +@click.option( + "--keep-vpn", + is_flag=True, + default=False, + help="Keep the VPN connected after switching projects. Skips both " + "VPN disconnect and hosts file cleanup.", +) @name_to_guid @click.pass_context def select_project( ctx: click.Context, project_name: str, project_guid: str, + keep_vpn: bool, ) -> None: """Switch to a different project in the current organization. The project will be set in the CLI's context and will be used for all the subsequent commands. + + By default, if a VPN is active it will be disconnected and the + hosts file will be cleaned up. Use --keep-vpn to suppress this, + for example when you have an active SSH session into a device on + the previous project. You can also set ``auto_disconnect_vpn: false`` + in the CLI config file (``~/.config/rio-cli/config.json`` on Linux, + ``~/Library/Application Support/rio-cli/config.json`` on macOS) + to permanently suppress auto-disconnect. """ ctx = get_root_context(ctx) + previous_project_id = ctx.obj.current_project_id ctx.obj.data["project_id"] = project_guid ctx.obj.data["project_name"] = project_name ctx.obj.save() - try: - cleanup_hosts_file() - except Exception as e: - click.secho( - f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}", - fg=Colors.YELLOW, - ) + if project_guid != previous_project_id: + disconnect_vpn_for_switch(ctx.obj, keep_vpn) click.secho( f"{Symbols.SUCCESS} Project {project_name} ({project_guid}) is selected!", diff --git a/riocli/vpn/util.py b/riocli/vpn/util.py index 03fcad3d..9a445165 100644 --- a/riocli/vpn/util.py +++ b/riocli/vpn/util.py @@ -87,6 +87,40 @@ def get_tailscale_status() -> dict: return json.loads(output) +def disconnect_vpn_for_switch(config: "Configuration", keep_vpn: bool) -> None: # noqa: F821 + """Disconnect VPN and clean up /etc/hosts when switching project or org. + + Both actions are skipped together when --keep-vpn is passed or when the + user has set ``auto_disconnect_vpn: false`` in the CLI config file + (``~/.config/rio-cli/config.json`` on Linux, + ``~/Library/Application Support/rio-cli/config.json`` on macOS, + or the path in ``$RIO_CONFIG``). + + /etc/hosts is only cleaned when Tailscale was already down or was + successfully disconnected — never when disconnect fails — to avoid a + state where hosts entries point at a dead VPN tunnel. + """ + if keep_vpn or not config.auto_disconnect_vpn: + return + + vpn_was_up = is_tailscale_up() + disconnected = True + if vpn_was_up: + disconnected = stop_tailscale() + if disconnected: + click.secho(f"{Symbols.SUCCESS} VPN disconnected.", fg=Colors.GREEN) + else: + click.secho(f"{Symbols.WARNING} Failed to disconnect VPN.", fg=Colors.YELLOW) + if not vpn_was_up or disconnected: + try: + cleanup_hosts_file() + except Exception as e: + click.secho( + f"{Symbols.WARNING} Failed to clean up hosts file: {str(e)}", + fg=Colors.YELLOW, + ) + + def install_vpn_tools(force: bool = False) -> None: if is_tailscale_installed(): return diff --git a/tests/unit/vpn/test_auto_disconnect.py b/tests/unit/vpn/test_auto_disconnect.py new file mode 100644 index 00000000..d411e67d --- /dev/null +++ b/tests/unit/vpn/test_auto_disconnect.py @@ -0,0 +1,163 @@ +from unittest.mock import MagicMock, patch + +from click.testing import CliRunner + + +def _make_project_ctx(auto_disconnect_vpn=True, project_id="old-guid"): + obj = MagicMock() + obj.auto_disconnect_vpn = auto_disconnect_vpn + obj.current_project_id = project_id + obj.data = { + "project_id": project_id, + "project_name": "old-project", + "organization_id": "org-guid", + } + return obj + + +def _make_org_ctx(auto_disconnect_vpn=True): + obj = MagicMock() + obj.auto_disconnect_vpn = auto_disconnect_vpn + obj.data = { + "organization_id": "different-org-guid", + "organization_name": "old-org", + "organization_short_id": "old-short", + } + return obj + + +class TestProjectSelectVpnDisconnect: + def _invoke(self, args, ctx_obj): + from riocli.project.select import select_project + + with ( + patch("riocli.project.util.new_v2_client"), + patch("riocli.project.util.find_project_guid", return_value="new-guid"), + patch("riocli.project.util.get_project_name", return_value="new-project"), + patch("riocli.project.select.get_root_context") as mock_get_ctx, + ): + mock_get_ctx.return_value.obj = ctx_obj + return CliRunner().invoke(select_project, args, obj=ctx_obj), mock_get_ctx + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale", return_value=True) + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_disconnects_vpn_when_tailscale_up(self, mock_cleanup, mock_stop, mock_is_up): + result, _ = self._invoke(["new-project"], _make_project_ctx()) + assert result.exit_code == 0 + mock_stop.assert_called_once() + mock_cleanup.assert_called_once() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=False) + @patch("riocli.vpn.util.stop_tailscale") + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_skips_stop_when_tailscale_not_up(self, mock_cleanup, mock_stop, mock_is_up): + result, _ = self._invoke(["new-project"], _make_project_ctx()) + assert result.exit_code == 0 + mock_stop.assert_not_called() + mock_cleanup.assert_called_once() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale") + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_keep_vpn_flag_skips_disconnect_and_cleanup( + self, mock_cleanup, mock_stop, mock_is_up + ): + result, _ = self._invoke(["new-project", "--keep-vpn"], _make_project_ctx()) + assert result.exit_code == 0 + mock_stop.assert_not_called() + mock_cleanup.assert_not_called() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale") + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_config_opt_out_skips_disconnect_and_cleanup( + self, mock_cleanup, mock_stop, mock_is_up + ): + result, _ = self._invoke( + ["new-project"], _make_project_ctx(auto_disconnect_vpn=False) + ) + assert result.exit_code == 0 + mock_stop.assert_not_called() + mock_cleanup.assert_not_called() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale", return_value=False) + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_hosts_not_cleaned_when_stop_fails(self, mock_cleanup, mock_stop, mock_is_up): + result, _ = self._invoke(["new-project"], _make_project_ctx()) + assert result.exit_code == 0 + mock_stop.assert_called_once() + mock_cleanup.assert_not_called() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale") + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_noop_switch_skips_vpn_disconnect(self, mock_cleanup, mock_stop, mock_is_up): + # Selecting the already-selected project must not touch VPN. + result, _ = self._invoke( + ["new-project"], _make_project_ctx(project_id="new-guid") + ) + assert result.exit_code == 0 + mock_stop.assert_not_called() + mock_cleanup.assert_not_called() + + +class TestOrgSelectVpnDisconnect: + def _invoke(self, args, ctx_obj): + from riocli.organization.select import select_organization + + with ( + patch("riocli.organization.util.new_v2_client"), + patch( + "riocli.organization.util.find_organization_guid", + return_value=("new-org-guid", "new-short"), + ), + patch("riocli.organization.select.get_root_context") as mock_get_ctx, + ): + mock_get_ctx.return_value.obj = ctx_obj + return CliRunner().invoke(select_organization, args), mock_get_ctx + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale", return_value=True) + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_disconnects_vpn_when_tailscale_up(self, mock_cleanup, mock_stop, mock_is_up): + result, _ = self._invoke(["new-org", "--no-interactive"], _make_org_ctx()) + assert result.exit_code == 0 + mock_stop.assert_called_once() + mock_cleanup.assert_called_once() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale") + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_keep_vpn_flag_skips_disconnect_and_cleanup( + self, mock_cleanup, mock_stop, mock_is_up + ): + result, _ = self._invoke( + ["new-org", "--keep-vpn", "--no-interactive"], _make_org_ctx() + ) + assert result.exit_code == 0 + mock_stop.assert_not_called() + mock_cleanup.assert_not_called() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale") + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_config_opt_out_skips_disconnect_and_cleanup( + self, mock_cleanup, mock_stop, mock_is_up + ): + result, _ = self._invoke( + ["new-org", "--no-interactive"], _make_org_ctx(auto_disconnect_vpn=False) + ) + assert result.exit_code == 0 + mock_stop.assert_not_called() + mock_cleanup.assert_not_called() + + @patch("riocli.vpn.util.is_tailscale_up", return_value=True) + @patch("riocli.vpn.util.stop_tailscale", return_value=False) + @patch("riocli.vpn.util.cleanup_hosts_file") + def test_hosts_not_cleaned_when_stop_fails(self, mock_cleanup, mock_stop, mock_is_up): + result, _ = self._invoke(["new-org", "--no-interactive"], _make_org_ctx()) + assert result.exit_code == 0 + mock_stop.assert_called_once() + mock_cleanup.assert_not_called()