From 25ffb620a90d560d3554f3b849af46f31a83fb2c Mon Sep 17 00:00:00 2001
From: nolash <dev@holbrook.no>
Date: Fri, 9 Apr 2021 16:05:16 +0200
Subject: [PATCH] Improve hex prefix handling for keyfile and dictstore

---
 crypto_dev_signer/keystore/dict.py      | 14 +++++++++-----
 crypto_dev_signer/keystore/keyfile.py   |  5 ++++-
 crypto_dev_signer/keystore/reference.py | 16 ++++++++++------
 crypto_dev_signer/runnable/keyfile.py   |  2 +-
 4 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/crypto_dev_signer/keystore/dict.py b/crypto_dev_signer/keystore/dict.py
index 4f7ae95..e8d2a16 100644
--- a/crypto_dev_signer/keystore/dict.py
+++ b/crypto_dev_signer/keystore/dict.py
@@ -2,7 +2,10 @@
 import logging
 
 # external imports
-from hexathon import strip_0x
+from hexathon import (
+        strip_0x,
+        add_0x,
+        )
 
 # local imports
 #from . import keyapi
@@ -20,12 +23,13 @@ class DictKeystore(Keystore):
 
 
     def get(self, address, password=None):
+        address_key = strip_0x(address).lower()
         if password != None:
             logg.debug('password ignored as dictkeystore doesnt do encryption')
         try:
-            return self.keys[address]
+            return self.keys[address_key]
         except KeyError:
-            raise UnknownAccountError(address)
+            raise UnknownAccountError(address_key)
 
 
     def list(self):
@@ -34,7 +38,7 @@ class DictKeystore(Keystore):
 
     def import_key(self, pk, password=None):
         address_hex = private_key_to_address(pk)
-        address_hex_clean = strip_0x(address_hex)
+        address_hex_clean = strip_0x(address_hex).lower()
         self.keys[address_hex_clean] = pk.secret
         logg.debug('added key {}'.format(address_hex))
-        return address_hex_clean
+        return add_0x(address_hex)
diff --git a/crypto_dev_signer/keystore/keyfile.py b/crypto_dev_signer/keystore/keyfile.py
index 55deb33..2cb94a3 100644
--- a/crypto_dev_signer/keystore/keyfile.py
+++ b/crypto_dev_signer/keystore/keyfile.py
@@ -6,6 +6,7 @@ import json
 import uuid
 
 # external imports
+import coincurve
 from Crypto.Cipher import AES
 from Crypto.Util import Counter
 import sha3
@@ -72,7 +73,9 @@ class Ciphers:
         return ciphertext
 
 
-def to_dict(private_key, passphrase=''):
+def to_dict(private_key_bytes, passphrase=''):
+   
+    private_key = coincurve.PrivateKey(secret=private_key_bytes)
 
     encryption_key = Hashes.from_scrypt(passphrase=passphrase)
 
diff --git a/crypto_dev_signer/keystore/reference.py b/crypto_dev_signer/keystore/reference.py
index e080ef6..d311458 100644
--- a/crypto_dev_signer/keystore/reference.py
+++ b/crypto_dev_signer/keystore/reference.py
@@ -10,7 +10,10 @@ from cryptography.fernet import Fernet
 from sqlalchemy import create_engine, text
 from sqlalchemy.orm import sessionmaker
 import sha3
-from hexathon import strip_0x
+from hexathon import (
+        strip_0x,
+        add_0x,
+        )
 
 # local imports
 from .interface import Keystore
@@ -54,7 +57,7 @@ class ReferenceKeystore(Keystore):
 
 
         def get(self, address, password=None):
-            safe_address = strip_0x(address)
+            safe_address = strip_0x(address).lower()
             s = text('SELECT key_ciphertext FROM ethereum WHERE wallet_address_hex = :a')
             r = self.db_session.execute(s, {
                 'a': safe_address,
@@ -64,14 +67,15 @@ class ReferenceKeystore(Keystore):
                 k = r.first()[0]
             except TypeError:
                 self.db_session.rollback()
-                raise UnknownAccountError(address)
+                raise UnknownAccountError(safe_address)
             self.db_session.commit()
-            return self._decrypt(k, password)
+            a = self._decrypt(k, password)
+            return a
 
 
         def import_key(self, pk, password=None):
             address_hex = private_key_to_address(pk)
-            address_hex_clean = strip_0x(address_hex)
+            address_hex_clean = strip_0x(address_hex).lower()
 
             c = self._encrypt(pk.secret, password)
             s = text('INSERT INTO ethereum (wallet_address_hex, key_ciphertext) VALUES (:a, :c)') #%s, %s)')
@@ -82,7 +86,7 @@ class ReferenceKeystore(Keystore):
                 )
             self.db_session.commit()
             logg.info('added private key for address {}'.format(address_hex_clean))
-            return address_hex
+            return add_0x(address_hex)
 
 
         def _encrypt(self, private_key, password):
diff --git a/crypto_dev_signer/runnable/keyfile.py b/crypto_dev_signer/runnable/keyfile.py
index 106eacb..80303ea 100644
--- a/crypto_dev_signer/runnable/keyfile.py
+++ b/crypto_dev_signer/runnable/keyfile.py
@@ -62,7 +62,7 @@ def main():
         else:
             pk_bytes = os.urandom(32)
         pk = coincurve.PrivateKey(secret=pk_bytes)
-        o = to_dict(pk, passphrase)
+        o = to_dict(pk_bytes, passphrase)
         r = json.dumps(o)
 
     print(r)