wallet: add simple key management

src/crypto.py

@@ -1,6 +1,10 @@
 """ Generic functions for the cryptographic primitives used in this project. """
 
+import os
+import os.path
+import tempfile
 from binascii import hexlify, unhexlify
+from typing import Iterator, Iterable
 
 from Crypto.Signature import PKCS1_PSS
 from Crypto.Hash import SHA512
@@ -80,3 +84,37 @@ class Signing:
         sign things.
         """
         return self.rsa.has_private()
+
+    @classmethod
+    def read_many_private(cls, file_contents: bytes) -> 'Iterator[Signing]':
+        """ Reads many private keys from the (binary) contents of a file written with `write_many_private`. """
+        end = b"-----END RSA PRIVATE KEY-----"
+        for key in file_contents.strip().split(end):
+            if not key:
+                continue
+
+            key = key.lstrip() + end
+            yield cls(key)
+
+    @staticmethod
+    def write_many_private(path: str, keys: 'Iterable[Signing]'):
+        """ Writes the private keys in `keys` to the file at `path`. """
+        dirname = os.path.dirname(path) or "."
+        with tempfile.NamedTemporaryFile("wb", delete=False, dir=dirname) as fp:
+            try:
+                for key in keys:
+                    fp.write(key.as_bytes(include_priv=True) + b"\n")
+
+                fp.flush()
+                os.fsync(fp.fileno())
+
+                os.rename(fp.name, path)
+            except Exception as e:
+                os.unlink(fp.name)
+                raise e
+
+        fd = os.open(dirname, os.O_DIRECTORY)
+        try:
+            os.fsync(fd)
+        finally:
+            os.close(fd)

wallet.py

@@ -16,8 +16,8 @@ from binascii import hexlify
 import logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)-8s %(message)s")
 
-from src.blockchain import Blockchain
 from src.block import Block
+from src.blockchain import Blockchain
 from src.transaction import Transaction, TransactionTarget, TransactionInput
 from src.crypto import Signing
 
@@ -80,26 +80,45 @@ def private_signing(path):
         raise ValueError("The specified key is not a private key.")
     return val
 
+def wallet_file(path):
+    try:
+        with open(path, "rb") as f:
+            contents = f.read()
+    except FileNotFoundError:
+        return [], path
+    return list(Signing.read_many_private(contents)), path
+
 def main():
     parser = argparse.ArgumentParser(description="Wallet.")
     parser.add_argument("--miner-port", default=40203, type=int,
                         help="The RPC port of the miner to connect to.")
+    parser.add_argument("--wallet", type=wallet_file, default=([],None),
+                        help="The wallet file containing the private keys to use.")
     subparsers = parser.add_subparsers(dest="command")
+
+    balance = subparsers.add_parser("create-address",
+                                    help="Creates new addresses and stores their secret keys in the wallet.")
+    balance.add_argument("file", nargs="+", type=argparse.FileType("wb"),
+                         help="Path to a file where the address should be stored.")
+
     balance = subparsers.add_parser("show-balance",
                                     help="Shows the current balance of the public key "
                                          "stored in the specified file.")
-    balance.add_argument("key", nargs="+", type=Signing.from_file)
+    balance.add_argument("key", nargs="*", type=Signing.from_file)
+
     trans = subparsers.add_parser("show-transactions",
                                   help="Shows all transactions involving the public key "
                                        "stored in the specified file.")
-    trans.add_argument("key", nargs="+", type=Signing.from_file)
+    trans.add_argument("key", nargs="*", type=Signing.from_file)
+
     subparsers.add_parser("show-network",
                           help="Prints networking information about the miner.")
+
     transfer = subparsers.add_parser("transfer", help="Transfer money.")
     transfer.add_argument("--private-key", type=private_signing,
-                          default=[], action="append", required=True,
+                          default=[], action="append", required=False,
                           help="The private key(s) whose coins should be used for the transfer.")
-    transfer.add_argument("--change-key", type=Signing.from_file, required=True,
+    transfer.add_argument("--change-key", type=Signing.from_file, required=False,
                           help="The private key where any remaining coins are sent to.")
     transfer.add_argument("--transaction-fee", type=int, default=0,
                           help="The transaction fee you want to pay to the miner.")
@@ -111,13 +130,31 @@ def main():
     url = "http://localhost:{}/".format(args.miner_port)
     s = requests.session()
 
+    def get_keys(keys):
+        all_keys = keys + args.wallet[0]
+        if not all_keys:
+            print("missing key or wallet", file=sys.stderr)
+            parser.parse_args(["--help"])
+        return all_keys
+
     if args.command == 'show-transactions':
-        for key in args.key:
+        for key in get_keys(args.key):
             for trans in get_transactions(s, url, key):
                 print(trans.to_json_compatible())
+            print()
+    elif args.command == "create-address":
+        if not args.wallet[1]:
+            print("no wallet specified", file=sys.stderr)
+            parser.parse_args(["--help"])
+
+        keys = [Signing.generate_private_key() for _ in args.file]
+        Signing.write_many_private(args.wallet[1], args.wallet[0] + keys)
+        for fp, key in zip(args.file, keys):
+            fp.write(key.as_bytes())
+            fp.close()
     elif args.command == 'show-balance':
         total = 0
-        for pubkey, balance in show_balance(s, url, args.key):
+        for pubkey, balance in show_balance(s, url, get_keys(args.key)):
             print("{}: {}".format(hexlify(pubkey.as_bytes()), balance))
             total += balance
         print()
@@ -130,7 +167,13 @@ def main():
             print("Missing amount to transfer for last target key.\n", file=sys.stderr)
             parser.parse_args(["--help"])
         targets = [TransactionTarget(k, a) for k, a in zip(args.target[::2], args.target[1::2])]
-        build_transaction(s, url, args.private_key, targets, args.change_key, args.transaction_fee)
+        change_key = args.change_key
+        if not change_key:
+            get_keys([]) # shows error if no wallet
+            change_key = Signing.generate_private_key()
+            Signing.write_many_private(args.wallet[1], args.wallet[0] + [change_key])
+
+        build_transaction(s, url, get_keys(args.private_key), targets, change_key, args.transaction_fee)
     else:
         print("You need to specify what to do.\n", file=sys.stderr)
         parser.parse_args(["--help"])