diff --git a/colab-restore.ps1 b/colab-restore.ps1 new file mode 100644 index 0000000..0d886df --- /dev/null +++ b/colab-restore.ps1 @@ -0,0 +1,87 @@ +# Colab CLI Windows Native — One-Click Restore +# ============================================== +# Usage: .\colab-restore.ps1 +# Installs the Windows-compatible fork, configures SSL certs, and verifies. +# +# Source: C:\Users\woodh\Documents\colab-cli-windows\ +# PR: https://github.com/googlecolab/google-colab-cli/pull/70 + +param( + [switch]$SkipInstall, + [switch]$SkipSSL, + [switch]$SkipVerify, + [switch]$SkipADC +) + +$ErrorActionPreference = "Stop" +$certBundle = "C:\anaconda3\Lib\site-packages\certifi\cacert.pem" + +Write-Host "=== Colab CLI Windows Restore ===" -ForegroundColor Cyan + +# ── 1. Install ────────────────────────────────────────── +if (-not $SkipInstall) { + Write-Host "[1/4] Installing colab CLI (Windows fork)..." -ForegroundColor Yellow + pip install git+https://github.com/woodhaha/google-colab-cli.git@windows-support --quiet 2>&1 | Out-Null + Write-Host " Installed: $(colab version 2>&1)" -ForegroundColor Green +} + +# ── 2. ADC Auth ────────────────────────────────────────── +if (-not $SkipADC) { + Write-Host "[2/4] Checking ADC auth..." -ForegroundColor Yellow + $adcFile = "$env:APPDATA\gcloud\application_default_credentials.json" + if (Test-Path $adcFile) { + Write-Host " ADC file exists: $adcFile" -ForegroundColor Green + } else { + Write-Host " No ADC credentials found. Run:" -ForegroundColor Red + Write-Host ' & "C:\Program Files (x86)\Google\Cloud SDK\google-cloud-sdk\bin\gcloud.cmd" auth application-default login --scopes=openid,https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/userinfo.email,https://www.googleapis.com/auth/colaboratory' -ForegroundColor White + } +} + +# ── 3. SSL cert env vars ───────────────────────────────── +if (-not $SkipSSL) { + Write-Host "[3/4] Configuring SSL cert env vars..." -ForegroundColor Yellow + $profilePath = $PROFILE.CurrentUserCurrentHost + $profileDir = Split-Path $profilePath -Parent + if (-not (Test-Path $profileDir)) { New-Item -ItemType Directory -Force $profileDir | Out-Null } + if (-not (Test-Path $profilePath)) { New-Item -ItemType File -Force $profilePath | Out-Null } + + $lines = @( + '$env:SSL_CERT_FILE = "C:\anaconda3\Lib\site-packages\certifi\cacert.pem"', + '$env:REQUESTS_CA_BUNDLE = "C:\anaconda3\Lib\site-packages\certifi\cacert.pem"' + ) + $existing = Get-Content $profilePath -Raw -ErrorAction SilentlyContinue + foreach ($line in $lines) { + if ($existing -notmatch [regex]::Escape($line)) { + Add-Content $profilePath $line + Write-Host " Added to profile: $line" -ForegroundColor Green + } else { + Write-Host " Already in profile: $line" -ForegroundColor Gray + } + } + + # Also set for current session + $env:SSL_CERT_FILE = $certBundle + $env:REQUESTS_CA_BUNDLE = $certBundle +} + +# ── 4. Verify ──────────────────────────────────────────── +if (-not $SkipVerify) { + Write-Host "[4/4] Verifying..." -ForegroundColor Yellow + $result = colab --auth=adc sessions 2>&1 + if ($LASTEXITCODE -eq 0) { + Write-Host " $result" -ForegroundColor Green + Write-Host "" + Write-Host "=== Colab CLI Ready ===" -ForegroundColor Green + } else { + Write-Host " $result" -ForegroundColor Red + Write-Host "" + Write-Host "=== Auth needed — see [2/4] above ===" -ForegroundColor Yellow + } +} + +Write-Host "" +Write-Host "Quick commands:" -ForegroundColor Cyan +Write-Host " colab --auth=adc new -s --gpu T4" -ForegroundColor White +Write-Host " colab --auth=adc exec -s -f script.py" -ForegroundColor White +Write-Host " colab --auth=adc upload -s local.file /content/" -ForegroundColor White +Write-Host " colab --auth=adc stop -s " -ForegroundColor White diff --git a/src/colab_cli/_terminal.py b/src/colab_cli/_terminal.py new file mode 100644 index 0000000..43fb1d5 --- /dev/null +++ b/src/colab_cli/_terminal.py @@ -0,0 +1,180 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Platform abstraction for terminal raw-mode handling. + +Provides a uniform API for putting a terminal into raw (character-at-a-time, +no-echo) mode and restoring it afterwards. On Linux/macOS it delegates to +``termios`` + ``tty``; on Windows it uses the Console API via ``ctypes``. +""" + +import logging +import os +import threading +import time + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def get_fd() -> int | None: + """Return the file descriptor for stdin if it is a TTY, else None.""" + if not os.isatty(0): + return None + return _get_fd() + + +def set_raw(fd: int): + """Put the terminal referenced by *fd* into raw mode. + + Returns an opaque *old_state* object that must be passed to + :func:`restore` when raw mode is no longer needed. + """ + return _set_raw(fd) + + +def restore(fd: int, old_state) -> None: + """Restore the terminal to the settings captured by :func:`set_raw`.""" + _restore(fd, old_state) + + +def register_resize_handler(callback) -> None: + """Register *callback* to be invoked when the terminal window is resized. + + The callback receives no arguments and should read the new size via + :func:`os.get_terminal_size`. + """ + _register_resize_handler(callback) + + +def unregister_resize_handler() -> None: + """Remove any resize handler registered by :func:`register_resize_handler`.""" + _unregister_resize_handler() + + +# --------------------------------------------------------------------------- +# Windows implementation (ctypes + msvcrt) +# --------------------------------------------------------------------------- + +if os.name == "nt": + import msvcrt + from ctypes import c_ulong, byref, windll, WINFUNCTYPE + + kernel32 = windll.kernel32 + + # Console mode flags + _ENABLE_PROCESSED_INPUT = 0x0001 + _ENABLE_LINE_INPUT = 0x0002 + _ENABLE_ECHO_INPUT = 0x0004 + _ENABLE_VIRTUAL_TERMINAL_INPUT = 0x0200 + + _resize_thread = None + _resize_stop = None + + def _get_handle(fd: int): + return msvcrt.get_osfhandle(fd) + + def _get_fd() -> int | None: + return 0 # stdin + + def _set_raw(fd: int): + handle = _get_handle(fd) + mode = c_ulong() + kernel32.GetConsoleMode(handle, byref(mode)) + old_mode = mode.value + + # Disable processed input (Ctrl+C handling), line input, and echo. + # Enable virtual terminal input so ANSI escape sequences from the + # remote TTY pass through. + new_mode = ( + old_mode + & ~_ENABLE_PROCESSED_INPUT + & ~_ENABLE_LINE_INPUT + & ~_ENABLE_ECHO_INPUT + | _ENABLE_VIRTUAL_TERMINAL_INPUT + ) + kernel32.SetConsoleMode(handle, new_mode) + return old_mode + + def _restore(fd: int, old_mode) -> None: + handle = _get_handle(fd) + kernel32.SetConsoleMode(handle, old_mode) + + def _resize_poll_loop(interval: float, callback): + """Background thread that polls terminal size and calls *callback* on change.""" + last = None + while not _resize_stop.is_set(): + try: + current = os.get_terminal_size() + if last is not None and current != last: + try: + callback() + except Exception: + logger.debug("Resize callback failed", exc_info=True) + last = current + except Exception: + pass + _resize_stop.wait(interval) + + def _register_resize_handler(callback) -> None: + global _resize_thread, _resize_stop + _unregister_resize_handler() + _resize_stop = threading.Event() + _resize_thread = threading.Thread( + target=_resize_poll_loop, + args=(0.5, callback), + daemon=True, + ) + _resize_thread.start() + + def _unregister_resize_handler() -> None: + global _resize_thread, _resize_stop + if _resize_stop is not None: + _resize_stop.set() + if _resize_thread is not None: + _resize_thread.join(timeout=1.0) + _resize_thread = None + _resize_stop = None + +# --------------------------------------------------------------------------- +# Unix implementation (termios + tty) +# --------------------------------------------------------------------------- + +else: + import signal + import termios + import tty + + def _get_fd() -> int | None: + return 0 # stdin + + def _set_raw(fd: int): + old = termios.tcgetattr(fd) + tty.setraw(fd, termios.TCSANOW) + return old + + def _restore(fd: int, old) -> None: + termios.tcsetattr(fd, termios.TCSANOW, old) + + def _register_resize_handler(callback) -> None: + # Wrap so we swallow the signum/frame arguments the callback doesn't need. + def handler(signum, frame): + callback() + + signal.signal(signal.SIGWINCH, handler) + + def _unregister_resize_handler() -> None: + signal.signal(signal.SIGWINCH, signal.SIG_DFL) diff --git a/src/colab_cli/commands/automation.py b/src/colab_cli/commands/automation.py index 18c02a5..ad0718c 100644 --- a/src/colab_cli/commands/automation.py +++ b/src/colab_cli/commands/automation.py @@ -102,7 +102,8 @@ def drivefs_hook(deserialize_msg, wsclient): state.history.log_event(s.name, "drive_auth_needed", {"uri": uri}) sys.stdout.write("Press Enter after you have granted access... ") sys.stdout.flush() - with open("/dev/tty") as tty: + tty_path = "CON" if os.name == "nt" else "/dev/tty" + with open(tty_path) as tty: tty.readline() typer.echo("[colab] Authorizing VM...") diff --git a/src/colab_cli/console.py b/src/colab_cli/console.py index 0766e1d..6b7c123 100644 --- a/src/colab_cli/console.py +++ b/src/colab_cli/console.py @@ -15,16 +15,14 @@ import json import logging import os -import signal import sys -import termios import threading import time -import tty from urllib.parse import urlparse import websocket +from colab_cli import _terminal from colab_cli.state import SessionState logger = logging.getLogger(__name__) @@ -133,8 +131,8 @@ def connect_console(session: SessionState): ws_url = f"{ws_scheme}://{parsed.netloc}/colab/tty?colab-runtime-proxy-token={session.token}" is_tty = sys.stdin.isatty() - fd = sys.stdin.fileno() if is_tty else None - old_settings = termios.tcgetattr(fd) if is_tty else None + fd = _terminal.get_fd() if is_tty else None + old_settings = None ws = websocket.WebSocketApp( url=ws_url, @@ -144,15 +142,15 @@ def connect_console(session: SessionState): on_close=on_close, ) - def handle_sigwinch(signum, frame): + def handle_resize(): """Handle window resize events.""" if _is_running: send_terminal_size(ws) try: - if is_tty: - tty.setraw(fd, termios.TCSANOW) - signal.signal(signal.SIGWINCH, handle_sigwinch) + if is_tty and fd is not None: + old_settings = _terminal.set_raw(fd) + _terminal.register_resize_handler(handle_resize) # This is a blocking call until the connection is closed ws.run_forever() @@ -164,9 +162,9 @@ def handle_sigwinch(signum, frame): # We raise a standard exception that the caller can recognize raise RuntimeError(f"Connection failed: {err_msg}") finally: - if is_tty: + if is_tty and fd is not None and old_settings is not None: # Always ensure the terminal is restored to its original state - termios.tcsetattr(fd, termios.TCSANOW, old_settings) - # Restore the default signal handler for resize - signal.signal(signal.SIGWINCH, signal.SIG_DFL) + _terminal.restore(fd, old_settings) + # Stop the resize handler + _terminal.unregister_resize_handler() print("\r\nConnection closed.") diff --git a/tests/test_console.py b/tests/test_console.py index 310e98e..7638e76 100644 --- a/tests/test_console.py +++ b/tests/test_console.py @@ -15,7 +15,6 @@ import json import os import sys -import termios from unittest.mock import MagicMock, patch from colab_cli.console import connect_console, on_message, on_open @@ -34,27 +33,29 @@ def mock_session(): @patch("colab_cli.console.websocket.WebSocketApp") -@patch("colab_cli.console.tty.setraw") -@patch("colab_cli.console.termios.tcgetattr") -@patch("colab_cli.console.termios.tcsetattr") +@patch("colab_cli.console._terminal.unregister_resize_handler") +@patch("colab_cli.console._terminal.register_resize_handler") +@patch("colab_cli.console._terminal.restore") +@patch("colab_cli.console._terminal.set_raw") +@patch("colab_cli.console._terminal.get_fd") @patch("colab_cli.console.os.get_terminal_size") -@patch("colab_cli.console.sys.stdin.fileno") @patch("colab_cli.console.sys.stdin.isatty") def test_console_initialization( mock_isatty, - mock_fileno, mock_get_term_size, - mock_tcsetattr, - mock_tcgetattr, - mock_setraw, + mock_get_fd, + mock_set_raw, + mock_restore, + mock_register_resize, + mock_unregister_resize, mock_ws_app, mock_session, ): # Setup mocks mock_isatty.return_value = True - mock_fileno.return_value = 0 + mock_get_fd.return_value = 0 mock_get_term_size.return_value = os.terminal_size((80, 24)) - mock_tcgetattr.return_value = ["fake_attrs"] + mock_set_raw.return_value = 12345 # opaque old settings token mock_ws_instance = MagicMock() mock_ws_app.return_value = mock_ws_instance @@ -69,30 +70,31 @@ def test_console_initialization( mock_ws_app.assert_called_once() assert mock_ws_app.call_args[1]["url"] == expected_url - # 2. Verify raw mode setup and teardown - mock_tcgetattr.assert_called_once_with(sys.stdin.fileno()) - mock_setraw.assert_called_once_with(sys.stdin.fileno(), termios.TCSANOW) + # 2. Verify raw mode setup via the platform abstraction + mock_get_fd.assert_called_once() + mock_set_raw.assert_called_once_with(0) + mock_register_resize.assert_called_once() - # Teardown should happen in a finally block - mock_tcsetattr.assert_called_once_with( - sys.stdin.fileno(), termios.TCSANOW, ["fake_attrs"] - ) + # 3. Teardown should happen in a finally block + mock_restore.assert_called_once_with(0, 12345) + mock_unregister_resize.assert_called_once() @patch("colab_cli.console.websocket.WebSocketApp") -@patch("colab_cli.console.tty.setraw") -@patch("colab_cli.console.termios.tcgetattr") -@patch("colab_cli.console.termios.tcsetattr") +@patch("colab_cli.console._terminal.set_raw") +@patch("colab_cli.console._terminal.get_fd") +@patch("colab_cli.console._terminal.restore") @patch("colab_cli.console.sys.stdin.isatty") def test_console_piped_input( mock_isatty, - mock_tcsetattr, - mock_tcgetattr, - mock_setraw, + mock_restore, + mock_get_fd, + mock_set_raw, mock_ws_app, mock_session, ): mock_isatty.return_value = False + mock_get_fd.return_value = None mock_ws_instance = MagicMock() mock_ws_app.return_value = mock_ws_instance mock_ws_instance.run_forever.return_value = None @@ -100,10 +102,9 @@ def test_console_piped_input( with patch("colab_cli.console.threading.Thread"): connect_console(mock_session) - # In a piped environment, we should not attempt to use termios or tty - mock_tcgetattr.assert_not_called() - mock_setraw.assert_not_called() - mock_tcsetattr.assert_not_called() + # In a piped environment, we should not attempt to use terminal raw mode + mock_set_raw.assert_not_called() + mock_restore.assert_not_called() @patch("colab_cli.console.os.get_terminal_size") @@ -194,7 +195,7 @@ def test_read_stdin_eof_tty_does_not_close_ws( ): """When stdin is a real TTY and read() returns empty (which happens on Ctrl-D in raw mode), we should NOT inject 'exit\\n' or close the websocket - \u2014 the user is in interactive mode and may have intended Ctrl-D as a literal + — the user is in interactive mode and may have intended Ctrl-D as a literal char. The websocket lifecycle is owned by the remote shell in this case. """ import colab_cli.console as console_mod