wip(transaction signature & verification)

mx · mx · commit 55daf2d94044 · 2019-10-30T19:35:02.000+01:00
diff --git a/test_wallet.py b/test_wallet.py
@@ -0,0 +1,27 @@
+from wallet import Wallet
+from transaction import Transaction
+from mx_crypto import MxCrypto
+
+wallet = Wallet(1024, True)
+guy_wallet = Wallet(1024, True)
+
+my_address = wallet.address
+guy_address = guy_wallet.address
+
+value = 10
+
+
+def test_signature():
+    # creation of a transaction with my wallet.
+    transaction = Transaction(
+        sender_address=my_address, recipient_address=guy_address, value=value)
+
+    signature = wallet.sign_transaction(transaction)
+
+    assert transaction.verify_signature(signature) == True
+
+    # guy sign my transaction...
+    fake_signature = guy_wallet.sign_transaction(transaction)
+    is_valid = transaction.verify_signature(fake_signature)
+
+    assert is_valid == False
diff --git a/transaction.py b/transaction.py
@@ -0,0 +1,45 @@
+from collections import OrderedDict
+import hashlib
+
+from Crypto.PublicKey import RSA
+from Crypto.Cipher import PKCS1_OAEP
+from Crypto.Signature import PKCS1_v1_5
+from Crypto import Random
+from base64 import b64encode, b64decode
+from Crypto.Hash import SHA256
+from mx_crypto import MxCrypto
+import json
+
+
+class Transaction:
+
+    def __init__(self, sender_address, recipient_address, value):
+        # @ = public keys
+        self.sender_address = sender_address
+        self.recipient_address = recipient_address
+
+        self.value = value
+
+    def to_dict(self):
+        return OrderedDict({'sender_address': self.sender_address,
+                            'recipient_address': self.recipient_address,
+                            'value': self.value})
+
+    def toJSON(self):
+        # return json.dumps({
+        #     'sender_address': self.sender_address,
+        #     'recipient_address': self.recipient_address,
+        #     'value': self.value
+        # }, sort_keys=True)
+        return json.dumps(self, default=lambda o: o.__dict__,
+                          sort_keys=True, indent=4)
+
+    def verify_signature(self, signature):
+        """
+        Check that the provided signature corresponds to transaction
+        signed by the public key (sender_address)
+        """
+
+        # sender_address = string. I need RSA Object.
+        pubkey = RSA.import_key(self.sender_address)
+        return MxCrypto.verify(pubkey, self.toJSON(), signature)
diff --git a/wallet.py b/wallet.py
@@ -6,25 +6,31 @@
 import hashlib
 from Crypto.Hash import SHA256
 from mx_crypto import MxCrypto
+from transaction import Transaction
 
 
 class Wallet:
     """
     Manage my priv/pub key.
     """
 
-    def __init__(self):
-        try:
-            public_key, private_key, str_public_key = Wallet._import_keys()
-        except FileNotFoundError:
-            # to handle!
-            public_key, private_key = MxCrypto.generate_keys()
+    def __init__(self, key_length=4096, regenerate=False):
+        if regenerate is True:
+            public_key, private_key = MxCrypto.generate_keys(key_length)
+        else:
+            try:
+                public_key, private_key = Wallet._import_keys()
+            except FileNotFoundError:
+                # to handle!
+                public_key, private_key = MxCrypto.generate_keys(key_length)
 
-            Wallet._save_key(public_key, private_key)
+                Wallet._save_key(public_key, private_key)
 
         self.public_key = public_key
         self.private_key = private_key
 
+        self.address = self.public_key.export_key().decode()
+
         # self.address = hashlib.sha256(
         #     str_public_key.encode('utf-8')).hexdigest()
 
@@ -50,14 +56,14 @@ def _import_keys():
         str_public_key = open('./keys/public.pem', 'r').read()
         public_key = RSA.import_key(str_public_key)
 
-        return public_key, private_key, str_public_key
+        return public_key, private_key
 
     def sign_transaction(self, transaction):
         """
         Sign transaction with private key
         """
 
-        if not isinstance(transaction):
+        if not isinstance(transaction, Transaction):
             raise ValueError('transaction should be Transation instance.')
 
-        return Wallet.sign(str(transaction.to_dict()), self.private_key)
+        return MxCrypto.sign(self.private_key, transaction.toJSON())
