funga/src/keystore/postgres.py

98 lines
2.7 KiB
Python
Raw Normal View History

2020-08-05 18:14:25 +02:00
import logging
import base64
import os
2020-08-05 18:14:25 +02:00
from cryptography.fernet import Fernet
import psycopg2
from psycopg2 import sql
from eth_keys import KeyAPI
from eth_keys.backends import NativeECCBackend
import sha3
from common import strip_hex_prefix
2020-08-06 11:07:18 +02:00
from keystore.interface import Keystore
keyapi = KeyAPI(NativeECCBackend)
2020-08-05 18:14:25 +02:00
logging.basicConfig(level=logging.DEBUG)
logg = logging.getLogger(__file__)
2020-08-05 19:51:22 +02:00
def to_bytes(x):
return x.encode('utf-8')
2020-08-05 18:14:25 +02:00
2020-08-06 11:07:18 +02:00
class ReferenceKeystore(Keystore):
schema = [
"""CREATE TABLE ethereum (
id SERIAL NOT NULL PRIMARY KEY,
key_ciphertext VARCHAR(256) NOT NULL,
wallet_address_hex CHAR(40) NOT NULL
);
""",
"""CREATE UNIQUE INDEX ethereum_address_idx ON ethereum ( wallet_address_hex );
""",
]
2020-08-05 18:14:25 +02:00
def __init__(self, dbname, **kwargs):
self.conn = psycopg2.connect('dbname=' + dbname)
2020-08-05 18:14:25 +02:00
self.cur = self.conn.cursor()
self.symmetric_key = kwargs.get('symmetric_key')
2020-08-05 18:14:25 +02:00
def get(self, address, password=None):
2020-08-06 11:07:18 +02:00
safe_address = strip_hex_prefix(address)
2020-08-05 18:14:25 +02:00
s = sql.SQL('SELECT key_ciphertext FROM ethereum WHERE wallet_address_hex = %s')
2020-08-06 11:07:18 +02:00
self.cur.execute(s, [ safe_address ] )
2020-08-05 18:14:25 +02:00
k = self.cur.fetchone()[0]
return self._decrypt(k, password)
def new(self, password=None):
b = os.urandom(32)
pk = keyapi.PrivateKey(b)
pubk = keyapi.private_key_to_public_key(pk)
address_hex = pubk.to_checksum_address()
address_hex_clean = strip_hex_prefix(address_hex)
logg.debug('address {}'.format(address_hex_clean))
c = self._encrypt(pk.to_bytes(), password)
s = sql.SQL('INSERT INTO ethereum (wallet_address_hex, key_ciphertext) VALUES (%s, %s)')
self.cur.execute(s, [ address_hex_clean, c.decode('utf-8') ])
self.conn.commit()
return address_hex
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:
2020-08-05 19:51:22 +02:00
password_bytes = to_bytes(password)
h.update(password_bytes)
g = h.digest()
return Fernet(base64.b64encode(g))
2020-08-05 18:14:25 +02:00
def _decrypt(self, c, password):
f = self._generate_encryption_engine(password)
return f.decrypt(c.encode('utf-8'))
2020-08-05 18:14:25 +02:00
def __del__(self):
logg.debug('closing database')
self.conn.commit()
2020-08-05 18:14:25 +02:00
self.cur.close()
self.conn.close()