From a38c0c2d93a5743e15dc8e8cbb16cc3d733cd2ca Mon Sep 17 00:00:00 2001
From: Brandon Robertz <brandon@bxroberts.org>
Date: Mon, 8 Jun 2015 19:40:58 -0500
Subject: [PATCH] modularize network version check/tests issues/1265

---
 lib/networks.js    | 12 +++++++++++-
 lib/privatekey.js  | 14 +++++++++-----
 test/networks.js   | 26 +++++++++++++++++++++++++-
 test/privatekey.js | 29 +++++++++++++++++++++++++++++
 4 files changed, 74 insertions(+), 7 deletions(-)

diff --git a/lib/networks.js b/lib/networks.js
index 3ea7c01..ba7246d 100644
--- a/lib/networks.js
+++ b/lib/networks.js
@@ -18,6 +18,15 @@ Network.prototype.toString = function toString() {
   return this.name;
 };
 
+/**
+ * @member Networks#all
+ * Retrieves all networks registered.
+ * @return Array
+ */
+function all() {
+  return networks;
+}
+
 /**
  * @function
  * @member Networks#get
@@ -172,5 +181,6 @@ module.exports = {
   livenet: livenet,
   mainnet: livenet,
   testnet: testnet,
-  get: get
+  get: get,
+  all: all
 };
diff --git a/lib/privatekey.js b/lib/privatekey.js
index 8560f68..2f56072 100644
--- a/lib/privatekey.js
+++ b/lib/privatekey.js
@@ -157,14 +157,18 @@ PrivateKey._transformBuffer = function(buf, network) {
   }
 
   info.network = Networks.get(buf[0], 'privatekey');
-  if (buf[0] === Networks.livenet.privatekey) {
-    info.network = Networks.livenet;
-  } else if (buf[0] === Networks.testnet.privatekey) {
-    info.network = Networks.testnet;
-  } else {
+
+  var allNetworks = Networks.all();
+  var matches = _.filter( allNetworks, function( network) {
+    return buf[0] === network.privatekey;
+  });
+
+  if (matches.length !== 1) {
     throw new Error('Invalid network');
   }
 
+  info.network = matches[0];
+
   if (network && info.network !== Networks.get(network)) {
     throw new TypeError('Private key network mismatch');
   }
diff --git a/test/networks.js b/test/networks.js
index 989e064..bd2b7b5 100644
--- a/test/networks.js
+++ b/test/networks.js
@@ -48,7 +48,31 @@ describe('Networks', function() {
     var net = networks.get('customnet');
     should.equal(net, undefined);
   });
-  
+
+  it('can return custom networks', function() {
+    var custom = {
+      name: 'customnet',
+      alias: 'mynet',
+      pubkeyhash: 0x10,
+      privatekey: 0x90,
+      scripthash: 0x08,
+      xpubkey: 0x0278b20e,
+      xprivkey: 0x0278ade4,
+      networkMagic: 0xe7beb4d4,
+      port: 20001,
+      dnsSeeds: [
+        'localhost',
+        'mynet.localhost'
+      ]
+    };
+    networks.add(custom);
+    customnet = networks.get('customnet');
+    var allNetworks = networks.all();
+    var customInOutput = allNetworks.indexOf(customnet) > -1;
+    should.equal(customInOutput, true);
+    networks.remove(customnet);
+  });
+
   it('should not set a network map for an undefined value', function() {
     var custom = {
       name: 'somenet',
diff --git a/test/privatekey.js b/test/privatekey.js
index 03aad24..cf0395f 100644
--- a/test/privatekey.js
+++ b/test/privatekey.js
@@ -44,6 +44,35 @@ describe('PrivateKey', function() {
     should.exist(a.bn);
   });
 
+  it('should create a private key from a custom network WIF string', function() {
+    var wifNamecoin = '74pxNKNpByQ2kMow4d9kF6Z77BYeKztQNLq3dSyU4ES1K5KLNiz';
+    var nmc = {
+      name: 'namecoin',
+      alias: 'namecoin',
+      pubkeyhash: 0x34,
+      privatekey: 0xB4,
+      // these below aren't the real NMC version numbers
+      scripthash: 0x08,
+      xpubkey: 0x0278b20e,
+      xprivkey: 0x0278ade4,
+      networkMagic: 0xf9beb4fe,
+      port: 20001,
+      dnsSeeds: [
+        'localhost',
+        'mynet.localhost'
+      ]
+    };
+    Networks.add(nmc);
+    var nmcNet = Networks.get('namecoin');
+    var a = new PrivateKey(
+      '74pxNKNpByQ2kMow4d9kF6Z77BYeKztQNLq3dSyU4ES1K5KLNiz',
+      nmcNet
+    );
+    should.exist(a);
+    should.exist(a.bn);
+    Networks.remove(nmcNet);
+  });
+
   it('should create a new random testnet private key with empty data', function() {
     var a = new PrivateKey(null, Networks.testnet);
     should.exist(a);