From 7e4eca8ee0fd858a9f2e89553246faffd1b5f35e Mon Sep 17 00:00:00 2001 From: Vasko Zdravevski Date: Wed, 6 May 2026 05:24:10 -0600 Subject: [PATCH 1/3] fix(login): preserve runpodctl credentials when saving flash api key flash login was clobbering the top-level apikey/apiurl values that runpodctl writes to ~/.runpod/config.toml, because it delegated to runpod-python's set_credentials() which opens the file in "w" mode. save_api_key() now does a line-based merge that updates only the [default].api_key field and preserves all other content. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/runpod_flash/core/credentials.py | 66 ++++++++++++++++++++++++++-- tests/unit/test_credentials.py | 59 +++++++++++++++++++++++++ 2 files changed, 122 insertions(+), 3 deletions(-) diff --git a/src/runpod_flash/core/credentials.py b/src/runpod_flash/core/credentials.py index 47350939..7d10721d 100644 --- a/src/runpod_flash/core/credentials.py +++ b/src/runpod_flash/core/credentials.py @@ -8,6 +8,7 @@ import logging import os +import re from pathlib import Path from typing import Optional @@ -15,11 +16,18 @@ 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. +_DEFAULT_HEADER_RE = re.compile(r"^\s*\[default\]\s*$") +_SECTION_HEADER_RE = re.compile(r"^\s*\[[^\]]+\]\s*$") +_API_KEY_LINE_RE = re.compile(r"^\s*api_key\s*=") + _OLD_XDG_PATH = Path.home() / ".config" / "runpod" / "credentials.toml" @@ -50,7 +58,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 +71,12 @@ 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) + + existing = path.read_text(encoding="utf-8") if path.exists() else "" + new_content = _upsert_default_api_key(existing, api_key) + path.write_text(new_content, encoding="utf-8") + try: os.chmod(path, 0o600) except OSError: @@ -67,6 +84,49 @@ def save_api_key(api_key: str) -> Path: return path +def _toml_quote(value: str) -> str: + escaped = value.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped}"' + + +def _upsert_default_api_key(content: str, api_key: str) -> str: + """Update `[default].api_key` in TOML text, leaving the rest intact.""" + new_line = f"api_key = {_toml_quote(api_key)}" + + if not content: + return f"[default]\n{new_line}\n" + + lines = content.splitlines(keepends=True) + + default_start: Optional[int] = None + default_end = len(lines) + for i, line in enumerate(lines): + if _DEFAULT_HEADER_RE.match(line): + default_start = i + for j in range(i + 1, len(lines)): + if _SECTION_HEADER_RE.match(lines[j]): + default_end = j + break + break + + if default_start is None: + suffix = "" if content.endswith("\n") else "\n" + separator = "\n" if content.strip() else "" + return f"{content}{suffix}{separator}[default]\n{new_line}\n" + + for i in range(default_start + 1, default_end): + if _API_KEY_LINE_RE.match(lines[i]): + ending = "\n" if lines[i].endswith("\n") else "" + lines[i] = new_line + ending + return "".join(lines) + + insert_idx = default_end + while insert_idx > default_start + 1 and lines[insert_idx - 1].strip() == "": + insert_idx -= 1 + lines.insert(insert_idx, new_line + "\n") + return "".join(lines) + + 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..effe8e4b 100644 --- a/tests/unit/test_credentials.py +++ b/tests/unit/test_credentials.py @@ -1,9 +1,15 @@ """Unit tests for credential storage and retrieval.""" 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 +75,56 @@ 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") + text = isolate_credentials_file.read_text() + assert "apikey = 'rpa_runpodctl_key'" in text + assert "apiurl = 'https://api.runpod.io/graphql'" in text + parsed = tomllib.loads(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"}} From db57f6cc44eac17060ae4f84308ba080507043ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Thu, 4 Jun 2026 07:30:30 -0700 Subject: [PATCH 2/3] refactor(login): use tomlkit for config.toml credential merge Replace the hand-rolled regex TOML merge in save_api_key with a tomlkit round-trip. The line-based approach mis-detected section boundaries when a header carried an inline comment, silently overwriting an unrelated profile's api_key and/or appending a duplicate [default] table. tomlkit parses structurally, so foreign top-level keys, sibling profile sections, comments, escaping, and line endings are all preserved by construction. save_api_key now reads/writes with newline="" so CRLF files keep their endings instead of collapsing to LF. - Drop _DEFAULT_HEADER_RE/_SECTION_HEADER_RE/_API_KEY_LINE_RE and _toml_quote; add _DEFAULT_SECTION/_API_KEY_FIELD constants - Declare tomlkit as a direct dependency (was transitive via runpod) - Add regression tests: inline-comment headers, duplicate-table guard, no-trailing-newline header, backslash/quote/control-char round-trip, CRLF preservation, [default] without api_key --- pyproject.toml | 1 + src/runpod_flash/core/credentials.py | 77 +++++++----------- tests/unit/test_credentials.py | 116 ++++++++++++++++++++++++++- 3 files changed, 144 insertions(+), 50 deletions(-) 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 7d10721d..9f43312f 100644 --- a/src/runpod_flash/core/credentials.py +++ b/src/runpod_flash/core/credentials.py @@ -8,11 +8,11 @@ import logging import os -import re from pathlib import Path from typing import Optional import runpod.cli.groups.config.functions as _runpod_config +import tomlkit from runpod.cli.groups.config.functions import ( get_credentials, @@ -23,10 +23,11 @@ # 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. -_DEFAULT_HEADER_RE = re.compile(r"^\s*\[default\]\s*$") -_SECTION_HEADER_RE = re.compile(r"^\s*\[[^\]]+\]\s*$") -_API_KEY_LINE_RE = re.compile(r"^\s*api_key\s*=") +# 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" @@ -73,9 +74,15 @@ def save_api_key(api_key: str) -> Path: path = get_credentials_path() path.parent.mkdir(parents=True, exist_ok=True) - existing = path.read_text(encoding="utf-8") if path.exists() else "" + # 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) - path.write_text(new_content, encoding="utf-8") + with path.open("w", encoding="utf-8", newline="") as f: + f.write(new_content) try: os.chmod(path, 0o600) @@ -84,47 +91,25 @@ def save_api_key(api_key: str) -> Path: return path -def _toml_quote(value: str) -> str: - escaped = value.replace("\\", "\\\\").replace('"', '\\"') - return f'"{escaped}"' +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. + """ + doc = tomlkit.parse(content) if content else 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 -def _upsert_default_api_key(content: str, api_key: str) -> str: - """Update `[default].api_key` in TOML text, leaving the rest intact.""" - new_line = f"api_key = {_toml_quote(api_key)}" - - if not content: - return f"[default]\n{new_line}\n" - - lines = content.splitlines(keepends=True) - - default_start: Optional[int] = None - default_end = len(lines) - for i, line in enumerate(lines): - if _DEFAULT_HEADER_RE.match(line): - default_start = i - for j in range(i + 1, len(lines)): - if _SECTION_HEADER_RE.match(lines[j]): - default_end = j - break - break - - if default_start is None: - suffix = "" if content.endswith("\n") else "\n" - separator = "\n" if content.strip() else "" - return f"{content}{suffix}{separator}[default]\n{new_line}\n" - - for i in range(default_start + 1, default_end): - if _API_KEY_LINE_RE.match(lines[i]): - ending = "\n" if lines[i].endswith("\n") else "" - lines[i] = new_line + ending - return "".join(lines) - - insert_idx = default_end - while insert_idx > default_start + 1 and lines[insert_idx - 1].strip() == "": - insert_idx -= 1 - lines.insert(insert_idx, new_line + "\n") - return "".join(lines) + return tomlkit.dumps(doc) def check_and_migrate_legacy_credentials() -> None: diff --git a/tests/unit/test_credentials.py b/tests/unit/test_credentials.py index effe8e4b..aa0041b9 100644 --- a/tests/unit/test_credentials.py +++ b/tests/unit/test_credentials.py @@ -86,10 +86,7 @@ def test_preserves_runpodctl_top_level_keys(self, isolate_credentials_file): 'api_key = "old-flash-key"\n' ) save_api_key("new-flash-key") - text = isolate_credentials_file.read_text() - assert "apikey = 'rpa_runpodctl_key'" in text - assert "apiurl = 'https://api.runpod.io/graphql'" in text - parsed = tomllib.loads(text) + 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" @@ -128,3 +125,114 @@ def test_creates_file_with_only_default_when_missing( 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 From d4db83c1d6922f58541a00b3947760dadd90eacc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Thu, 4 Jun 2026 07:58:11 -0700 Subject: [PATCH 3/3] fix(login): fall back to fresh config when existing TOML is malformed _upsert_default_api_key parsed the existing config.toml with tomlkit but did not handle parse failures, so a malformed file (hand-edited or partially written) would raise and abort flash login. A malformed file is already unloadable -- runpodctl cannot read it either -- so there is nothing to preserve. Catch TOMLKitError, log a warning, and fall back to a fresh [default] document so login still completes. Addresses Copilot review on PR #333. --- src/runpod_flash/core/credentials.py | 18 +++++++++++++++++- tests/unit/test_credentials.py | 15 +++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/runpod_flash/core/credentials.py b/src/runpod_flash/core/credentials.py index 9f43312f..a984bf32 100644 --- a/src/runpod_flash/core/credentials.py +++ b/src/runpod_flash/core/credentials.py @@ -98,8 +98,24 @@ def _upsert_default_api_key(content: str, api_key: str) -> str: `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. """ - doc = tomlkit.parse(content) if content else tomlkit.document() + 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)): diff --git a/tests/unit/test_credentials.py b/tests/unit/test_credentials.py index aa0041b9..5ef368a2 100644 --- a/tests/unit/test_credentials.py +++ b/tests/unit/test_credentials.py @@ -1,5 +1,6 @@ """Unit tests for credential storage and retrieval.""" +import logging import os import sys from pathlib import Path @@ -236,3 +237,17 @@ def test_preserves_comments_and_formatting(self, isolate_credentials_file): 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)