From 9119b03a07143c75bf910edd033aa448227a25d2 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 26 Jan 2024 13:18:13 -0800
Subject: [PATCH] Chmod 600 keys.json on creation, refs #351

---
 llm/cli.py         |  1 +
 tests/test_keys.py | 11 ++++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/llm/cli.py b/llm/cli.py
index 5accb5e..7e72283 100644
--- a/llm/cli.py
+++ b/llm/cli.py
@@ -517,6 +517,7 @@ def keys_set(name, value):
     path.parent.mkdir(parents=True, exist_ok=True)
     if not path.exists():
         path.write_text(json.dumps(default))
+        path.chmod(0o600)
     try:
         current = json.loads(path.read_text())
     except json.decoder.JSONDecodeError:
diff --git a/tests/test_keys.py b/tests/test_keys.py
index e1c503f..b172664 100644
--- a/tests/test_keys.py
+++ b/tests/test_keys.py
@@ -20,12 +20,17 @@ def test_keys_in_user_path(monkeypatch, env, user_path):
 
 
 def test_keys_set(monkeypatch, tmpdir):
-    user_path = str(tmpdir / "user/keys")
-    monkeypatch.setenv("LLM_USER_PATH", user_path)
+    user_path = tmpdir / "user/keys"
+    monkeypatch.setenv("LLM_USER_PATH", str(user_path))
+    keys_path = user_path / "keys.json"
+    assert not keys_path.exists()
     runner = CliRunner()
     result = runner.invoke(cli, ["keys", "set", "openai"], input="foo")
     assert result.exit_code == 0
-    content = open(user_path + "/keys.json").read()
+    assert keys_path.exists()
+    # Should be chmod 600
+    assert oct(keys_path.stat().mode)[-3:] == "600"
+    content = keys_path.read_text("utf-8")
     assert json.loads(content) == {
         "// Note": "This file stores secret API credentials. Do not share!",
         "openai": "foo",