From 88042886bedab048390cb2ee060b22a6f64ac78d Mon Sep 17 00:00:00 2001
From: David Shrewsbury <shrewsbury.dave@gmail.com>
Date: Mon, 13 Mar 2017 15:02:44 -0400
Subject: [PATCH] Record SSH public keys for new nodes in ZK

Change-Id: I3ad63196d584d8dc93a8bcdd9b211f8f6a65bf2f
Story: 2000897
---
 nodepool/nodepool.py            | 12 ++++++++++++
 nodepool/nodeutils.py           | 32 ++++++++++++++++++++++++++++++++
 nodepool/tests/test_nodepool.py |  1 +
 nodepool/tests/test_zk.py       |  4 ++++
 nodepool/zk.py                  |  6 +++++-
 5 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/nodepool/nodepool.py b/nodepool/nodepool.py
index 6aedcfd81..b518445bf 100644
--- a/nodepool/nodepool.py
+++ b/nodepool/nodepool.py
@@ -69,6 +69,10 @@ class LaunchAuthException(Exception):
     statsd_key = 'error.auth'
 
 
+class LaunchKeyscanException(Exception):
+    statsd_key = 'error.keyscan'
+
+
 class StatsReporter(object):
     '''
     Class adding statsd reporting functionality.
@@ -356,6 +360,14 @@ class NodeLauncher(threading.Thread, StatsReporter):
         if not host:
             raise LaunchAuthException("Unable to connect via ssh")
 
+        # Get the SSH public keys for the new node and record in ZooKeeper
+        self.log.debug("Gathering host keys for node %s", self._node.id)
+        host_keys = utils.keyscan(preferred_ip)
+        if not host_keys:
+            raise LaunchKeyscanException("Unable to gather host keys")
+        self._node.host_keys = host_keys
+        self._zk.storeNode(self._node)
+
         self._writeNodepoolInfo(host, preferred_ip, self._node)
         if self._label.ready_script:
             self._runReadyScript(host, hostname, self._label.ready_script)
diff --git a/nodepool/nodeutils.py b/nodepool/nodeutils.py
index 13a8384b5..ba2c6c348 100644
--- a/nodepool/nodeutils.py
+++ b/nodepool/nodeutils.py
@@ -16,6 +16,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
 import errno
 import time
 import socket
@@ -73,3 +74,34 @@ def ssh_connect(ip, username, connect_kwargs={}, timeout=60):
     if "access okay" in out:
         return client
     return None
+
+
+def keyscan(ip):
+    '''
+    Scan the IP address for public SSH keys.
+
+    Keys are returned formatted as: "<type> <base64_string>"
+    '''
+    if 'fake' in ip:
+        return ['ssh-rsa FAKEKEY']
+
+    keys = []
+
+    key = None
+    try:
+        t = paramiko.transport.Transport('%s:%s' % (ip, "22"))
+        t.start_client()
+        key = t.get_remote_server_key()
+        t.close()
+    except Exception as e:
+        log.exception("ssh-keyscan failure: %s", e)
+
+    # Paramiko, at this time, seems to return only the ssh-rsa key, so
+    # only the single key is placed into the list.
+    if key:
+        keys.append(
+            "%s %s" % (key.get_name(),
+                       base64.encodestring(str(key)).replace('\n', ''))
+        )
+
+    return keys
diff --git a/nodepool/tests/test_nodepool.py b/nodepool/tests/test_nodepool.py
index 16df1d72b..a75e41b98 100644
--- a/nodepool/tests/test_nodepool.py
+++ b/nodepool/tests/test_nodepool.py
@@ -202,6 +202,7 @@ class TestNodepool(tests.DBTestCase):
         self.assertEqual(len(nodes), 1)
         self.assertEqual(nodes[0].provider, 'fake-provider')
         self.assertEqual(nodes[0].type, 'fake-label')
+        self.assertNotEqual(nodes[0].host_keys, [])
 
     def test_disabled_label(self):
         """Test that a node is not created with min-ready=0"""
diff --git a/nodepool/tests/test_zk.py b/nodepool/tests/test_zk.py
index 6168057c1..77fd0d822 100644
--- a/nodepool/tests/test_zk.py
+++ b/nodepool/tests/test_zk.py
@@ -772,6 +772,7 @@ class TestZKModel(tests.BaseTestCase):
         o.external_id = 'ABCD'
         o.hostname = 'xyz'
         o.comment = 'comment'
+        o.host_keys = ['key1', 'key2']
 
         d = o.toDict()
         self.assertNotIn('id', d)
@@ -790,6 +791,7 @@ class TestZKModel(tests.BaseTestCase):
         self.assertEqual(d['external_id'], o.external_id)
         self.assertEqual(d['hostname'], o.hostname)
         self.assertEqual(d['comment'], o.comment)
+        self.assertEqual(d['host_keys'], o.host_keys)
 
     def test_Node_fromDict(self):
         now = int(time.time())
@@ -810,6 +812,7 @@ class TestZKModel(tests.BaseTestCase):
             'external_id': 'ABCD',
             'hostname': 'xyz',
             'comment': 'comment',
+            'host_keys': ['key1', 'key2'],
         }
 
         o = zk.Node.fromDict(d, node_id)
@@ -829,3 +832,4 @@ class TestZKModel(tests.BaseTestCase):
         self.assertEqual(o.external_id, d['external_id'])
         self.assertEqual(o.hostname , d['hostname'])
         self.assertEqual(o.comment , d['comment'])
+        self.assertEqual(o.host_keys , d['host_keys'])
diff --git a/nodepool/zk.py b/nodepool/zk.py
index 98e07f913..acb2326fe 100644
--- a/nodepool/zk.py
+++ b/nodepool/zk.py
@@ -416,6 +416,7 @@ class Node(BaseModel):
         self.external_id = None
         self.hostname = None
         self.comment = None
+        self.host_keys = []
 
     def __repr__(self):
         d = self.toDict()
@@ -440,7 +441,8 @@ class Node(BaseModel):
                     self.created_time == other.created_time and
                     self.external_id == other.external_id and
                     self.hostname == other.hostname and
-                    self.comment == other.comment)
+                    self.comment == other.comment,
+                    self.host_keys == other.host_keys)
         else:
             return False
 
@@ -462,6 +464,7 @@ class Node(BaseModel):
         d['external_id'] = self.external_id
         d['hostname'] = self.hostname
         d['comment'] = self.comment
+        d['host_keys'] = self.host_keys
         return d
 
     @staticmethod
@@ -489,6 +492,7 @@ class Node(BaseModel):
         o.external_id = d.get('external_id')
         o.hostname = d.get('hostname')
         o.comment = d.get('comment')
+        o.host_keys = d.get('host_keys', [])
         return o