Compare commits
3 Commits
main
...
67cb51352a
| Author | SHA1 | Date | |
|---|---|---|---|
| 67cb51352a | |||
| 2e1074b659 | |||
| 9b29c77547 |
@@ -19,6 +19,7 @@ from .commit import cmd_commit
|
|||||||
from .edit import cmd_edit
|
from .edit import cmd_edit
|
||||||
from .info import cmd_info
|
from .info import cmd_info
|
||||||
from .init import cmd_init
|
from .init import cmd_init
|
||||||
|
from .login import cmd_login
|
||||||
from .resume import cmd_resume
|
from .resume import cmd_resume
|
||||||
from .start import cmd_start
|
from .start import cmd_start
|
||||||
from .supervise import cmd_supervise
|
from .supervise import cmd_supervise
|
||||||
@@ -33,6 +34,7 @@ COMMANDS = {
|
|||||||
"info": cmd_info,
|
"info": cmd_info,
|
||||||
"init": cmd_init,
|
"init": cmd_init,
|
||||||
"list": cmd_list,
|
"list": cmd_list,
|
||||||
|
"login": cmd_login,
|
||||||
"resume": cmd_resume,
|
"resume": cmd_resume,
|
||||||
"start": cmd_start,
|
"start": cmd_start,
|
||||||
"supervise": cmd_supervise,
|
"supervise": cmd_supervise,
|
||||||
@@ -43,7 +45,7 @@ COMMANDS = {
|
|||||||
# the host (TAP pool, /dev/kvm, firecracker) and never opens the store, so
|
# 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
|
# 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.
|
# 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:
|
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(" 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(" 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(" list list available agents or active containers\n")
|
||||||
|
sys.stderr.write(" login register this host with a bot-bottle console\n")
|
||||||
sys.stderr.write(
|
sys.stderr.write(
|
||||||
" resume re-launch a bottle by its identity "
|
" resume re-launch a bottle by its identity "
|
||||||
"(continues state from PRD 0016)\n"
|
"(continues state from PRD 0016)\n"
|
||||||
|
|||||||
@@ -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 Exception:
|
||||||
|
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 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)
|
||||||
|
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}/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
|
||||||
@@ -0,0 +1,260 @@
|
|||||||
|
"""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))
|
||||||
|
|
||||||
|
|
||||||
|
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]]:
|
||||||
|
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("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")) 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()
|
||||||
Reference in New Issue
Block a user