From b20853312d53c68fa8292fe0d04624149fcc6d6c Mon Sep 17 00:00:00 2001
From: idaapayo <idaapayo@gmail.com>
Date: Mon, 24 Jan 2022 13:23:12 +0300
Subject: [PATCH] review changes with tests

---
 funga/eth/keystore/keyfile.py | 24 ++++++++++--------------
 funga/eth/runnable/keyfile.py |  9 +++++++--
 tests/test_cli.py             | 10 +++++++++-
 3 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/funga/eth/keystore/keyfile.py b/funga/eth/keystore/keyfile.py
index 0d1664a..a4b0f58 100644
--- a/funga/eth/keystore/keyfile.py
+++ b/funga/eth/keystore/keyfile.py
@@ -28,7 +28,7 @@ hash_keywords = [
     'pbkdf2'
 ]
 
-default_kdfparams = {
+default_scrypt_kdfparams = {
     'dklen': 32,
     'n': 1 << 18,
     'p': 1,
@@ -53,7 +53,7 @@ def to_mac(mac_key, ciphertext_bytes):
 class Hashes:
 
     @staticmethod
-    def from_scrypt(kdfparams=default_kdfparams, passphrase=''):
+    def from_scrypt(kdfparams=default_scrypt_kdfparams, passphrase=''):
         dklen = int(kdfparams['dklen'])
         n = int(kdfparams['n'])
         p = int(kdfparams['p'])
@@ -65,20 +65,16 @@ class Hashes:
 
     @staticmethod
     def from_pbkdf2(kdfparams=pbkdf2_kdfparams, passphrase=''):
-        hashname = kdfparams['prf']
-        pwd = passphrase.encode('utf-8')
-        salt = bytes.fromhex(kdfparams['salt'])
-        itr = int(kdfparams['c'])
-        dklen = int(kdfparams['dklen'])
+        if kdfparams['prf'] == 'hmac-sha256':
+            kdfparams['prf'].replace('sha256')
 
         derived_key = hashlib.pbkdf2_hmac(
-            hash_name=hashname,
-            password=pwd,
-            salt=salt,
-            iterations=itr,
-            dklen=dklen
+            hash_name=kdfparams['prf'],
+            password=passphrase.encode('utf-8'),
+            salt=bytes.fromhex(kdfparams['salt']),
+            iterations=int(kdfparams['c']),
+            dklen=int(kdfparams['dklen'])
         )
-
         return derived_key
 
 
@@ -106,7 +102,7 @@ def to_dict(private_key_bytes, kdf :str, passphrase=''):
 
     if kdf == 'scrypt':
         encryption_key = Hashes.from_scrypt(passphrase=passphrase)
-        kdfparams = default_kdfparams
+        kdfparams = default_scrypt_kdfparams
 
     elif kdf == 'pbkdf2':
         encryption_key = Hashes.from_pbkdf2(passphrase=passphrase)
diff --git a/funga/eth/runnable/keyfile.py b/funga/eth/runnable/keyfile.py
index c708123..b1f2a39 100644
--- a/funga/eth/runnable/keyfile.py
+++ b/funga/eth/runnable/keyfile.py
@@ -16,10 +16,15 @@ from funga.eth.keystore.keyfile import (
         from_file,
         to_dict,
         )
-# from testkeyfile import (
+# from pathlib import Path
+# import sys
+# path_root = Path(__file__).parents[1]
+# sys.path.append(str(path_root))
+# from keystore.testkeyfile import (
 #     from_file,
 #     to_dict
 # )
+
 from funga.eth.encoding import (
         private_key_to_address,
         private_key_from_bytes,
@@ -81,7 +86,7 @@ def main():
         else:
             pk_bytes = os.urandom(32)
         pk = coincurve.PrivateKey(secret=pk_bytes)
-        o = to_dict(pk_bytes, 'pbkdf2', passphrase)
+        o = to_dict(pk_bytes, 'scrypt', passphrase)
         r = json.dumps(o)
 
     print(r) 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 27a8aa6..978835b 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -5,10 +5,18 @@ import os
 
 # external imports
 from hexathon import strip_0x, add_0x
+from pathlib import Path
+
+import sys
+path_root = Path('/home/vincent/ida/grassroots/funga-eth/funga/eth/keystore')
+sys.path.append(str(path_root))
+print(sys.path)
 
 # local imports
 from funga.eth.signer import EIP155Signer
 from funga.eth.keystore.dict import DictKeystore
+
+# import dicttest as d
 from funga.eth.cli.handle import SignRequestHandler
 from funga.eth.transaction import EIP155Transaction
 
@@ -73,7 +81,7 @@ class TestCli(unittest.TestCase):
                          'f86c2a8504a817c8008252089435353535353535353535353535353535353535358203e884deadbeef82466aa0b7c1bbf52f736ada30fe253c7484176f44d6fd097a9720dc85ae5bbc7f060e54a07afee2563b0cf6d00333df51cc62b0d13c63108b2bce54ce2ad24e26ce7b4f25')
 
 
-    def pbkdf2_test(self):
+    def test_pbkdf2(self):
         keystore_filepath = os.path.join(data_dir, 'foo2.json')
         address_hex = self.keystore.import_keystore_file(keystore_filepath)
         logg.debug('getting {}'.format(address_hex))