From e5185b6fba35db0cf675a37d43c5108d98711fc2 Mon Sep 17 00:00:00 2001
From: Luke Childs <lukechilds123@gmail.com>
Date: Fri, 10 May 2019 15:44:29 +0700
Subject: [PATCH] Abstract key derivation into key format interface

---
 src/index.js       | 18 +++++++++++++-----
 src/p2pkh.js       | 15 +++------------
 src/p2wpkh-p2sh.js | 15 +++------------
 src/p2wpkh.js      | 15 +++------------
 src/wif.js         |  9 +++++++++
 5 files changed, 31 insertions(+), 41 deletions(-)
 create mode 100644 src/wif.js

diff --git a/src/index.js b/src/index.js
index 854d63a..7a7997b 100644
--- a/src/index.js
+++ b/src/index.js
@@ -2,6 +2,10 @@ const Emitter = require('tiny-emitter');
 
 const ONE_SECOND = 1000;
 
+const keyFormats = new Map(Object.entries({
+	wif: require('./wif')
+}));
+
 const addressFormats = new Map(Object.entries({
 	p2pkh: require('./p2pkh'),
 	'p2wpkh-p2sh': require('./p2wpkh-p2sh'),
@@ -9,8 +13,9 @@ const addressFormats = new Map(Object.entries({
 }));
 
 class Vain extends Emitter {
-	constructor({addressFormat = 'p2pkh', prefix}) {
+	constructor({keyFormat = 'wif', addressFormat = 'p2pkh', prefix}) {
 		super();
+		this.keyFormat = keyFormats.get(keyFormat);
 		this.addressFormat = addressFormats.get(addressFormat);
 
 		if (typeof prefix !== 'string' || prefix.length === 0) {
@@ -32,15 +37,17 @@ class Vain extends Emitter {
 
 			let found;
 			let attempts = 0;
-			let data;
+			let keyData;
+			let address;
 			let lastUpdate = Date.now();
 
 			while (!found) {
 				attempts++;
 
-				data = this.addressFormat.derive();
+				keyData = this.keyFormat.generate();
+				address = this.addressFormat.derive(keyData.publicKey);
 
-				if (data.address.startsWith(this.prefix)) {
+				if (address.startsWith(this.prefix)) {
 					found = true;
 				}
 
@@ -64,7 +71,8 @@ class Vain extends Emitter {
 			const result = {
 				duration,
 				addressesPerSecond,
-				...this.addressFormat.format(data)
+				address,
+				...this.keyFormat.format(keyData)
 			};
 			this.emit('found', result);
 			resolve(result);
diff --git a/src/p2pkh.js b/src/p2pkh.js
index a9cf33f..dbf7bf7 100644
--- a/src/p2pkh.js
+++ b/src/p2pkh.js
@@ -6,19 +6,10 @@ const p2pkh = {
 	charset: base58
 };
 
-p2pkh.derive = () => {
-	const keyPair = bitcoin.ECPair.makeRandom();
-	const {address} = bitcoin.payments.p2pkh({pubkey: keyPair.publicKey});
+p2pkh.derive = pubkey => {
+	const {address} = bitcoin.payments.p2pkh({pubkey});
 
-	return {
-		address,
-		keyPair
-	};
+	return address;
 };
 
-p2pkh.format = ({address, keyPair}) => ({
-	address,
-	wif: keyPair.toWIF()
-});
-
 module.exports = p2pkh;
diff --git a/src/p2wpkh-p2sh.js b/src/p2wpkh-p2sh.js
index b531b41..f2a0241 100644
--- a/src/p2wpkh-p2sh.js
+++ b/src/p2wpkh-p2sh.js
@@ -6,21 +6,12 @@ const p2wpkhp2sh = {
 	charset: base58
 };
 
-p2wpkhp2sh.derive = () => {
-	const keyPair = bitcoin.ECPair.makeRandom();
+p2wpkhp2sh.derive = pubkey => {
 	const {address} = bitcoin.payments.p2sh({
-		redeem: bitcoin.payments.p2wpkh({pubkey: keyPair.publicKey})
+		redeem: bitcoin.payments.p2wpkh({pubkey})
 	});
 
-	return {
-		address,
-		keyPair
-	};
+	return address;
 };
 
-p2wpkhp2sh.format = ({address, keyPair}) => ({
-	address,
-	wif: keyPair.toWIF()
-});
-
 module.exports = p2wpkhp2sh;
diff --git a/src/p2wpkh.js b/src/p2wpkh.js
index f844f18..cc14aad 100644
--- a/src/p2wpkh.js
+++ b/src/p2wpkh.js
@@ -6,19 +6,10 @@ const p2wpkh = {
 	charset: bech32
 };
 
-p2wpkh.derive = () => {
-	const keyPair = bitcoin.ECPair.makeRandom();
-	const {address} = bitcoin.payments.p2wpkh({pubkey: keyPair.publicKey});
+p2wpkh.derive = pubkey => {
+	const {address} = bitcoin.payments.p2wpkh({pubkey});
 
-	return {
-		address,
-		keyPair
-	};
+	return address;
 };
 
-p2wpkh.format = ({address, keyPair}) => ({
-	address,
-	wif: keyPair.toWIF()
-});
-
 module.exports = p2wpkh;
diff --git a/src/wif.js b/src/wif.js
new file mode 100644
index 0000000..eb97387
--- /dev/null
+++ b/src/wif.js
@@ -0,0 +1,9 @@
+const bitcoin = require('bitcoinjs-lib');
+
+const wif = {};
+
+wif.generate = bitcoin.ECPair.makeRandom;
+
+wif.format = keyPair => ({wif: keyPair.toWIF()});
+
+module.exports = wif;