diff --git a/bot_bottle/cli/__init__.py b/bot_bottle/cli/__init__.py index ae421fc..cf62f68 100644 --- a/bot_bottle/cli/__init__.py +++ b/bot_bottle/cli/__init__.py @@ -19,6 +19,7 @@ from .commit import cmd_commit from .edit import cmd_edit from .info import cmd_info from .init import cmd_init +from .login import cmd_login from .resume import cmd_resume from .start import cmd_start from .supervise import cmd_supervise @@ -33,6 +34,7 @@ COMMANDS = { "info": cmd_info, "init": cmd_init, "list": cmd_list, + "login": cmd_login, "resume": cmd_resume, "start": cmd_start, "supervise": cmd_supervise, @@ -43,7 +45,7 @@ COMMANDS = { # the host (TAP pool, /dev/kvm, firecracker) and never opens the store, so # gating it on the schema breaks preflight on a fresh CI runner where stdin # isn't a TTY and the migration prompt can't be answered. -NO_MIGRATION_COMMANDS = frozenset({"backend"}) +NO_MIGRATION_COMMANDS = frozenset({"backend", "login"}) def usage() -> None: @@ -56,6 +58,7 @@ def usage() -> None: sys.stderr.write(" info print env, skills, and prompt details for a named agent\n") sys.stderr.write(" init interactively create a new agent and add it to bot-bottle.json\n") sys.stderr.write(" list list available agents or active containers\n") + sys.stderr.write(" login register this host with a bot-bottle console\n") sys.stderr.write( " resume re-launch a bottle by its identity " "(continues state from PRD 0016)\n" @@ -111,7 +114,3 @@ def main(argv: list[str] | None = None) -> int: return e.code if isinstance(e.code, int) else 1 except KeyboardInterrupt: return 130 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/bot_bottle/cli/__main__.py b/bot_bottle/cli/__main__.py new file mode 100644 index 0000000..3cf5f19 --- /dev/null +++ b/bot_bottle/cli/__main__.py @@ -0,0 +1,15 @@ +"""Entry point for `python -m bot_bottle.cli`. + +`cli.py` at the repo root is the usual way in; this makes the package +runnable too, so the CLI works from an installed copy where there is no +`cli.py` on disk to point at. +""" + +from __future__ import annotations + +import sys + +from . import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/bot_bottle/cli/login.py b/bot_bottle/cli/login.py new file mode 100644 index 0000000..3d00f17 --- /dev/null +++ b/bot_bottle/cli/login.py @@ -0,0 +1,168 @@ +"""bb login — register this host with a bot-bottle console. + +Opens a device-authorization flow against the target console, waits for the +operator to approve, then writes access and refresh tokens to +~/.bot-bottle/console.json (or $BOT_BOTTLE_ROOT/console.json). + +Usage: + bb login [--console-url URL] [--label LABEL] + +Flags: + --console-url URL Target console URL (overrides BB_CONSOLE_URL env var) + --label LABEL Host label shown in the console (default: hostname) +""" + +from __future__ import annotations + +import json +import os +import socket +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +from ..paths import bot_bottle_root + +_CONSOLE_URL_ENV = "BB_CONSOLE_URL" +_POLL_SLEEP = 2 # seconds between polls; matches console's poll_interval default + + +def _usage() -> None: + sys.stderr.write( + "usage: bb login [--console-url URL] [--label LABEL]\n" + "\n" + "Options:\n" + " --console-url URL Console base URL (or BB_CONSOLE_URL env var)\n" + " --label LABEL Host label shown in the console (default: hostname)\n" + ) + + +def _flag(argv: list[str], name: str) -> str | None: + for i, arg in enumerate(argv): + if arg == name and i + 1 < len(argv): + return argv[i + 1] + if arg.startswith(f"{name}="): + return arg[len(name) + 1:] + return None + + +def _post(url: str, payload: dict[str, Any]) -> dict[str, Any]: + data = json.dumps(payload).encode() + req = urllib.request.Request( + url, data=data, headers={"Content-Type": "application/json"} + ) + with urllib.request.urlopen(req, timeout=10) as resp: + return json.loads(resp.read()) + + +def _get(url: str) -> tuple[int, dict[str, Any]]: + req = urllib.request.Request(url) + try: + with urllib.request.urlopen(req, timeout=10) as resp: + return resp.status, json.loads(resp.read()) + except urllib.error.HTTPError as e: + return e.code, {} + + +def _save_credentials( + console_url: str, host_id: str, access_token: str, refresh_token: str +) -> Path: + path = bot_bottle_root() / "console.json" + path.parent.mkdir(parents=True, exist_ok=True) + content = ( + json.dumps( + { + "url": console_url, + "host_id": host_id, + "access_token": access_token, + "refresh_token": refresh_token, + }, + indent=2, + ) + + "\n" + ) + fd, tmp_path_str = tempfile.mkstemp(dir=path.parent, prefix=".console-") + tmp = Path(tmp_path_str) + try: + tmp.chmod(0o600) + with os.fdopen(fd, "w") as f: + f.write(content) + os.replace(tmp, path) + except OSError: + try: + tmp.unlink() + except OSError: + pass + raise + return path + + +def cmd_login(argv: list[str]) -> int: + if "--help" in argv or "-h" in argv: + _usage() + return 0 + + console_url = _flag(argv, "--console-url") or os.environ.get(_CONSOLE_URL_ENV) + if not console_url: + sys.stderr.write( + "bb login: --console-url or BB_CONSOLE_URL is required\n" + ) + return 1 + console_url = console_url.rstrip("/") + + label = _flag(argv, "--label") or socket.gethostname() + + try: + resp = _post(f"{console_url}/api/v1/hosts/authorize", {"label": label}) + except (OSError, ValueError) as exc: + sys.stderr.write(f"bb login: failed to start authorization: {exc}\n") + return 1 + + device_code = resp["device_code"] + user_code = resp["user_code"] + expires_in = resp.get("expires_in", 300) + poll_sleep = max(1, min(int(resp.get("poll_interval", _POLL_SLEEP)), 60)) + + sys.stderr.write( + f"\nOpen this URL in your browser to authorize this host:\n\n" + f" {console_url}/hosts/authorize?code={user_code}\n\n" + f"Waiting for approval" + ) + + deadline = time.monotonic() + expires_in + while time.monotonic() < deadline: + sys.stderr.write(".") + sys.stderr.flush() + time.sleep(poll_sleep) + + try: + code, result = _get( + f"{console_url}/api/v1/hosts/authorize/{device_code}" + ) + except (OSError, ValueError): + continue + + if code == 410: + break + + st = result.get("status") + if st == "approved": + sys.stderr.write("\n\nApproved.\n") + path = _save_credentials( + console_url, + result["host_id"], + result["access_token"], + result["refresh_token"], + ) + sys.stderr.write(f"Credentials saved to {path}\n") + return 0 + if st == "denied": + sys.stderr.write("\n\nDenied by operator.\n") + return 1 + + sys.stderr.write("\n\nAuthorization timed out.\n") + return 1 diff --git a/tests/unit/test_cli_login.py b/tests/unit/test_cli_login.py new file mode 100644 index 0000000..0d4c13c --- /dev/null +++ b/tests/unit/test_cli_login.py @@ -0,0 +1,268 @@ +"""Unit tests for bb login command.""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +import urllib.error +from email.message import Message +from typing import Any +from unittest.mock import MagicMock, patch + + +class TestFlagParsing(unittest.TestCase): + def test_console_url_flag(self) -> None: + from bot_bottle.cli.login import _flag + self.assertEqual(_flag(["--console-url", "http://x"], "--console-url"), "http://x") + + def test_console_url_equals_form(self) -> None: + from bot_bottle.cli.login import _flag + self.assertEqual( + _flag(["--console-url=http://x"], "--console-url"), "http://x" + ) + + def test_label_flag(self) -> None: + from bot_bottle.cli.login import _flag + self.assertEqual(_flag(["--label", "my-mac"], "--label"), "my-mac") + + def test_missing_flag_returns_none(self) -> None: + from bot_bottle.cli.login import _flag + self.assertIsNone(_flag([], "--console-url")) + + +class TestHttpHelpers(unittest.TestCase): + def test_post_sends_json_and_decodes_response(self) -> None: + from bot_bottle.cli.login import _post + + response = MagicMock() + response.__enter__.return_value.read.return_value = b'{"ok": true}' + with patch("urllib.request.urlopen", return_value=response) as urlopen: + self.assertEqual( + _post("http://console/start", {"label": "host"}), {"ok": True} + ) + + request = urlopen.call_args.args[0] + self.assertEqual(request.data, b'{"label": "host"}') + self.assertEqual(request.get_header("Content-type"), "application/json") + + def test_get_decodes_success_response(self) -> None: + from bot_bottle.cli.login import _get + + response = MagicMock() + response.__enter__.return_value.status = 200 + response.__enter__.return_value.read.return_value = b'{"status": "pending"}' + with patch("urllib.request.urlopen", return_value=response): + self.assertEqual( + _get("http://console/status"), (200, {"status": "pending"}) + ) + + def test_get_returns_http_error_status(self) -> None: + from bot_bottle.cli.login import _get + + error = urllib.error.HTTPError( + "http://console/status", 410, "gone", Message(), None + ) + with patch("urllib.request.urlopen", side_effect=error): + self.assertEqual(_get("http://console/status"), (410, {})) + + +class TestSaveCredentials(unittest.TestCase): + def test_writes_json_and_sets_perms(self) -> None: + from bot_bottle.cli.login import _save_credentials + + with tempfile.TemporaryDirectory() as tmp: + with patch.dict(os.environ, {"BOT_BOTTLE_ROOT": tmp}): + path = _save_credentials("http://c", "hid", "at", "rt") + self.assertTrue(path.exists()) + data = json.loads(path.read_text()) + self.assertEqual(data["url"], "http://c") + self.assertEqual(data["host_id"], "hid") + self.assertEqual(data["access_token"], "at") + self.assertEqual(data["refresh_token"], "rt") + self.assertEqual(oct(path.stat().st_mode & 0o777), oct(0o600)) + + def test_temp_file_is_private_before_replace(self) -> None: + """Temp file must be 0600 at the moment os.replace is called.""" + from bot_bottle.cli.login import _save_credentials + from pathlib import Path as _Path + + tmp_perms_at_replace: list[int] = [] + real_replace = os.replace + + def _spy_replace( + src: str | os.PathLike[str], dst: str | os.PathLike[str] + ) -> None: + tmp_perms_at_replace.append(_Path(src).stat().st_mode & 0o777) + real_replace(src, dst) + + with tempfile.TemporaryDirectory() as tmp: + with patch.dict(os.environ, {"BOT_BOTTLE_ROOT": tmp}): + with patch("os.replace", side_effect=_spy_replace): + path = _save_credentials("http://c", "hid", "at", "rt") + self.assertEqual(len(tmp_perms_at_replace), 1) + self.assertEqual(oct(tmp_perms_at_replace[0]), oct(0o600)) + self.assertEqual(oct(path.stat().st_mode & 0o777), oct(0o600)) + + def test_cleanup_on_write_failure(self) -> None: + """Temp file is removed and no credentials remain if replace fails.""" + from bot_bottle.cli.login import _save_credentials + + with tempfile.TemporaryDirectory() as tmp: + with patch.dict(os.environ, {"BOT_BOTTLE_ROOT": tmp}): + with patch("os.replace", side_effect=OSError("disk full")): + with self.assertRaises(OSError): + _save_credentials("http://c", "hid", "at", "rt") + leftovers = [f for f in os.listdir(tmp) if f.startswith(".console-")] + self.assertEqual(leftovers, []) + + +class TestCmdLoginMissingUrl(unittest.TestCase): + def test_help_returns_0(self) -> None: + from bot_bottle.cli.login import cmd_login + + self.assertEqual(cmd_login(["--help"]), 0) + + def test_returns_1_without_url(self) -> None: + from bot_bottle.cli.login import cmd_login + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("BB_CONSOLE_URL", None) + result = cmd_login([]) + self.assertEqual(result, 1) + + def test_reads_env_var(self) -> None: + """Exits 1 (network error) not because of missing URL when env var is set.""" + from bot_bottle.cli.login import cmd_login + + def _fail_post(_url: str, _payload: dict[str, Any]) -> dict[str, Any]: + raise OSError("connection refused") + + with tempfile.TemporaryDirectory() as tmp: + with patch.dict( + os.environ, + { + "BB_CONSOLE_URL": "http://localhost:9999", + "BOT_BOTTLE_ROOT": tmp, + }, + ): + with patch("bot_bottle.cli.login._post", side_effect=_fail_post): + result = cmd_login([]) + self.assertEqual(result, 1) + + +class TestCmdLoginFlow(unittest.TestCase): + def _run_with_mocks( + self, poll_responses: list[dict[str, Any]], tmp: str + ) -> int: + from bot_bottle.cli.login import cmd_login + + start_resp = { + "device_code": "dc123", + "user_code": "ABC-DEF", + "expires_in": 300, + "poll_interval": 0, + } + + poll_iter = iter(poll_responses) + + def _fake_post( + _url: str, _payload: dict[str, Any] + ) -> dict[str, Any]: + return start_resp + + def _fake_get(_url: str) -> tuple[int, dict[str, Any]]: + return 200, next(poll_iter, {"status": "pending"}) + + with patch.dict(os.environ, {"BOT_BOTTLE_ROOT": tmp}): + with patch("bot_bottle.cli.login._post", side_effect=_fake_post): + with patch("bot_bottle.cli.login._get", side_effect=_fake_get): + with patch("time.sleep"): + return cmd_login(["--console-url", "http://console"]) + + def test_approved_flow_returns_0(self) -> None: + approved = { + "status": "approved", + "host_id": "hid", + "access_token": "at", + "refresh_token": "rt", + } + with tempfile.TemporaryDirectory() as tmp: + result = self._run_with_mocks( + [{"status": "pending"}, approved], tmp + ) + self.assertEqual(result, 0) + with open(os.path.join(tmp, "console.json"), encoding="utf-8") as f: + creds = json.loads(f.read()) + self.assertEqual(creds["host_id"], "hid") + + def test_denied_flow_returns_1(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + result = self._run_with_mocks([{"status": "denied"}], tmp) + self.assertEqual(result, 1) + + def test_timeout_returns_1(self) -> None: + from bot_bottle.cli.login import cmd_login + + start_resp = { + "device_code": "dc", + "user_code": "ZZZ-ZZZ", + "expires_in": 0, # already expired; loop never runs + "poll_interval": 2, + } + + with tempfile.TemporaryDirectory() as tmp: + with patch.dict(os.environ, {"BOT_BOTTLE_ROOT": tmp}): + with patch("bot_bottle.cli.login._post", return_value=start_resp): + result = cmd_login(["--console-url", "http://console"]) + self.assertEqual(result, 1) + + def test_poll_interval_from_server_is_used(self) -> None: + """time.sleep must be called with the server-provided poll_interval.""" + from bot_bottle.cli.login import cmd_login + + server_interval = 7 + start_resp = { + "device_code": "dc", + "user_code": "ABC-DEF", + "expires_in": 300, + "poll_interval": server_interval, + } + approved = { + "status": "approved", + "host_id": "hid", + "access_token": "at", + "refresh_token": "rt", + } + poll_iter = iter([{"status": "pending"}, approved]) + + def _fake_get(_url: str) -> tuple[int, dict[str, str]]: + return 200, next(poll_iter) + + with tempfile.TemporaryDirectory() as tmp: + with patch.dict(os.environ, {"BOT_BOTTLE_ROOT": tmp}): + with patch("bot_bottle.cli.login._post", return_value=start_resp): + with patch( + "bot_bottle.cli.login._get", side_effect=_fake_get + ): + with patch("time.sleep") as mock_sleep: + result = cmd_login(["--console-url", "http://console"]) + + self.assertEqual(result, 0) + self.assertTrue(mock_sleep.called) + for call in mock_sleep.call_args_list: + self.assertEqual(call.args[0], server_interval) + + +class TestDispatcherRegistration(unittest.TestCase): + def test_login_in_commands(self) -> None: + from bot_bottle.cli import COMMANDS + self.assertIn("login", COMMANDS) + + def test_login_in_no_migration(self) -> None: + from bot_bottle.cli import NO_MIGRATION_COMMANDS + self.assertIn("login", NO_MIGRATION_COMMANDS) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_cli_module_entry.py b/tests/unit/test_cli_module_entry.py new file mode 100644 index 0000000..2ff003b --- /dev/null +++ b/tests/unit/test_cli_module_entry.py @@ -0,0 +1,40 @@ +"""The CLI package is runnable as `python -m bot_bottle.cli`.""" + +from __future__ import annotations + +import subprocess +import sys +import unittest +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _run(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, "-m", "bot_bottle.cli", *args], + cwd=_REPO_ROOT, + capture_output=True, + text=True, + check=False, + ) + + +class TestModuleEntry(unittest.TestCase): + def test_help_exits_zero(self) -> None: + result = _run("--help") + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("login", result.stderr) + + def test_no_args_prints_usage(self) -> None: + # main() returns 2 with no command, matching the cli.py entry point. + self.assertEqual(_run().returncode, 2) + + def test_subcommand_help_reaches_handler(self) -> None: + result = _run("login", "--help") + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("--console-url", result.stderr) + + +if __name__ == "__main__": + unittest.main()