From 9200479159bc43aa88fe24cd65284294c20172b4 Mon Sep 17 00:00:00 2001
From: Wei Lu <luwei.here@gmail.com>
Date: Thu, 13 Mar 2014 19:05:55 +0800
Subject: [PATCH] getPrivateKey derives key on the fly

also added getPrivateKeyFromAddress
---
 src/wallet.js  | 20 +++++++++++++++-----
 test/wallet.js | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/src/wallet.js b/src/wallet.js
index 995a823..a7fccfc 100644
--- a/src/wallet.js
+++ b/src/wallet.js
@@ -182,13 +182,23 @@ var Wallet = function (seed, options) {
     this.getMasterKey = function() { return masterkey }
 
     this.getPrivateKey = function(index) {
-        if (typeof index == "string")
-            return keys.filter(function(i,k){ return addresses[i] == index })[0]
-        else
-            return keys[index]
+        return this.externalAccount.derive(index)
     }
 
-    this.getPrivateKeys = function() { return keys }
+    this.getInternalPrivateKey = function(index) {
+        return this.internalAccount.derive(index)
+    }
+
+    this.getPrivateKeyForAddress = function(address) {
+      var index;
+      if((index = this.addresses.indexOf(address)) > -1) {
+        return this.getPrivateKey(index)
+      } else if((index = this.changeAddresses.indexOf(address)) > -1) {
+        return this.getInternalPrivateKey(index)
+      } else {
+        throw new Error('Unknown address. Make sure the address is from the keychain and has been generated.')
+      }
+    }
 };
 
 module.exports = Wallet;
diff --git a/test/wallet.js b/test/wallet.js
index 39a8500..3b1a2ca 100644
--- a/test/wallet.js
+++ b/test/wallet.js
@@ -74,4 +74,47 @@ describe('Wallet', function() {
       assert.deepEqual(wallet.changeAddresses, expectedAddresses)
     })
   })
+
+  describe('getPrivateKey', function(){
+    it('returns the private key at the given index of external account', function(){
+      var wallet = new Wallet(seed, {network: 'testnet'})
+
+      assertPrivateKeyEqual(wallet.getPrivateKey(0), wallet.externalAccount.derive(0))
+      assertPrivateKeyEqual(wallet.getPrivateKey(1), wallet.externalAccount.derive(1))
+    })
+  })
+
+  describe('getInternalPrivateKey', function(){
+    it('returns the private key at the given index of internal account', function(){
+      var wallet = new Wallet(seed, {network: 'testnet'})
+
+      assertPrivateKeyEqual(wallet.getInternalPrivateKey(0), wallet.internalAccount.derive(0))
+      assertPrivateKeyEqual(wallet.getInternalPrivateKey(1), wallet.internalAccount.derive(1))
+    })
+  })
+
+  describe('getPrivateKeyForAddress', function(){
+    it('returns the private key for the given address', function(){
+      var wallet = new Wallet(seed, {network: 'testnet'})
+      wallet.generateChangeAddress()
+      wallet.generateAddress()
+      wallet.generateAddress()
+
+      assertPrivateKeyEqual(wallet.getPrivateKeyForAddress("n2fiWrHqD6GM5GiEqkbWAc6aaZQp3ba93X"),
+                   wallet.externalAccount.derive(1))
+      assertPrivateKeyEqual(wallet.getPrivateKeyForAddress("mnXiDR4MKsFxcKJEZjx4353oXvo55iuptn"),
+                   wallet.internalAccount.derive(0))
+    })
+
+    it('raises an error when address is not found', function(){
+      var wallet = new Wallet(seed, {network: 'testnet'})
+      assert.throws(function() {
+        wallet.getPrivateKeyForAddress("n2fiWrHqD6GM5GiEqkbWAc6aaZQp3ba93X")
+      }, Error, 'Unknown address. Make sure the address is from the keychain and has been generated.')
+    })
+  })
+
+  function assertPrivateKeyEqual(key1, key2){
+    assert.equal(key1.toString(), key2.toString())
+  }
 })