diff --git a/bot_bottle/cli/__init__.py b/bot_bottle/cli/__init__.py index ae421fc..0061724 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" diff --git a/bot_bottle/cli/login.py b/bot_bottle/cli/login.py new file mode 100644 index 0000000..6410266 --- /dev/null +++ b/bot_bottle/cli/login.py @@ -0,0 +1,153 @@ +"""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 time +import urllib.error +import urllib.request +from pathlib import Path + +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) -> dict: + 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]: + 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) + path.write_text( + json.dumps( + { + "url": console_url, + "host_id": host_id, + "access_token": access_token, + "refresh_token": refresh_token, + }, + indent=2, + ) + + "\n" + ) + path.chmod(0o600) + 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 Exception 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) + + sys.stderr.write( + f"\nOpen this URL in your browser to authorize this host:\n\n" + f" {console_url}/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 Exception: + 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..0cb2c9b --- /dev/null +++ b/tests/unit/test_cli_login.py @@ -0,0 +1,147 @@ +"""Unit tests for bb login command.""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +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 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)) + + +class TestCmdLoginMissingUrl(unittest.TestCase): + 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, payload): + 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], 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, payload): + return start_resp + + def _fake_get(url): + try: + resp = next(poll_iter) + except StopIteration: + return 200, {"status": "pending"} + return 200, resp + + 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("bot_bottle.cli.login._POLL_SLEEP", 0): + 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")) 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 + "poll_interval": 0, + } + + 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._POLL_SLEEP", 0): + result = cmd_login(["--console-url", "http://console"]) + self.assertEqual(result, 1) + + +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()