| 1 | """`setup --store-key <NAME>` (U6): persist one allowlisted credential from stdin. |
| 2 | |
| 3 | This is the only way the Grok Bot flow persists X_BEARER_TOKEN: the model |
| 4 | pipes the value on stdin, the engine appends it to the global .env as a 0o600 |
| 5 | secret through setup_wizard.write_api_key, and stdout carries only the masked |
| 6 | name plus a JSON status line. The value never reaches stdout or stderr; a name |
| 7 | outside env.KEYCHAIN_KEYS or an empty value exits 2 without echoing anything. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | import io |
| 13 | import json |
| 14 | import os |
| 15 | import stat |
| 16 | import sys |
| 17 | from contextlib import redirect_stderr, redirect_stdout |
| 18 | from pathlib import Path |
| 19 | from unittest import mock |
| 20 | |
| 21 | import pytest |
| 22 | |
| 23 | import last30days as cli |
| 24 | from lib import env |
| 25 | |
| 26 | DUMMY = "dummy-x-bearer-token-not-real-0123456789" |
| 27 | |
| 28 | |
| 29 | def _run(argv: list[str], stdin_text: str, env_path: Path) -> tuple[int, str, str]: |
| 30 | stdout, stderr = io.StringIO(), io.StringIO() |
| 31 | with ( |
| 32 | mock.patch.object(cli.env, "CONFIG_FILE", env_path), |
| 33 | mock.patch.object(cli.env, "get_config", side_effect=AssertionError("store-key must not load config")), |
| 34 | mock.patch.object(sys, "stdin", io.StringIO(stdin_text)), |
| 35 | mock.patch.object(sys, "argv", ["last30days.py"] + argv), |
| 36 | ): |
| 37 | with redirect_stdout(stdout), redirect_stderr(stderr): |
| 38 | try: |
| 39 | rc = cli.main() |
| 40 | except SystemExit as exc: # argparse parser.error path |
| 41 | rc = int(exc.code or 0) |
| 42 | return rc, stdout.getvalue(), stderr.getvalue() |
| 43 | |
| 44 | |
| 45 | def _mode(path: Path) -> int: |
| 46 | return stat.S_IMODE(path.stat().st_mode) |
| 47 | |
| 48 | |
| 49 | @pytest.mark.parametrize("argv", [ |
| 50 | ["setup", "--store-key", "X_BEARER_TOKEN"], |
| 51 | ["setup", "--store-key=X_BEARER_TOKEN"], |
| 52 | ]) |
| 53 | def test_store_key_writes_allowlisted_key_at_0600_and_masks_stdout(tmp_path, argv): |
| 54 | env_path = tmp_path / "cfg" / ".env" |
| 55 | rc, out, err = _run(argv, DUMMY + "\n", env_path) |
| 56 | |
| 57 | assert rc == 0, (out, err) |
| 58 | assert env_path.exists() |
| 59 | assert _mode(env_path) == 0o600 |
| 60 | assert env.load_env_file(env_path)["X_BEARER_TOKEN"] == DUMMY |
| 61 | |
| 62 | lines = out.strip().splitlines() |
| 63 | assert lines[0] == "X_BEARER_TOKEN=****" |
| 64 | assert json.loads(lines[-1]) == {"persisted": True, "key": "X_BEARER_TOKEN"} |
| 65 | assert DUMMY not in out |
| 66 | assert DUMMY not in err |
| 67 | |
| 68 | |
| 69 | def test_store_key_strips_surrounding_whitespace(tmp_path): |
| 70 | env_path = tmp_path / ".env" |
| 71 | rc, out, err = _run(["setup", "--store-key", "X_BEARER_TOKEN"], f" {DUMMY}\t\r\n", env_path) |
| 72 | assert rc == 0 |
| 73 | assert env.load_env_file(env_path)["X_BEARER_TOKEN"] == DUMMY |
| 74 | assert DUMMY not in out and DUMMY not in err |
| 75 | |
| 76 | |
| 77 | def test_store_key_reads_exactly_one_line(tmp_path): |
| 78 | env_path = tmp_path / ".env" |
| 79 | rc, out, err = _run( |
| 80 | ["setup", "--store-key", "X_BEARER_TOKEN"], |
| 81 | DUMMY + "\nsecond-line-must-be-ignored\n", |
| 82 | env_path, |
| 83 | ) |
| 84 | assert rc == 0 |
| 85 | assert env.load_env_file(env_path)["X_BEARER_TOKEN"] == DUMMY |
| 86 | assert "second-line" not in env_path.read_text() |
| 87 | |
| 88 | |
| 89 | def test_store_key_replaces_an_existing_value_in_place(tmp_path): |
| 90 | """A second store-key rotates the credential: a rejected token must not |
| 91 | survive a 'persisted: true' receipt (review finding).""" |
| 92 | env_path = tmp_path / ".env" |
| 93 | env_path.write_text("SETUP_COMPLETE=true\n# note\nX_BEARER_TOKEN=old-dummy\nXAI_API_KEY=other\n") |
| 94 | os.chmod(env_path, 0o600) |
| 95 | rotated = "rotated-dummy-value-not-real" |
| 96 | rc, out, err = _run(["setup", "--store-key", "X_BEARER_TOKEN"], rotated + "\n", env_path) |
| 97 | assert rc == 0 |
| 98 | assert json.loads(out.strip().splitlines()[-1]) == {"persisted": True, "key": "X_BEARER_TOKEN"} |
| 99 | content = env_path.read_text() |
| 100 | assert content.count("X_BEARER_TOKEN=") == 1 |
| 101 | assert "old-dummy" not in content |
| 102 | loaded = env.load_env_file(env_path) |
| 103 | assert loaded["X_BEARER_TOKEN"] == rotated |
| 104 | assert loaded["XAI_API_KEY"] == "other" |
| 105 | assert loaded["SETUP_COMPLETE"] == "true" |
| 106 | assert "# note" in content |
| 107 | assert _mode(env_path) == 0o600 |
| 108 | assert not (tmp_path / ".env.tmp").exists() |
| 109 | assert rotated not in out and rotated not in err |
| 110 | |
| 111 | |
| 112 | def test_write_api_key_default_still_keeps_an_existing_value(tmp_path): |
| 113 | from lib import setup_wizard |
| 114 | env_path = tmp_path / ".env" |
| 115 | assert setup_wizard.write_api_key(env_path, DUMMY, key_name="X_BEARER_TOKEN") |
| 116 | assert setup_wizard.write_api_key(env_path, "other-dummy", key_name="X_BEARER_TOKEN") |
| 117 | assert env.load_env_file(env_path)["X_BEARER_TOKEN"] == DUMMY |
| 118 | |
| 119 | |
| 120 | def test_store_key_reads_a_bounded_line(tmp_path): |
| 121 | env_path = tmp_path / ".env" |
| 122 | huge = "a" * (cli.STORE_KEY_MAX_BYTES * 2) |
| 123 | rc, out, err = _run(["setup", "--store-key", "X_BEARER_TOKEN"], huge + "\n", env_path) |
| 124 | assert rc == 0 |
| 125 | assert len(env.load_env_file(env_path)["X_BEARER_TOKEN"]) == cli.STORE_KEY_MAX_BYTES |
| 126 | |
| 127 | |
| 128 | def test_store_key_tightens_a_loose_existing_file(tmp_path): |
| 129 | env_path = tmp_path / ".env" |
| 130 | env_path.write_text("SETUP_COMPLETE=true\n") |
| 131 | os.chmod(env_path, 0o644) |
| 132 | rc, _, _ = _run(["setup", "--store-key", "X_BEARER_TOKEN"], DUMMY + "\n", env_path) |
| 133 | assert rc == 0 |
| 134 | assert _mode(env_path) == 0o600 |
| 135 | loaded = env.load_env_file(env_path) |
| 136 | assert loaded["SETUP_COMPLETE"] == "true" |
| 137 | assert loaded["X_BEARER_TOKEN"] == DUMMY |
| 138 | |
| 139 | |
| 140 | @pytest.mark.parametrize("name", ["NOT_A_REAL_KEY", "PATH", "x_bearer_token", "X_BEARER_TOKEN=evil"]) |
| 141 | def test_store_key_rejects_name_outside_allowlist_with_exit_2(tmp_path, name): |
| 142 | env_path = tmp_path / ".env" |
| 143 | rc, out, err = _run(["setup", "--store-key", name], DUMMY + "\n", env_path) |
| 144 | assert rc == 2 |
| 145 | assert not env_path.exists() |
| 146 | assert DUMMY not in out and DUMMY not in err |
| 147 | assert "store-key" in err |
| 148 | # R4: the failure hint must not enumerate legacy credential names. |
| 149 | for banned in ("AUTH_TOKEN", "CT0", "XQUIK_API_KEY"): |
| 150 | assert banned not in err |
| 151 | |
| 152 | |
| 153 | def test_store_key_without_name_exits_2(tmp_path): |
| 154 | env_path = tmp_path / ".env" |
| 155 | rc, out, err = _run(["setup", "--store-key"], DUMMY + "\n", env_path) |
| 156 | assert rc == 2 |
| 157 | assert not env_path.exists() |
| 158 | assert DUMMY not in out and DUMMY not in err |
| 159 | |
| 160 | |
| 161 | @pytest.mark.parametrize("stdin_text", ["", "\n", " \n"]) |
| 162 | def test_store_key_rejects_empty_value_with_exit_2(tmp_path, stdin_text): |
| 163 | env_path = tmp_path / ".env" |
| 164 | rc, out, err = _run(["setup", "--store-key", "X_BEARER_TOKEN"], stdin_text, env_path) |
| 165 | assert rc == 2 |
| 166 | assert not env_path.exists() |
| 167 | assert "empty" in err.lower() |
| 168 | assert "X_BEARER_TOKEN" in err |
| 169 | |
| 170 | |
| 171 | def test_store_key_works_for_every_allowlisted_name(tmp_path): |
| 172 | """The allowlist is env.KEYCHAIN_KEYS, not a bearer-only special case.""" |
| 173 | env_path = tmp_path / ".env" |
| 174 | for name in env.KEYCHAIN_KEYS: |
| 175 | rc, out, _ = _run(["setup", "--store-key", name], f"dummy-{name.lower()}-not-real\n", env_path) |
| 176 | assert rc == 0, name |
| 177 | assert out.strip().splitlines()[0] == f"{name}=****" |
| 178 | loaded = env.load_env_file(env_path) |
| 179 | assert set(env.KEYCHAIN_KEYS) <= set(loaded) |
| 180 | assert _mode(env_path) == 0o600 |
| 181 | |
| 182 | |
| 183 | def test_store_key_is_a_declared_setup_passthrough_flag(): |
| 184 | assert "--store-key" in cli.SETUP_PASSTHROUGH_FLAGS |
| 185 | |
| 186 | |
| 187 | def test_store_key_reports_persist_failure_as_false(tmp_path): |
| 188 | env_path = tmp_path / ".env" |
| 189 | with mock.patch("lib.setup_wizard.write_api_key", return_value=False) as w: |
| 190 | rc, out, err = _run(["setup", "--store-key", "X_BEARER_TOKEN"], DUMMY + "\n", env_path) |
| 191 | assert rc == 1 |
| 192 | w.assert_called_once_with(env_path, DUMMY, key_name="X_BEARER_TOKEN", replace=True) |
| 193 | lines = out.strip().splitlines() |
| 194 | assert lines[0] == "X_BEARER_TOKEN=****" |
| 195 | assert json.loads(lines[-1]) == {"persisted": False, "key": "X_BEARER_TOKEN"} |
| 196 | assert DUMMY not in out and DUMMY not in err |
| 197 |