Skip to content

feat: implement basic keystore persistence and CLI wallet commands #114

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SIDDHANTCOOKIE wants to merge 1 commit into
base: main
Choose a base branch
from
+70 −5
Open

feat: implement basic keystore persistence and CLI wallet commands #114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 69 additions & 4 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
import logging
import re
import sys
import os
import json

from nacl.signing import SigningKey
from nacl.encoding import HexEncoder
Expand All @@ -41,9 +43,38 @@
# Wallet helpers
# ──────────────────────────────────────────────

def create_wallet():
def load_or_create_wallet(datadir: str | None):
path = datadir or "."
keystore_path = os.path.join(path, "keystore.json")

# Security Warning:
# Keys are currently stored in unencrypted raw hex format for minimality.
# In a production environment, this file should be encrypted with a "spending password"
# so that the private key only lives in memory when decrypted by the user.
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)
Comment on lines +71 to +73

@coderabbitai coderabbitai Bot Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 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.

Suggested change
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.

logger.info("Created new wallet at %s", keystore_path)
except Exception as e:
logger.warning("Failed to save keystore: %s", e)

Comment on lines +54 to +77

@coderabbitai coderabbitai Bot Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.

Suggested change
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.

return sk, pk


Expand Down Expand Up @@ -292,6 +323,8 @@ def gradient_text(text: str, c1: tuple[int, int, int], c2: tuple[int, int, int])
{C_CYAN}║{C_RESET} {C_GREEN}peers{C_RESET} - show connected peers {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}connect <multiaddr>{C_RESET} - connect to a peer {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}address{C_RESET} - show your public key {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}export_key{C_RESET} - show your private key {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}import_key <hex>{C_RESET} - import a private key {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}chain{C_RESET} - show chain summary {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}help{C_RESET} - show this help {C_CYAN}║{C_RESET}
{C_CYAN}║{C_RESET} {C_GREEN}quit{C_RESET} - shut down {C_CYAN}║{C_RESET}
Expand All @@ -303,7 +336,11 @@ async def cli_loop(sk, pk, chain, mempool, network):
"""Read commands from stdin asynchronously."""
loop = asyncio.get_event_loop()
print(HELP_TEXT)
print(f" {C_YELLOW}Your address:{C_RESET} {C_BOLD}{pk}{C_RESET}\n")

def print_prompt_info(current_pk):
print(f" {C_YELLOW}Your address:{C_RESET} {C_BOLD}{current_pk}{C_RESET}\n")

print_prompt_info(pk)

while True:
try:
Expand Down Expand Up @@ -450,6 +487,34 @@ async def cli_loop(sk, pk, chain, mempool, network):
elif cmd == "address":
print(f" {pk}")

# ── export_key ──
elif cmd == "export_key":
sk_hex = sk.encode(encoder=HexEncoder).decode()
print(f" {C_YELLOW}Private Key:{C_RESET} {sk_hex}")
print(f" {C_RED}Warning: Never share this key!{C_RESET}")

# ── import_key ──
elif cmd == "import_key":
if len(parts) < 2:
print(" Usage: import_key <hex_private_key>")
continue

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)
Comment on lines +509 to +510

@coderabbitai coderabbitai Bot Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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


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}")
Comment on lines +502 to +516

@coderabbitai coderabbitai Bot Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.py

Repository: 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.py

Repository: 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.


# ── chain ──
elif cmd == "chain":
print(f" Chain length: {len(chain.chain)} blocks")
Expand Down Expand Up @@ -506,7 +571,7 @@ async def cli_loop(sk, pk, chain, mempool, network):

async def run_node(port: int, host: str, connect_to: str | None, fund: int, datadir: str | None):
"""Boot the node, optionally connect to a peer, then enter the CLI."""
sk, pk = create_wallet()
sk, pk = load_or_create_wallet(datadir)

# Load existing chain from disk, or start fresh
chain = None
Expand Down Expand Up @@ -589,7 +654,7 @@ def main():
parser.add_argument("--port", type=int, default=9000, help="TCP port to listen on (default: 9000)")
parser.add_argument("--connect", type=str, default=None, help="Peer address to connect to (multiaddr)")
parser.add_argument("--fund", type=int, default=100, help="Initial coins to fund this wallet (default: 100)")
parser.add_argument("--datadir", type=str, default=None, help="Directory to save/load blockchain state (enables persistence)")
parser.add_argument("--datadir", type=str, default=".minichain", help="Directory to save/load blockchain state (enables persistence)")
args = parser.parse_args()

logging.basicConfig(
Expand Down
2 changes: 1 addition & 1 deletion tests/test_persistence_runtime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ async def fake_cli_loop(sk, pk, chain, mempool, network):

with patch.object(main_module, "P2PNetwork", FakeNetwork), patch.object(
main_module, "cli_loop", fake_cli_loop
), patch.object(main_module, "create_wallet", return_value=(fixed_sk, fixed_pk)):
), patch.object(main_module, "load_or_create_wallet", return_value=(fixed_sk, fixed_pk)):
await main_module.run_node(
port=9401,
host="127.0.0.1",
Expand Down