getsshkeys.py

   1 #!/usr/bin/python
   2
   3 import os
   4 import sys
   5 import string
   6 import time
   7 import soltesz
   8 import plc
   9
  10 args = {}
  11 args['known_hosts'] = "/home/soltesz/.ssh/known_hosts"
  12
  13 class SSHKnownHosts:
  14         def __init__(self, args = args):
  15                 self.args = args
  16                 self.read_knownhosts()
  17
  18         def _split_kh_entry(self, line):
  19                 s = line.split(' ')
  20                 try:
  21                         (host,ip) = s[0].split(',')
  22                 except:
  23                         ip = s[0]
  24                         host = ""
  25
  26                 key = ' '.join(s[1:3])
  27                 comment = ' '.join(s[3:])
  28                 return (host, ip, key, comment)
  29
  30         def _get_index(self, host, ip):
  31                 index = ""
  32                 if host is not "":
  33                         index = "%s,%s" % (host,ip)
  34                 else:
  35                         index = ip
  36                 return index
  37
  38         def read_knownhosts(self):
  39                 kh_read = open(self.args["known_hosts"], 'r')
  40                 self.pl_keys = {}
  41                 self.other_keys = {}
  42                 for line in kh_read:
  43                         (host, ip, key, comment) = self._split_kh_entry(line[:-1])
  44                         rec = { self._get_index(host, ip) : "%s %s" % (key, comment) }
  45                         if 'PlanetLab' in comment:
  46                                 self.pl_keys.update(rec)
  47                         else:
  48                                 self.other_keys.update(rec)
  49
  50                 #for i in self.pl_keys:
  51                 #       print i
  52                 #       print self.pl_keys[i]
  53
  54                 return
  55
  56         def write(self):
  57                 self.write_knownhosts()
  58
  59         def write_knownhosts(self):
  60                 f = open(self.args['known_hosts'], 'w')
  61                 for index in self.pl_keys:
  62                         print >>f, "%s %s" % (index, self.pl_keys[index])
  63                 for index in self.other_keys:
  64                         print >>f, "%s %s" % (index, self.other_keys[index])
  65                 f.close()
  66
  67         def updateAll(self):
  68                 l_nodes = plc.getNodes()
  69                 d_nodes = {}
  70                 nokey_list = []
  71                 for node in l_nodes:
  72                         name = node['hostname']
  73                         d_nodes[name] = node
  74
  75                 for host in d_nodes:
  76                         node = d_nodes[host]
  77                         (host, ip, key, comment) = self._record_from_node(node, nokey_list)
  78                         rec = { "%s,%s" % (host,ip) : "%s %s" % (key, comment) }
  79                         self.pl_keys.update(rec)
  80
  81                 return nokey_list
  82
  83         def delete(self, host):
  84                 node = plc.getNodes(host)
  85                 (host, ip, _, _) = self._record_from_node(node[0])
  86                 index = "%s,%s" % (host,ip)
  87                 if index in self.pl_keys:
  88                         del self.pl_keys[index]
  89                 if index in self.other_keys:
  90                         del self.other_keys[index]
  91
  92         def updateDirect(self, host):
  93                 cmd = os.popen("/usr/bin/ssh-keyscan -t rsa %s 2>/dev/null" % host)
  94                 line = cmd.read()
  95                 (h,  ip,  key,  comment) = self._split_kh_entry(line[:-1])
  96                 node = plc.getNodes(host)
  97                 (host2, ip2, x, x) = self._record_from_node(node[0])
  98                 rec = { self._get_index(host2, ip2) : "%s %s" % (key, "DIRECT") }
  99
 100                 self.delete(host)
 101                 self.other_keys.update(rec)
 102
 103         def update(self, host):
 104                 node = plc.getNodes(host)
 105                 ret = self._record_from_node(node[0])
 106                 (host, ip, key, comment)  = ret
 107                 if ip == None:
 108                         self.updateDirect(host)
 109                 else:
 110                         rec = { "%s,%s" % (host,ip) : "%s %s" % (key, comment) }
 111                         self.pl_keys.update(rec)
 112
 113         def _record_from_node(self, node, nokey_list=None):
 114                 host = node['hostname']
 115                 key = node['ssh_rsa_key']
 116
 117                 l_nw = plc.getNodeNetworks({'nodenetwork_id':node['nodenetwork_ids']})
 118                 if len(l_nw) == 0:
 119                         # No network for this node. So, skip it.
 120                         return (host, None, None, None)
 121
 122                 ip = l_nw[0]['ip']
 123
 124                 if key == None:
 125                         if nokey_list is not None: nokey_list += [node]
 126                         return (host, ip, None, None)
 127
 128                 key = key.strip()
 129                 # TODO: check for '==' at end of key.
 130                 if key[-1] != '=':
 131                         print "Host with corrupt key! for %s %s" % (node['boot_state'], node['hostname'])
 132
 133                 s_date = time.strftime("%Y/%m/%d_%H:%M:%S",time.gmtime(time.time()))
 134                 #rec = { "%s,%s" % (host,ip) : "%s %s" % (key, "PlanetLab_%s" % (s_date)) }
 135                 #return rec
 136                 return (host, ip, key, "PlanetLab_%s" % s_date)
 137
 138
 139 def main():
 140         k = SSHKnownHosts()
 141         nokey_list = k.updateAll()
 142
 143         for node in nokey_list:
 144                 print "%5s %s" % (node['boot_state'], node['hostname'])
 145
 146 if __name__ == '__main__':
 147         #main()
 148         k = SSHKnownHosts()
 149         #print "update"
 150         #k.update('planetlab-4.cs.princeton.edu')
 151         #print "updateDirect"
 152         k.update(sys.argv[1])
 153         #k.updateDirect(sys.argv[1])
 154         k.write()