convert GetNodeNetworks to GetInterfaces and nodenetwork_ids to interface_ids
[monitor.git] / getsshkeys.py
index 540ba31..7b74759 100755 (executable)
@@ -4,15 +4,33 @@ import os
 import sys
 import string
 import time
-import plc
+import xml, xmlrpclib
+try:
+       from monitor import config
+       auth = {'Username'   : config.API_AUTH_USER,
+               'AuthMethod' : "password",
+                       'AuthString' : config.API_AUTH_PASSWORD}
+except:
+       import traceback
+       print traceback.print_exc()
+       auth = {'AuthMethod' : "anonymous"}
 
 args = {}
-args['known_hosts'] =  os.environ['HOME'] + "/.ssh/known_hosts"
+args['known_hosts'] =  os.environ['HOME'] + os.sep + ".ssh" + os.sep + "known_hosts"
+try:
+       from monitor import config
+       args['XMLRPC_SERVER'] = config.API_SERVER
+except:
+       args['XMLRPC_SERVER'] = 'https://boot.planet-lab.org/PLCAPI/'
+       print "Using default API server %s" %  args['XMLRPC_SERVER']
 
 class SSHKnownHosts:
        def __init__(self, args = args):
                self.args = args
                self.read_knownhosts()
+               self.auth = auth
+               self.api = xmlrpclib.Server(args['XMLRPC_SERVER'], verbose=False, allow_none=True)
+               self.nodenetworks = {}
 
        def _split_kh_entry(self, line):
                s = line.split(' ')
@@ -64,7 +82,7 @@ class SSHKnownHosts:
                f.close()
 
        def updateAll(self):
-               l_nodes = plc.getNodes() 
+               l_nodes = self.getNodes() 
                d_nodes = {}
                nokey_list = []
                for node in l_nodes:
@@ -80,19 +98,21 @@ class SSHKnownHosts:
                return nokey_list
 
        def delete(self, host):
-               node = plc.getNodes(host) 
-               (host, ip, _, _) = self._record_from_node(node[0])
-               index = "%s,%s" % (host,ip)
-               if index in self.pl_keys:
-                       del self.pl_keys[index]
-               if index in self.other_keys:
-                       del self.other_keys[index]
+               node = self.getNodes(host) 
+               if len(node) > 0:
+                       (host, ip, _, _) = self._record_from_node(node[0])
+                       index = "%s,%s" % (host,ip)
+                       if index in self.pl_keys:
+                               del self.pl_keys[index]
+                       if index in self.other_keys:
+                               del self.other_keys[index]
+               return node
 
        def updateDirect(self, host):
                cmd = os.popen("/usr/bin/ssh-keyscan -t rsa %s 2>/dev/null" % host)
                line = cmd.read()
                (h,  ip,  key,  comment) = self._split_kh_entry(line[:-1])
-               node = plc.getNodes(host)
+               node = self.getNodes(host)
                (host2, ip2, x, x) = self._record_from_node(node[0])
                rec = { self._get_index(host2, ip2) : "%s %s" % (key, "DIRECT") }
 
@@ -100,25 +120,46 @@ class SSHKnownHosts:
                self.other_keys.update(rec)
 
        def update(self, host):
-               node = plc.getNodes(host) 
-               ret = self._record_from_node(node[0])
-               (host, ip, key, comment)  = ret
-               if ip == None:
-                       self.updateDirect(host)
-               else:
-                       rec = { "%s,%s" % (host,ip) : "%s %s" % (key, comment) }
-                       self.pl_keys.update(rec)
+               node = self.delete(host)
+               #node = self.getNodes(host) 
+               if node is not []:
+                       ret = self._record_from_node(node[0])
+                       (host, ip, key, comment)  = ret
+                       if ip == None:
+                               self.updateDirect(host)
+                       else:
+                               rec = { "%s,%s" % (host,ip) : "%s %s" % (key, comment) }
+                               self.pl_keys.update(rec)
+
+       def getNodes(self, host=None):
+               if type(host) == type(""): host = [host]
+
+               # get the node(s) info
+               nodes = self.api.GetNodes(self.auth,host,["hostname","ssh_rsa_key","interface_ids"])
+
+               # for each node's node network, update the self.nodenetworks cache
+               nodenetworks = []
+               for node in nodes:
+                       for net in node["interface_ids"]:
+                               nodenetworks.append(net)
+
+               plcnodenetworks = self.api.GetInterfaces(self.auth,nodenetworks,["nodenetwork_id","ip"])
+               for n in plcnodenetworks:
+                       self.nodenetworks[n["nodenetwork_id"]]=n
+               return nodes
 
        def _record_from_node(self, node, nokey_list=None):
                host = node['hostname']
                key = node['ssh_rsa_key']
 
-               l_nw = plc.getNodeNetworks({'nodenetwork_id':node['nodenetwork_ids']})
-               if len(l_nw) == 0:
-                       # No network for this node. So, skip it.
-                       return (host, None, None, None)
+               nodenetworks = node['interface_ids']
+               if len(nodenetworks)==0: return (host, None, None, None)
 
-               ip = l_nw[0]['ip']
+               # the [0] subscript to node['interface_ids'] means
+               # that this function wont work with multihomed nodes
+               l_nw = self.nodenetworks.get(nodenetworks[0],None)
+               if l_nw is None: return (host, None, None, None)
+               ip = l_nw['ip']
 
                if key == None:
                        if nokey_list is not None: nokey_list += [node]
@@ -126,7 +167,7 @@ class SSHKnownHosts:
 
                key = key.strip()
                # TODO: check for '==' at end of key.
-               if key[-1] != '=':
+               if len(key) > 0 and key[-1] != '=':
                        print "Host with corrupt key! for %s %s" % (node['boot_state'], node['hostname'])
 
                s_date = time.strftime("%Y/%m/%d_%H:%M:%S",time.gmtime(time.time()))
@@ -135,19 +176,14 @@ class SSHKnownHosts:
                return (host, ip, key, "PlanetLab_%s" % s_date) 
 
 
-def main():
+def main(hosts):
        k = SSHKnownHosts()
-       nokey_list = k.updateAll()
+       if len (hosts) > 0:
+               for host in hosts:
+                       k.updateDirect(host)
+       else:
+               k.updateAll()
+       k.write()
 
-       for node in nokey_list:
-               print "%5s %s" % (node['boot_state'], node['hostname'])
-       
 if __name__ == '__main__':
-       #main()
-       k = SSHKnownHosts()
-       #print "update"
-       #k.update('planetlab-4.cs.princeton.edu')
-       #print "updateDirect"
-       k.update(sys.argv[1])
-       #k.updateDirect(sys.argv[1])
-       k.write()
+       main(sys.argv[1:])