fix(cli-auth): write the session-JWT file privately instead of chmod-ing it after (#3441)

* fix(cli-auth): write the session-JWT file privately instead of chmod-ing it after

_store_entry's docstring already promised the file is written "with user-only
read/write permissions (0o600) - the file may hold session JWTs, which are
sensitive". The implementation did not deliver that:

    path.parent.mkdir(parents=True, exist_ok=True)
    ...
    path.write_text(json.dumps(data, indent=2))
    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)

write_text creates a missing file at the process umask, so on the very first
login - exactly when a session JWT is first persisted - the token sat on disk
readable by every local user until the chmod landed. Measured:

    dir  mode after mkdir     : 0o755
    file mode after write_text: 0o644   <- JWT is on disk at this mode
    file mode after chmod     : 0o600

The parent ~/.omnigent was also left world-traversable, and clear_token
rewrote the same file with no chmod of its own, relying on the mode of a file
it may not have created.

Routes both writers through _write_tokens_file, mirroring the pattern already
used in claude_native_bridge._atomic_write_user_json: a tempfile beside the
target (created owner-only by tempfile before any bytes are written), fsync,
chmod, then os.replace. The directory is created 0o700.

The rename also fixes a robustness bug: write_text truncated in place, so a
write that failed partway left a truncated file, and the JSONDecodeError
handler in _store_entry treats that as {} - silently discarding every stored
token for every server. The temp is discarded on failure and the previous file
is left intact.

Tests: tests/test_cli_auth_token_file_mode.py. Three of the eight fail on the
previous code (the on-disk window, the directory mode, and token loss on a
failed write); the rest pin the final mode, round-tripping, trailing-slash
normalisation and selective clearing.

Signed-off-by: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com>

* style: satisfy ruff format

Pre-commit's ruff-format hook flagged the skipif decorator in the new test
module; it fits on one line.

Signed-off-by: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com>

* refactor(cli-auth): hoist state-dir hardening into _write_tokens_file

Move the 0o700 mkdir + chmod from _store_entry into _write_tokens_file
so every writer routes through it. Previously only _store_entry
hardened the directory, so a clear_token-only interaction left a
pre-existing world-traversable (0o755) ~/.omnigent untightened. Adds a
regression test pinning that clear_token now hardens the dir.

Co-authored-by: Isaac

---------

Signed-off-by: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com>
Co-authored-by: Pat Sukprasert <pat.sukprasert@databricks.com>

omnigent/cli_auth.py

@@ -18,10 +18,12 @@ See ``designs/OIDC_AUTH.md`` §CLI Login Flow.
 
 from __future__ import annotations
 
+import contextlib
 import json
 import logging
 import os
 import stat
+import tempfile
 import time
 from pathlib import Path
 
@@ -53,6 +55,58 @@ def _normalize_server_url(server_url: str) -> str:
     return server_url.rstrip("/")
 
 
+def _write_tokens_file(path: Path, data: dict[str, dict[str, str | float]]) -> None:
+    """Atomically write the auth-tokens file, never exposing it readable.
+
+    The previous sequence was ``path.write_text(...)`` followed by ``os.chmod``.
+    ``write_text`` creates a missing file at the process umask (``0o644`` on a
+    typical box), so on the very first login — exactly when a session JWT is
+    first persisted — the token sat world-readable on disk until the ``chmod``
+    landed. ``clear_token`` did not ``chmod`` at all, relying on the mode of a
+    file that may not have been created here.
+
+    Mirrors ``claude_native_bridge._atomic_write_user_json``: write a
+    ``0o600`` temp beside the target, ``fsync``, then ``os.replace``. The temp
+    is created by :mod:`tempfile` with owner-only permissions before any bytes
+    are written, and the rename means a crash mid-write can no longer truncate
+    the file and lose every stored token.
+
+    Hardens the parent to ``0o700`` first (``mkdir`` alone won't
+    re-permission an existing ``0o755`` dir), so every writer routes
+    through here — including ``clear_token``.
+
+    :param path: Destination file, i.e. ``~/.omnigent/auth_tokens.json``.
+    :param data: The full token map to serialise.
+    :returns: None.
+    :raises OSError: If the temp cannot be written or replaced into place.
+    """
+    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    with contextlib.suppress(OSError):
+        os.chmod(path.parent, stat.S_IRWXU)
+
+    tmp_path: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            "w",
+            encoding="utf-8",
+            dir=path.parent,
+            prefix=f".{path.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as handle:
+            tmp_path = Path(handle.name)
+            json.dump(data, handle, indent=2)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.chmod(tmp_path, stat.S_IRUSR | stat.S_IWUSR)
+        os.replace(tmp_path, path)
+        tmp_path = None
+    finally:
+        if tmp_path is not None:
+            with contextlib.suppress(FileNotFoundError):
+                tmp_path.unlink()
+
+
 def _store_entry(server_url: str, entry: dict[str, str | float]) -> None:
     """Create or update a server's record in the auth-tokens file.
 
@@ -66,7 +120,6 @@ def _store_entry(server_url: str, entry: dict[str, str | float]) -> None:
         ``{"token": "...", "user_id": "...", "expires_at": 1750000000.0}``.
     """
     path = _token_file_path()
-    path.parent.mkdir(parents=True, exist_ok=True)
 
     data: dict[str, dict[str, str | float]] = {}
     if path.exists():
@@ -77,8 +130,7 @@ def _store_entry(server_url: str, entry: dict[str, str | float]) -> None:
 
     data[_normalize_server_url(server_url)] = entry
 
-    path.write_text(json.dumps(data, indent=2))
-    os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
+    _write_tokens_file(path, data)
 
 
 def store_token(
@@ -318,4 +370,4 @@ def clear_token(server_url: str) -> None:
     key = _normalize_server_url(server_url)
     if key in data:
         del data[key]
-        path.write_text(json.dumps(data, indent=2))
+        _write_tokens_file(path, data)

tests/test_cli_auth_token_file_mode.py

@@ -0,0 +1,129 @@
+"""``~/.omnigent/auth_tokens.json`` holds session JWTs and must never be readable.
+
+``_store_entry``'s docstring already promised ``0o600``, but the implementation was
+``path.write_text(...)`` followed by ``os.chmod``. ``write_text`` creates a missing file at
+the process umask, so on first login — exactly when a JWT is first persisted — the token was
+world-readable until the chmod landed. ``clear_token`` rewrote the file with no chmod at all.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import stat
+import sys
+
+import pytest
+
+from omnigent import cli_auth
+
+posix_only = pytest.mark.skipif(sys.platform == "win32", reason="POSIX mode bits")
+
+SERVER = "http://localhost:6767"
+
+
+@pytest.fixture
+def state(tmp_path, monkeypatch):
+    """Point the module's token path at a scratch dir with nothing pre-created."""
+    target = tmp_path / ".omnigent" / "auth_tokens.json"
+    monkeypatch.setattr(cli_auth, "_token_file_path", lambda: target)
+    return target
+
+
+@posix_only
+def test_first_login_never_leaves_the_jwt_readable(state):
+    """The regression: the file must be 0o600 from creation, not after a chmod."""
+    cli_auth.store_token(SERVER, "eyJ-SESSION-JWT", "alice@example.com", 1.0)
+    assert stat.S_IMODE(state.stat().st_mode) == 0o600
+
+
+@posix_only
+def test_state_directory_is_not_world_traversable(state):
+    cli_auth.store_token(SERVER, "eyJ-JWT", "alice@example.com", 1.0)
+    assert stat.S_IMODE(state.parent.stat().st_mode) == 0o700
+
+
+@posix_only
+def test_the_jwt_is_never_on_disk_group_or_world_readable(state, monkeypatch):
+    """The regression itself.
+
+    Hooks ``os.chmod``, which both the old and new writer call, and samples the mode of the
+    path being restricted at that moment. Old code chmod-ed the *final* file after
+    ``write_text`` had already created it at the umask default, so this observes 0o644 with
+    the JWT in it. New code chmods a temp that ``tempfile`` already created owner-only, so
+    the same sample sees 0o600.
+    """
+    observed: dict[str, int] = {}
+    real_chmod = cli_auth.os.chmod
+
+    def spy(path, mode, *a, **kw):
+        try:
+            if os.path.isfile(path):
+                observed[str(path)] = stat.S_IMODE(os.stat(path).st_mode)
+        except OSError:
+            pass
+        return real_chmod(path, mode, *a, **kw)
+
+    monkeypatch.setattr(cli_auth.os, "chmod", spy)
+    cli_auth.store_token(SERVER, "eyJ-SESSION-JWT", "alice@example.com", 1.0)
+
+    assert observed, "expected the writer to restrict a file holding the token"
+    for name, mode in observed.items():
+        assert mode & (stat.S_IRGRP | stat.S_IROTH) == 0, (
+            f"{name} held the JWT at {oct(mode)} before being restricted"
+        )
+
+
+@posix_only
+def test_clear_token_keeps_the_file_private(state):
+    """clear_token rewrote the file without any chmod of its own."""
+    cli_auth.store_token(SERVER, "eyJ-A", "alice@example.com", 1.0)
+    cli_auth.store_token("http://other:1", "eyJ-B", "bob@example.com", 1.0)
+    cli_auth.clear_token(SERVER)
+    assert stat.S_IMODE(state.stat().st_mode) == 0o600
+
+
+@posix_only
+def test_clear_token_hardens_a_preexisting_world_traversable_dir(state):
+    """A dir that predates this fix at 0o755 is tightened even on a clear-only path."""
+    cli_auth.store_token(SERVER, "eyJ-A", "alice@example.com", 1.0)
+    cli_auth.store_token("http://other:1", "eyJ-B", "bob@example.com", 1.0)
+    os.chmod(state.parent, 0o755)
+    cli_auth.clear_token(SERVER)
+    assert stat.S_IMODE(state.parent.stat().st_mode) == 0o700
+
+
+def test_tokens_round_trip(state):
+    cli_auth.store_token(SERVER, "eyJ-A", "alice@example.com", 9e9)
+    cli_auth.store_token("http://other:1", "eyJ-B", "bob@example.com", 9e9)
+    assert cli_auth.load_token(SERVER) == "eyJ-A"
+    assert cli_auth.load_token("http://other:1") == "eyJ-B"
+
+
+def test_trailing_slash_is_the_same_entry(state):
+    cli_auth.store_token(SERVER + "/", "eyJ-A", "alice@example.com", 9e9)
+    assert cli_auth.load_token(SERVER) == "eyJ-A"
+
+
+def test_clear_token_removes_only_the_named_server(state):
+    cli_auth.store_token(SERVER, "eyJ-A", "alice@example.com", 9e9)
+    cli_auth.store_token("http://other:1", "eyJ-B", "bob@example.com", 9e9)
+    cli_auth.clear_token(SERVER)
+    assert cli_auth.load_token(SERVER) is None
+    assert cli_auth.load_token("http://other:1") == "eyJ-B"
+
+
+def test_a_failed_write_does_not_destroy_existing_tokens(state, monkeypatch):
+    """write_text truncated in place, so a crash mid-write lost every token."""
+    cli_auth.store_token(SERVER, "eyJ-ORIGINAL", "alice@example.com", 9e9)
+
+    def boom(*a, **kw):
+        raise OSError("disk full")
+
+    monkeypatch.setattr(cli_auth.os, "replace", boom)
+    with pytest.raises(OSError):
+        cli_auth.store_token(SERVER, "eyJ-REPLACEMENT", "alice@example.com", 9e9)
+
+    assert json.loads(state.read_text())[SERVER]["token"] == "eyJ-ORIGINAL"
+    leftovers = [p.name for p in state.parent.iterdir() if p.name != state.name]
+    assert leftovers == [], f"temp not cleaned up: {leftovers}"