Create and test the Wallet class

Author: 15Dkatz

backend/config.py

@@ -4,3 +4,5 @@
 SECONDS = 1000 * MILLISECONDS
 
 MINE_RATE = 4 * SECONDS
+
+STARTING_BALANCE = 1000

backend/tests/wallet/test_wallet.py

@@ -0,0 +1,15 @@
+from backend.wallet.wallet import Wallet
+
+def test_verify_valid_signature():
+    data = { 'foo': 'test_data' }
+    wallet = Wallet()
+    signature = wallet.sign(data)
+
+    assert Wallet.verify(wallet.public_key, data, signature)
+
+def test_verify_invalid_signature():
+    data = { 'foo': 'test_data' }
+    wallet = Wallet()
+    signature = wallet.sign(data)
+
+    assert not Wallet.verify(Wallet().public_key, data, signature)

backend/wallet/__init__.py

@@ -0,0 +1,64 @@
+import json
+import uuid
+
+from backend.config import STARTING_BALANCE
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives import hashes
+from cryptography.exceptions import InvalidSignature
+
+class Wallet:
+    """
+    An individual wallet for a miner.
+    Keeps track of the miner's balance.
+    Allows a miner to authorize transactions.
+    """
+    def __init__(self):
+        self.address = str(uuid.uuid4())[0:8]
+        self.balance = STARTING_BALANCE
+        self.private_key = ec.generate_private_key(
+            ec.SECP256K1(),
+            default_backend()
+        )
+        self.public_key = self.private_key.public_key()
+
+    def sign(self, data):
+        """
+        Generate a signature based on the data using the local private key.
+        """
+        return self.private_key.sign(
+            json.dumps(data).encode('utf-8'),
+            ec.ECDSA(hashes.SHA256())
+        )
+
+    @staticmethod
+    def verify(public_key, data, signature):
+        """
+        Verify a signature based on the original public key and data.
+        """
+        try:
+            public_key.verify(
+                signature,
+                json.dumps(data).encode('utf-8'),
+                ec.ECDSA(hashes.SHA256())    
+            )
+            return True
+        except InvalidSignature:
+            return False
+
+def main():
+    wallet = Wallet()
+    print(f'wallet.__dict__: {wallet.__dict__}')
+
+    data = { 'foo': 'bar' }
+    signature = wallet.sign(data)
+    print(f'signature: {signature}')
+
+    should_be_valid = Wallet.verify(wallet.public_key, data, signature)
+    print(f'should_be_valid: {should_be_valid}')
+
+    should_be_invalid = Wallet.verify(Wallet().public_key, data, signature)
+    print(f'should_be_invalid: {should_be_invalid}')
+
+if __name__ == '__main__':
+    main()

backend/wallet/wallet.py

@@ -1,4 +1,5 @@
 pytest==5.1.2
 Flask==1.1.1
 pubnub==4.1.6
-requests==2.22.0
+requests==2.22.0
+cryptography==2.7