Add gen subcommand for BIP39 mnemonic generation with dice entropy, master fingerprint to outputs, and test subcommand with Trezor vector

src/hdwallet_recovery.py

@@ -4,12 +4,14 @@ hdwallet_recovery.py (Python 3.12)
 
 Subcommands:
   - fetchkey (online): download ASCII-armored PGP pubkey, print SHA256 + fingerprint
-  - derive (default): derive ADDRESSES ONLY for ETH/SOL/BTC
+  - gen: generate BIP39 mnemonic (12/15/18/24 words) with optional passphrase and dice entropy
+  - recover (default): derive ADDRESSES ONLY for ETH/SOL/BTC from mnemonic/seed
       - optionally encrypt a secret payload using a PGP public key
 
 Security invariant:
-  - derive never performs network I/O
+  - recover never performs network I/O
   - fetchkey refuses mnemonic/seed/passphrase inputs
+  - gen generates secure mnemonics compliant with BIP39
 """
 
 import argparse
@@ -47,6 +49,14 @@ def require_for_derive(export_private: bool, chains: List[str]):
         _require("base58", "base58")
 
 
+def get_master_fingerprint(seed_bytes: bytes) -> str:
+    from bip_utils import Bip32Slip10Secp256k1, Hash160
+    master = Bip32Slip10Secp256k1.FromSeed(seed_bytes)
+    pubkey_bytes = master.PublicKey().RawCompressed().ToBytes()
+    hash160 = Hash160.QuickDigest(pubkey_bytes)
+    return hash160[:4].hex().upper()
+
+
 # -----------------------------------------------------------------------------
 # PGP helpers (PGPy)
 # -----------------------------------------------------------------------------
@@ -267,7 +277,88 @@ def format_addresses_human(result: Dict[str, Any]) -> str:
     return "\n".join(lines)
 
 
-def cmd_derive(args):
+def cmd_gen(args):
+    require_for_derive(False, args.chains)
+
+    import secrets
+    from bip_utils import Bip39MnemonicGenerator, Bip39Languages, Bip39SeedGenerator
+
+    words_to_entropy = {12: 128, 15: 160, 18: 192, 24: 256}
+    entropy_bits = words_to_entropy[args.words]
+    entropy_bytes_len = entropy_bits // 8
+
+    if args.dice_rolls:
+        rolls = args.dice_rolls.strip().split()
+        if not rolls or not all(r.isdigit() and 1 <= int(r) <= 6 for r in rolls):
+            raise ValueError("--dice-rolls must be space-separated integers 1-6")
+        dice_bytes = ' '.join(rolls).encode('utf-8')
+        crypto_bytes = secrets.token_bytes(entropy_bytes_len)
+        combined = dice_bytes + crypto_bytes
+        entropy = hashlib.sha256(combined).digest()[:entropy_bytes_len]
+    else:
+        entropy = secrets.token_bytes(entropy_bytes_len)
+
+    generator = Bip39MnemonicGenerator(Bip39Languages.ENGLISH)
+    mnemonic = generator.FromEntropy(entropy)
+
+    print(f"📍 Generated {args.words}-word BIP39 mnemonic...")
+    seed_bytes = Bip39SeedGenerator(mnemonic).Generate(args.passphrase or "")
+    fingerprint_with = get_master_fingerprint(seed_bytes)
+    fingerprint_without = None
+    if args.passphrase:
+        seed_without = Bip39SeedGenerator(mnemonic).Generate("")
+        fingerprint_without = get_master_fingerprint(seed_without)
+    result = derive_addresses_and_maybe_secrets(seed_bytes, args.chains, args.addresses, False)
+
+    if args.output == "json":
+        data = {
+            "mnemonic": mnemonic,
+            "passphrase_set": bool(args.passphrase),
+            "master_fingerprint": fingerprint_with,
+            "addresses": result["addresses"]
+        }
+        if fingerprint_without:
+            data["master_fingerprint_no_passphrase"] = fingerprint_without
+        if args.dice_rolls:
+            data["dice_rolls_used"] = True
+        out_text = json.dumps(data, indent=2)
+    else:
+        fp_text = f"Master Fingerprint: {fingerprint_with}"
+        if fingerprint_without:
+            fp_text += f"\nMaster Fingerprint (no passphrase): {fingerprint_without}"
+        dice_note = "\nDice rolls used for extra entropy: Yes" if args.dice_rolls else ""
+        out_text = f"Generated Mnemonic ({args.words} words):\n{mnemonic}\n\nPassphrase set: {bool(args.passphrase)}\n{fp_text}{dice_note}\n\n" + format_addresses_human(result)
+
+    print(out_text)
+
+    if args.file:
+        with open(args.file, "w", encoding="utf-8") as f:
+            f.write(out_text)
+        print(f"✅ Saved to {args.file}")
+
+    if args.pgp_pubkey_file:
+        with open(args.pgp_pubkey_file, "r", encoding="utf-8") as f:
+            pgp_pub = f.read()
+
+        payload = {
+            "version": "v4",
+            "purpose": "hdwallet generated mnemonic",
+            "mnemonic": mnemonic,
+            "passphrase": args.passphrase or "",
+            "master_fingerprint": fingerprint_with,
+            "dice_rolls_used": bool(args.dice_rolls),
+            "addresses": result["addresses"]
+        }
+        if fingerprint_without:
+            payload["master_fingerprint_no_passphrase"] = fingerprint_without
+
+        armored = pgp_encrypt_ascii_armored(pgp_pub, payload)
+        print("\n===== BEGIN PGP ENCRYPTED PAYLOAD =====")
+        print(armored)
+        print("===== END PGP ENCRYPTED PAYLOAD =====\n")
+
+
+def cmd_recover(args):
     require_for_derive(args.export_private, args.chains)
 
     from bip_utils import Bip39MnemonicValidator, Bip39SeedGenerator
@@ -303,14 +394,26 @@ def cmd_derive(args):
             raise ValueError(f"Seed must be 64 bytes (128 hex chars), got {len(b)}")
         seed_bytes = b
 
-    print(f"📍 Deriving {len(args.chains)} chain(s), {args.addresses} address(es) each...")
+    fingerprint_with = get_master_fingerprint(seed_bytes)
+    fingerprint_without = None
+    if args.passphrase and mnemonic:
+        seed_without = Bip39SeedGenerator(mnemonic).Generate("")
+        fingerprint_without = get_master_fingerprint(seed_without)
+
+    print(f"📍 Recovering {len(args.chains)} chain(s), {args.addresses} address(es) each...")
     result = derive_addresses_and_maybe_secrets(seed_bytes, args.chains, args.addresses, args.export_private)
 
     # Always print addresses-only output (safe)
     if args.output == "json":
-        out_text = json.dumps({"addresses": result["addresses"]}, indent=2)
+        data = {"addresses": result["addresses"], "master_fingerprint": fingerprint_with}
+        if fingerprint_without:
+            data["master_fingerprint_no_passphrase"] = fingerprint_without
+        out_text = json.dumps(data, indent=2)
     else:
-        out_text = format_addresses_human({"addresses": result["addresses"]})
+        fp_text = f"Master Fingerprint: {fingerprint_with}"
+        if fingerprint_without:
+            fp_text += f"\nMaster Fingerprint (no passphrase): {fingerprint_without}"
+        out_text = fp_text + "\n\n" + format_addresses_human({"addresses": result["addresses"]})
     print(out_text)
 
     if args.file:
@@ -326,7 +429,10 @@ def cmd_derive(args):
         payload: Dict[str, Any] = {
             "version": "v4",
             "purpose": "hdwallet recovery secret payload",
+            "master_fingerprint": fingerprint_with,
         }
+        if fingerprint_without:
+            payload["master_fingerprint_no_passphrase"] = fingerprint_without
 
         # Always include mnemonic/seed_hex if present (your requirement)
         if mnemonic:
@@ -349,12 +455,67 @@ def cmd_derive(args):
         print("===== END PGP ENCRYPTED PAYLOAD =====\n")
 
 
+def cmd_test(args):
+    require_for_derive(False, ["ethereum"])  # minimal
+
+    from bip_utils import Bip39SeedGenerator, Bip39MnemonicValidator
+
+    print("🧪 Running Trezor BIP39/BIP32 test vectors...")
+
+    # Trezor vector with passphrase
+    mnemonic = "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about"
+    passphrase = "crypto"
+    expected_seed_hex = "92c58d3f4fe52f0111d314f3fa8f10ba498751c37e7c36475c2a5b60145b29708576e11bf6c5c46efac1add0cc26e07c69eb65476742082f9c6460f132181cfe"
+    expected_fp = "AECCC350"
+
+    tests = [
+        {
+            "mnemonic": mnemonic,
+            "passphrase": passphrase,
+            "expected_seed_hex": expected_seed_hex,
+            "expected_fp": expected_fp
+        }
+    ]
+
+    all_passed = True
+
+    for i, test in enumerate(tests, 1):
+        print(f"\nTest {i}: With passphrase '{test['passphrase']}'")
+
+        if not Bip39MnemonicValidator().IsValid(test["mnemonic"]):
+            print(f"❌ Invalid mnemonic")
+            all_passed = False
+            continue
+
+        seed_bytes = Bip39SeedGenerator(test["mnemonic"]).Generate(test["passphrase"])
+        seed_hex = seed_bytes.hex()
+
+        if seed_hex != test["expected_seed_hex"]:
+            print(f"❌ Seed mismatch: got {seed_hex}, expected {test['expected_seed_hex']}")
+            all_passed = False
+        else:
+            print(f"✅ Seed correct")
+
+        fp = get_master_fingerprint(seed_bytes)
+        if fp != test["expected_fp"]:
+            print(f"❌ Fingerprint mismatch: got {fp}, expected {test['expected_fp']}")
+            all_passed = False
+        else:
+            print(f"✅ Fingerprint correct: {fp}")
+
+    if all_passed:
+        print("\n🎉 All tests passed!")
+    else:
+        print("\n❌ Some tests failed!")
+        sys.exit(1)
+
+
 # -----------------------------------------------------------------------------
 # Main
 # -----------------------------------------------------------------------------
 
 def main():
-    parser = argparse.ArgumentParser(description="HD wallet recovery (addresses only) + fetchkey helper")
+    parser = argparse.ArgumentParser(description="HD wallet recovery + fetchkey + mnemonic generation")
     subparsers = parser.add_subparsers(dest="cmd")
 
     p_fetch = subparsers.add_parser("fetchkey", help="Download ASCII-armored PGP pubkey from URL")
@@ -362,18 +523,30 @@ def main():
     p_fetch.add_argument("--out", help="Write key to file (recommended)", default=None)
     p_fetch.add_argument("--timeout", type=int, default=15, help="HTTP timeout seconds")
 
-    # derive mode args (default when no subcommand is given)
+    p_test = subparsers.add_parser("test", help="Run tests against Trezor BIP39/BIP32 vectors")
+
+    p_gen = subparsers.add_parser("gen", help="Generate BIP39 mnemonic with optional dice entropy")
+    p_recover = subparsers.add_parser("recover", help="Derive addresses from mnemonic/seed")
+
+    # Common args for gen and recover
+    for p in [p_gen, p_recover, parser]:  # parser for default recover
+        p.add_argument("--passphrase", default="", help="BIP39 passphrase (optional)")
+        p.add_argument("--chains", nargs="+", choices=["ethereum", "solana", "bitcoin"],
+                        default=["ethereum", "solana", "bitcoin"])
+        p.add_argument("--addresses", type=int, default=5)
+        p.add_argument("--output", choices=["text", "json"], default="text")
+        p.add_argument("--file", help="Save output to file", default=None)
+        p.add_argument("--pgp-pubkey-file", help="ASCII-armored PGP public key file to encrypt payload", default=None)
+
+    # Gen specific
+    p_gen.add_argument("--words", type=int, choices=[12,15,18,24], default=12, help="Number of words for mnemonic")
+    p_gen.add_argument("--dice-rolls", help="Space-separated die rolls (1-6) for extra entropy")
+
+    # Recover specific
     parser.add_argument("--mnemonic", help="BIP39 mnemonic (12/24 words)")
     parser.add_argument("--seed", help="64-byte seed hex (128 hex chars)")
     parser.add_argument("--interactive", action="store_true", help="Prompt for input via stdin")
-    parser.add_argument("--passphrase", default="", help="BIP39 passphrase (optional)")
     parser.add_argument("--passphrase-hint", default="", help="Hint/reminder for passphrase (stored only in encrypted payload when --export-private)")
-    parser.add_argument("--chains", nargs="+", choices=["ethereum", "solana", "bitcoin"],
-                        default=["ethereum", "solana", "bitcoin"])
-    parser.add_argument("--addresses", type=int, default=5)
-    parser.add_argument("--output", choices=["text", "json"], default="text")
-    parser.add_argument("--file", help="Save address output to file", default=None)
-    parser.add_argument("--pgp-pubkey-file", help="ASCII-armored PGP public key file to encrypt secrets", default=None)
     parser.add_argument("--export-private", action="store_true",
                         help="Encrypt derived private keys into the PGP payload (never printed). Requires --pgp-pubkey-file.")
 
@@ -389,16 +562,29 @@ def main():
                 getattr(args, "passphrase_hint", ""),
                 getattr(args, "pgp_pubkey_file", None),
                 getattr(args, "export_private", False),
+                getattr(args, "words", 12) != 12,
+                getattr(args, "dice_rolls", None),
             ])
             if forbidden:
-                raise ValueError("fetchkey mode must not be used with mnemonic/seed/passphrase/export-private options")
+                raise ValueError("fetchkey mode must not be used with mnemonic/seed/passphrase/export-private/words/dice-rolls options")
             cmd_fetchkey(args.url, args.out, args.timeout)
             return
 
-        if not (args.mnemonic or args.seed or args.interactive):
+        elif args.cmd == "gen":
-            raise ValueError("Provide --mnemonic or --seed or --interactive (or use subcommand fetchkey)")
+            if any([getattr(args, "mnemonic", None), getattr(args, "seed", None), getattr(args, "interactive", False), getattr(args, "export_private", False), getattr(args, "passphrase_hint", "")]):
+                raise ValueError("gen mode must not be used with mnemonic/seed/interactive/export-private/passphrase-hint options")
+            cmd_gen(args)
+            return
 
-        cmd_derive(args)
+        elif args.cmd == "test":
+            cmd_test(args)
+            return
+
+        # recover (default or explicit)
+        if not (args.mnemonic or args.seed or args.interactive):
+            raise ValueError("Provide --mnemonic or --seed or --interactive (or use subcommand gen/fetchkey)")
+
+        cmd_recover(args)
 
     except Exception as e:
         print(f"❌ Error: {e}", file=sys.stderr)