feat: implement basic keystore persistence and CLI wallet commands#114
feat: implement basic keystore persistence and CLI wallet commands#114SIDDHANTCOOKIE wants to merge 1 commit into
Conversation
WalkthroughWallet initialization in main.py switches from ephemeral key generation to a persisted keystore (keystore.json) via a new load_or_create_wallet function. CLI gains export_key and import_key commands, prompt logic is updated, default datadir changes to ".minichain", and a test is updated to match. ChangesPersisted wallet keystore
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant User
participant CLI as main.py CLI
participant Keystore as keystore.json
User->>CLI: start node with --datadir
CLI->>Keystore: load_or_create_wallet(datadir)
alt keystore exists
Keystore-->>CLI: return private key hex
else no keystore
CLI->>CLI: generate new signing keypair
CLI->>Keystore: write keystore.json
end
CLI-->>User: node running with wallet pk
User->>CLI: export_key
CLI-->>User: print private key hex + warning
User->>CLI: import_key <hex>
CLI->>CLI: derive new public key
CLI->>Keystore: overwrite keystore.json
CLI-->>User: refreshed prompt with new pk
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@main.py`:
- Around line 509-510: The keystore write in the async wallet flow blocks the
event loop and duplicates the existing persistence logic used by
load_or_create_wallet. Refactor the repeated keystore serialization into a
shared save_keystore helper used by both the wallet creation path and the
load_or_create_wallet path, and perform the actual file write asynchronously
using asyncio.to_thread or run_in_executor. Also ensure the shared helper
preserves the keystore permission hardening by creating the file with 0o600
access.
- Around line 54-77: The wallet loader in the key-creation routine should not
silently generate and save a new keypair when an existing keystore cannot be
read or parsed. Update the logic around the keystore load path so that if
os.path.exists(keystore_path) is true and the load in the try block fails or the
private_key field is missing, the function aborts instead of falling through to
SigningKey.generate() and overwriting the file. Keep the existing behavior only
for the true “no keystore exists yet” case in the wallet/keystore loading
function.
- Around line 71-73: Restrict the keystore write path in main.py so the private
key file is created with owner-only permissions instead of relying on the
default umask. Update the logic around the open(..., "w") call that writes the
JSON keystore to set the file mode to 0600, keeping the change localized to the
keystore creation flow and the code that handles keystore_path so other local
users cannot read the saved key.
- Around line 502-516: `cli_loop()` is using `datadir` inside the `import_key`
path without receiving it, so the `NameError` is caught and key import never
persists. Update `cli_loop` to accept and use `datadir` as an argument, and pass
the same value from `run_node()` when calling it so the `import_key` block can
build `keystore_path` and write `keystore.json` successfully.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: af31a870-745b-471c-8a8a-c996407ec4d0
📒 Files selected for processing (2)
main.pytests/test_persistence_runtime.py
| if os.path.exists(keystore_path): | ||
| try: | ||
| with open(keystore_path, "r") as f: | ||
| data = json.load(f) | ||
| sk_hex = data.get("private_key") | ||
| if sk_hex: | ||
| sk = SigningKey(bytes.fromhex(sk_hex)) | ||
| pk = sk.verify_key.encode(encoder=HexEncoder).decode() | ||
| logger.info("Loaded existing wallet from %s", keystore_path) | ||
| return sk, pk | ||
| except Exception as e: | ||
| logger.warning("Failed to load keystore: %s", e) | ||
|
|
||
| sk = SigningKey.generate() | ||
| pk = sk.verify_key.encode(encoder=HexEncoder).decode() | ||
|
|
||
| os.makedirs(path, exist_ok=True) | ||
| try: | ||
| with open(keystore_path, "w") as f: | ||
| json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f) | ||
| logger.info("Created new wallet at %s", keystore_path) | ||
| except Exception as e: | ||
| logger.warning("Failed to save keystore: %s", e) | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Do not overwrite an existing keystore when loading fails — risks permanent key loss.
If keystore.json exists but reading/parsing fails (corruption, transient permission error, or a missing private_key field), control falls through to generating a brand-new keypair and writing over the same path (Line 72). For a wallet, silently clobbering the only copy of the private key means irrecoverable loss of funds. On a load failure of an existing file, abort instead of regenerating.
🛡️ Proposed guard
if os.path.exists(keystore_path):
try:
with open(keystore_path, "r") as f:
data = json.load(f)
sk_hex = data.get("private_key")
if sk_hex:
sk = SigningKey(bytes.fromhex(sk_hex))
pk = sk.verify_key.encode(encoder=HexEncoder).decode()
logger.info("Loaded existing wallet from %s", keystore_path)
return sk, pk
+ raise ValueError("keystore.json exists but has no 'private_key'")
except Exception as e:
- logger.warning("Failed to load keystore: %s", e)
+ logger.error("Failed to load existing keystore %s: %s", keystore_path, e)
+ raise📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if os.path.exists(keystore_path): | |
| try: | |
| with open(keystore_path, "r") as f: | |
| data = json.load(f) | |
| sk_hex = data.get("private_key") | |
| if sk_hex: | |
| sk = SigningKey(bytes.fromhex(sk_hex)) | |
| pk = sk.verify_key.encode(encoder=HexEncoder).decode() | |
| logger.info("Loaded existing wallet from %s", keystore_path) | |
| return sk, pk | |
| except Exception as e: | |
| logger.warning("Failed to load keystore: %s", e) | |
| sk = SigningKey.generate() | |
| pk = sk.verify_key.encode(encoder=HexEncoder).decode() | |
| os.makedirs(path, exist_ok=True) | |
| try: | |
| with open(keystore_path, "w") as f: | |
| json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f) | |
| logger.info("Created new wallet at %s", keystore_path) | |
| except Exception as e: | |
| logger.warning("Failed to save keystore: %s", e) | |
| if os.path.exists(keystore_path): | |
| try: | |
| with open(keystore_path, "r") as f: | |
| data = json.load(f) | |
| sk_hex = data.get("private_key") | |
| if sk_hex: | |
| sk = SigningKey(bytes.fromhex(sk_hex)) | |
| pk = sk.verify_key.encode(encoder=HexEncoder).decode() | |
| logger.info("Loaded existing wallet from %s", keystore_path) | |
| return sk, pk | |
| raise ValueError("keystore.json exists but has no 'private_key'") | |
| except Exception as e: | |
| logger.error("Failed to load existing keystore %s: %s", keystore_path, e) | |
| raise | |
| sk = SigningKey.generate() | |
| pk = sk.verify_key.encode(encoder=HexEncoder).decode() | |
| os.makedirs(path, exist_ok=True) | |
| try: | |
| with open(keystore_path, "w") as f: | |
| json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f) | |
| logger.info("Created new wallet at %s", keystore_path) | |
| except Exception as e: | |
| logger.warning("Failed to save keystore: %s", e) |
🧰 Tools
🪛 ast-grep (0.44.0)
[warning] 55-55: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(keystore_path, "r")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
[warning] 71-71: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(keystore_path, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🪛 Ruff (0.15.20)
[warning] 56-56: Unnecessary mode argument
Remove mode argument
(UP015)
[warning] 64-64: Do not catch blind exception: Exception
(BLE001)
[warning] 75-75: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@main.py` around lines 54 - 77, The wallet loader in the key-creation routine
should not silently generate and save a new keypair when an existing keystore
cannot be read or parsed. Update the logic around the keystore load path so that
if os.path.exists(keystore_path) is true and the load in the try block fails or
the private_key field is missing, the function aborts instead of falling through
to SigningKey.generate() and overwriting the file. Keep the existing behavior
only for the true “no keystore exists yet” case in the wallet/keystore loading
function.
| try: | ||
| with open(keystore_path, "w") as f: | ||
| json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Restrict keystore file permissions.
The private key is persisted in plaintext hex with default umask permissions (typically world/group-readable). Restrict the file to owner-only (0600) so other local users cannot read the key.
🔒 Proposed change
- os.makedirs(path, exist_ok=True)
+ os.makedirs(path, exist_ok=True)
try:
- with open(keystore_path, "w") as f:
+ fd = os.open(keystore_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+ with os.fdopen(fd, "w") as f:
json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f)
logger.info("Created new wallet at %s", keystore_path)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| try: | |
| with open(keystore_path, "w") as f: | |
| json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f) | |
| try: | |
| fd = os.open(keystore_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) | |
| with os.fdopen(fd, "w") as f: | |
| json.dump({"private_key": sk.encode(encoder=HexEncoder).decode()}, f) |
🧰 Tools
🪛 ast-grep (0.44.0)
[warning] 71-71: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(keystore_path, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@main.py` around lines 71 - 73, Restrict the keystore write path in main.py so
the private key file is created with owner-only permissions instead of relying
on the default umask. Update the logic around the open(..., "w") call that
writes the JSON keystore to set the file mode to 0600, keeping the change
localized to the keystore creation flow and the code that handles keystore_path
so other local users cannot read the saved key.
| try: | ||
| new_sk = SigningKey(bytes.fromhex(parts[1])) | ||
| new_pk = new_sk.verify_key.encode(encoder=HexEncoder).decode() | ||
|
|
||
| path = datadir or "." | ||
| keystore_path = os.path.join(path, "keystore.json") | ||
| os.makedirs(path, exist_ok=True) | ||
| with open(keystore_path, "w") as f: | ||
| json.dump({"private_key": new_sk.encode(encoder=HexEncoder).decode()}, f) | ||
|
|
||
| sk, pk = new_sk, new_pk | ||
| print(f" {C_GREEN}✅ Key imported successfully!{C_RESET}") | ||
| print_prompt_info(pk) | ||
| except Exception as e: | ||
| print(f" {C_RED}❌ Failed to import key: {e}{C_RESET}") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm cli_loop signature and all call sites.
rg -nP '\b(async\s+def\s+cli_loop|cli_loop\s*\()' main.pyRepository: StabilityNexus/MiniChain
Length of output: 279
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant regions around cli_loop, import_key, and run_node.
sed -n '320,360p' main.py
printf '\n----\n'
sed -n '494,520p' main.py
printf '\n----\n'
sed -n '620,645p' main.pyRepository: StabilityNexus/MiniChain
Length of output: 4451
Thread datadir into cli_loop so import_key can persist the key main.py:335, 506, 636 — cli_loop() uses datadir, but it isn’t in the function signature or the run_node() call. That raises NameError, gets swallowed by the except, and makes every import_key attempt fail without writing keystore.json.
🧰 Tools
🪛 ast-grep (0.44.0)
[warning] 508-508: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(keystore_path, "w")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').
(open-filename-from-request)
🪛 Ruff (0.15.20)
[error] 506-506: Undefined name datadir
(F821)
[warning] 509-509: Async functions should not open files with blocking methods like open
(ASYNC230)
[warning] 515-515: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@main.py` around lines 502 - 516, `cli_loop()` is using `datadir` inside the
`import_key` path without receiving it, so the `NameError` is caught and key
import never persists. Update `cli_loop` to accept and use `datadir` as an
argument, and pass the same value from `run_node()` when calling it so the
`import_key` block can build `keystore_path` and write `keystore.json`
successfully.
| with open(keystore_path, "w") as f: | ||
| json.dump({"private_key": new_sk.encode(encoder=HexEncoder).decode()}, f) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Blocking file write inside async coroutine (and repeats keystore logic).
open()/json.dump here block the event loop (Ruff ASYNC230), and this duplicates the persistence logic in load_or_create_wallet (Lines 70-73), including the missing 0o600 permission hardening. Consider extracting a small save_keystore(path) helper reused by both paths and offloading the write via run_in_executor (or asyncio.to_thread).
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 509-509: Async functions should not open files with blocking methods like open
(ASYNC230)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@main.py` around lines 509 - 510, The keystore write in the async wallet flow
blocks the event loop and duplicates the existing persistence logic used by
load_or_create_wallet. Refactor the repeated keystore serialization into a
shared save_keystore helper used by both the wallet creation path and the
load_or_create_wallet path, and perform the actual file write asynchronously
using asyncio.to_thread or run_in_executor. Also ensure the shared helper
preserves the keystore permission hardening by creating the file with 0o600
access.
Source: Linters/SAST tools
1 participant
Addressed Issues:
Fixes #(TODO:issue number)
Screenshots/Recordings:
TODO: If applicable, add screenshots or recordings that demonstrate the interface before and after the changes.
Additional Notes:
AI Usage Disclosure:
We encourage contributors to use AI tools responsibly when creating Pull Requests. While AI can be a valuable aid, it is essential to ensure that your contributions meet the task requirements, build successfully, include relevant tests, and pass all linters. Submissions that do not meet these standards may be closed without warning to maintain the quality and integrity of the project. Please take the time to understand the changes you are proposing and their impact. AI slop is strongly discouraged and may lead to banning and blocking. Do not spam our repos with AI slop.
Check one of the checkboxes below:
I have used the following AI models and tools: TODO
Checklist
Summary by CodeRabbit
New Features
Bug Fixes