diff --git a/pyproject.toml b/pyproject.toml index 7ceb59c7..c3f6e74d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ dependencies = [ "typer>=0.12.0", "questionary>=2.0.0", "pathspec>=0.11.0", + "tomlkit>=0.13.0", "tomli>=2.0.0; python_version < '3.11'", ] diff --git a/src/runpod_flash/core/credentials.py b/src/runpod_flash/core/credentials.py index 47350939..a984bf32 100644 --- a/src/runpod_flash/core/credentials.py +++ b/src/runpod_flash/core/credentials.py @@ -12,14 +12,23 @@ from typing import Optional import runpod.cli.groups.config.functions as _runpod_config +import tomlkit from runpod.cli.groups.config.functions import ( get_credentials, - set_credentials, ) log = logging.getLogger(__name__) +# runpodctl writes top-level `apikey`/`apiurl` keys into the same config.toml +# that runpod-python uses for its `[default]` profile. We must preserve those +# (and any other unrelated content) when updating flash's api_key, so flash +# login does not clobber runpodctl's credentials. tomlkit round-trips comments, +# foreign keys, sibling profiles, and line endings, so we only mutate the one +# value we own. +_DEFAULT_SECTION = "default" +_API_KEY_FIELD = "api_key" + _OLD_XDG_PATH = Path.home() / ".config" / "runpod" / "credentials.toml" @@ -50,7 +59,11 @@ def get_api_key() -> Optional[str]: def save_api_key(api_key: str) -> Path: - """Save API key to ~/.runpod/config.toml via runpod-python. + """Save API key into the [default] section of ~/.runpod/config.toml. + + Updates only flash's `[default].api_key` value, preserving any other + content in the file (notably runpodctl's top-level `apikey`/`apiurl` + keys and other profile sections). Args: api_key: The API key to save. @@ -59,7 +72,18 @@ def save_api_key(api_key: str) -> Path: Path to the credentials file. """ path = get_credentials_path() - set_credentials(api_key, overwrite=True) + path.parent.mkdir(parents=True, exist_ok=True) + + # newline="" disables universal-newline translation so tomlkit sees (and + # preserves) the file's original line endings rather than collapsing CRLF. + existing = "" + if path.exists(): + with path.open("r", encoding="utf-8", newline="") as f: + existing = f.read() + new_content = _upsert_default_api_key(existing, api_key) + with path.open("w", encoding="utf-8", newline="") as f: + f.write(new_content) + try: os.chmod(path, 0o600) except OSError: @@ -67,6 +91,43 @@ def save_api_key(api_key: str) -> Path: return path +def _upsert_default_api_key(content: str, api_key: str) -> str: + """Update `[default].api_key` in TOML text, leaving the rest intact. + + Parses with tomlkit so foreign top-level keys (runpodctl's + `apikey`/`apiurl`), sibling profile sections, comments, and the file's + original line endings are preserved. Only `[default].api_key` is mutated; + a missing `[default]` table is created. + + If the existing file is malformed TOML it is already unloadable (runpodctl + cannot read it either), so we cannot preserve it. Rather than block login, + fall back to a fresh document and warn -- the discarded content held no + recoverable credentials. + """ + if content: + try: + doc = tomlkit.parse(content) + except tomlkit.exceptions.TOMLKitError: + log.warning( + "Existing credentials file is not valid TOML; replacing it " + "with a fresh [default] section. Unrelated content was lost.", + exc_info=True, + ) + doc = tomlkit.document() + else: + doc = tomlkit.document() + + section = doc.get(_DEFAULT_SECTION) + if isinstance(section, (tomlkit.items.Table, tomlkit.items.InlineTable)): + section[_API_KEY_FIELD] = api_key + else: + table = tomlkit.table() + table[_API_KEY_FIELD] = api_key + doc[_DEFAULT_SECTION] = table + + return tomlkit.dumps(doc) + + def check_and_migrate_legacy_credentials() -> None: """Check for credentials at old XDG path and migrate if needed. diff --git a/tests/unit/test_credentials.py b/tests/unit/test_credentials.py index 6ce4d1c0..5ef368a2 100644 --- a/tests/unit/test_credentials.py +++ b/tests/unit/test_credentials.py @@ -1,9 +1,16 @@ """Unit tests for credential storage and retrieval.""" +import logging import os +import sys from pathlib import Path from unittest.mock import patch +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + from runpod_flash.core.credentials import ( get_api_key, get_credentials_path, @@ -69,3 +76,178 @@ def test_sets_restrictive_permissions(self, isolate_credentials_file): save_api_key("secret") mode = oct(isolate_credentials_file.stat().st_mode & 0o777) assert mode == "0o600" + + def test_preserves_runpodctl_top_level_keys(self, isolate_credentials_file): + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + "apikey = 'rpa_runpodctl_key'\n" + "apiurl = 'https://api.runpod.io/graphql'\n" + "\n" + "[default]\n" + 'api_key = "old-flash-key"\n' + ) + save_api_key("new-flash-key") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["apikey"] == "rpa_runpodctl_key" + assert parsed["apiurl"] == "https://api.runpod.io/graphql" + assert parsed["default"]["api_key"] == "new-flash-key" + + def test_adds_default_section_when_missing(self, isolate_credentials_file): + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + "apikey = 'rpa_runpodctl_key'\napiurl = 'https://api.runpod.io/graphql'\n" + ) + save_api_key("flash-key") + text = isolate_credentials_file.read_text() + parsed = tomllib.loads(text) + assert parsed["apikey"] == "rpa_runpodctl_key" + assert parsed["apiurl"] == "https://api.runpod.io/graphql" + assert parsed["default"]["api_key"] == "flash-key" + + def test_preserves_other_profile_sections(self, isolate_credentials_file): + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + "[default]\n" + 'api_key = "old"\n' + "\n" + "[staging]\n" + 'api_key = "staging-key"\n' + 'extra = "preserved"\n' + ) + save_api_key("new-default") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == "new-default" + assert parsed["staging"]["api_key"] == "staging-key" + assert parsed["staging"]["extra"] == "preserved" + + def test_creates_file_with_only_default_when_missing( + self, isolate_credentials_file + ): + save_api_key("first-key") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed == {"default": {"api_key": "first-key"}} + + def test_preserves_inline_comment_on_other_section_header( + self, isolate_credentials_file + ): + """An inline comment on a later section header must not redirect the + update onto that section. Regression for the regex section-boundary bug. + """ + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + "[default]\n" + 'some_other_field = "x"\n' + "\n" + "[staging] # production environment\n" + 'api_key = "staging-key"\n' + ) + save_api_key("new-flash-key") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == "new-flash-key" + assert parsed["staging"]["api_key"] == "staging-key" + + def test_updates_default_with_inline_comment_header(self, isolate_credentials_file): + """`[default] # comment` must update in place, not append a duplicate + `[default]` table that tomllib would refuse to load. + """ + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + '[default] # flash profile\napi_key = "old"\n' + ) + save_api_key("new-key") + text = isolate_credentials_file.read_text() + parsed = tomllib.loads(text) + assert parsed["default"]["api_key"] == "new-key" + assert text.count("[default]") == 1 + + def test_handles_default_header_without_trailing_newline( + self, isolate_credentials_file + ): + """A `[default]` header with no trailing newline must not concatenate + the new key onto the header line. + """ + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text("[default]") + save_api_key("new-key") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == "new-key" + + def test_roundtrips_api_key_with_backslash_and_quote( + self, isolate_credentials_file + ): + """Keys containing backslash and double-quote must survive a write/read + round-trip through a real TOML parser. + """ + weird_key = r'a\b"c' + save_api_key(weird_key) + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == weird_key + + def test_roundtrips_api_key_with_control_char(self, isolate_credentials_file): + """Keys containing control characters (e.g. a tab) must produce valid + TOML rather than a file the next load rejects. + """ + weird_key = "tok\ten\nwith-controls" + save_api_key(weird_key) + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == weird_key + + def test_preserves_crlf_line_endings(self, isolate_credentials_file): + """A CRLF-edited config must not gain a stray LF-only line.""" + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_bytes( + b"apikey = 'rpa_runpodctl_key'\r\n\r\n[default]\r\napi_key = \"old\"\r\n" + ) + save_api_key("new-key") + raw = isolate_credentials_file.read_bytes() + assert b"\r\n" in raw + assert b"\n" not in raw.replace(b"\r\n", b"") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == "new-key" + assert parsed["apikey"] == "rpa_runpodctl_key" + + def test_inserts_api_key_into_default_without_existing_key( + self, isolate_credentials_file + ): + """`[default]` with no `api_key` field gains one without disturbing + sibling fields or other sections. + """ + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + '[default]\nfoo = "bar"\n\n[other]\nx = 1\n' + ) + save_api_key("new-key") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed["default"]["api_key"] == "new-key" + assert parsed["default"]["foo"] == "bar" + assert parsed["other"]["x"] == 1 + + def test_preserves_comments_and_formatting(self, isolate_credentials_file): + """Comments and unrelated content survive the round-trip.""" + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text( + "# managed by runpodctl\n" + "apikey = 'rpa_key'\n" + "\n" + "[default]\n" + "# flash credentials\n" + 'api_key = "old"\n' + ) + save_api_key("new-key") + text = isolate_credentials_file.read_text() + assert "# managed by runpodctl" in text + assert "# flash credentials" in text + + def test_recovers_from_corrupt_existing_file( + self, isolate_credentials_file, caplog + ): + """A malformed config (already unloadable) must not block login: fall + back to a fresh minimal document and warn rather than raise. + """ + isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) + isolate_credentials_file.write_text("not valid toml {{{{\n") + with caplog.at_level(logging.WARNING): + save_api_key("new-key") + parsed = tomllib.loads(isolate_credentials_file.read_text()) + assert parsed == {"default": {"api_key": "new-key"}} + assert any(record.levelno == logging.WARNING for record in caplog.records)