Add encrypt and password to symmetric key hash in postgres package

src/keystore/postgres.py

@@ -1,9 +1,15 @@
 import logging
 import base64
+import os
 
 from cryptography.fernet import Fernet
 import psycopg2
 from psycopg2 import sql
+from eth_keys import KeyAPI
+from eth_keys.backends import NativeECCBackend
+import sha3
+
+keyapi = KeyAPI(NativeECCBackend)
 
 logging.basicConfig(level=logging.DEBUG)
 logg = logging.getLogger(__file__)
@@ -16,25 +22,44 @@ class ReferenceDatabase:
             logg.debug(kwargs)
             self.conn = psycopg2.connect('dbname='+dbname)
             self.cur = self.conn.cursor()
-            self.cryptengine = None
-            if kwargs.get('symmetric_key') != None:
-                be = kwargs.get('symmetric_key')
-                self.cryptengine = Fernet(base64.b64encode(be))
+            self.symmetric_key = kwargs.get('symmetric_key')
 
 
-        def get(self, address):
+        def get(self, address, password=None):
             s = sql.SQL('SELECT key_ciphertext FROM ethereum WHERE wallet_address_hex = %s')
             logg.debug(address)
             self.cur.execute(s, [ address ] )
             k = self.cur.fetchone()[0]
-            return self.decrypt(k)
+            return self._decrypt(k, password)
 
 
-        def decrypt(self, c):
-            if self.cryptengine == None:
-                return c
-            logg.debug('decryption')
-            return self.cryptengine.decrypt(c.encode('utf-8'))
+        def new(self, address, password=None):
+            b = os.urandom(32)
+            pk = keyapi.PrivateKey(b)
+            logg.debug('pk {}'.format(pk.to_hex()))
+            c = self._encrypt(pk.to_bytes(), password)
+            logg.debug('pkc {} {}'.format(c, len(pk.to_bytes())))
+            s = sql.SQL('INSERT INTO ethereum (wallet_address_hex, key_ciphertext) VALUES (%s, %s)')
+            self.cur.execute(s, [ address, c.decode('utf-8') ])
+
+
+        def _encrypt(self, private_key, password):
+            f = self._generate_encryption_engine(password)
+            return f.encrypt(private_key)
+
+
+        def _generate_encryption_engine(self, password):
+            h = sha3.keccak_256()
+            h.update(self.symmetric_key)
+            if password != None:
+                h.update(password)
+            g = h.digest()
+            return Fernet(base64.b64encode(g))
+
+
+        def _decrypt(self, c, password):
+            f = self._generate_encryption_engine(password)
+            return f.decrypt(c.encode('utf-8'))
              
 
         def __exit__(self):