Merge pull request #1 from riltsken/add-multiple-sshkeys

Allow multiple sshkeys for hesiod53

hesiod53/ssh.py

@@ -1,26 +1,36 @@
 #!/usr/bin/env python
-
 import argparse
 import dns.resolver
 import time
 
-def fetch_ssh_keys(username, hesiod_domain, tries=3):
-    fqdn = username + ".ssh." + hesiod_domain
+# Meant for handling dns query failures
+def retry(func, max_retries=3):
+   def func_wrapper(*args, **kwargs):
+        for attempt in range(0, max_retries):
+            try:
+                return func(*args, **kwargs)
+            except:
+                time.sleep(0.1)
+        return
+   return func_wrapper
+
+def concatenate_txt_record(answers):
     results = []
+    for rdata in answers:
+        results.append("".join(rdata.strings))
+    return results
 
-    try:
-        # use TCP because ssh keys are too big for UDP anyways, so we'll back
-        # off and try TCP eventually anyways
-        answers = dns.resolver.query(fqdn, "TXT", tcp=True)
-        for rdata in answers:
-            results.append("".join(rdata.strings))
-    except:
-        # retry, in case of dns failure
-        if tries > 1:
-            time.sleep(0.3)
-            return fetch_ssh_keys(username, hesiod_domain, tries - 1)
+@retry
+def fetch_ssh_key_count(username, hesiod_domain):
+    fqdn = "{username}.count.ssh.{domain}".format(username=username, domain=hesiod_domain)
+    answers = dns.resolver.query(fqdn, "TXT", tcp=True)
+    return int(concatenate_txt_record(answers)[0])
 
-    return results
+@retry
+def fetch_ssh_key(username, hesiod_domain, _id):
+    fqdn = "{username}.{id}.ssh.{domain}".format(username=username, id=_id, domain=hesiod_domain)
+    answers = dns.resolver.query(fqdn, "TXT", tcp=True)
+    return concatenate_txt_record(answers)
 
 def find_hesiod_domain(hesiod_conf_file):
     lhs = None
@@ -58,8 +68,9 @@ def main():
     username = args.username
     hesiod_domain = find_hesiod_domain(args.hesiod_conf_file)
 
-    for key in fetch_ssh_keys(username, hesiod_domain):
-        print key
+    for _id in range(0, fetch_ssh_key_count(username, hesiod_domain)):
+        for key in fetch_ssh_key(username, hesiod_domain, _id):
+            print key
 
 if __name__ == "__main__":
     main()

hesiod53/sync.py

@@ -10,7 +10,7 @@
 import yaml
 
 # a DNS record for hesiod
-# fqdn should includ the trailing .
+# fqdn should include the trailing .
 # value should contain the value without quotes
 DNSRecord = namedtuple("DNSRecord", "fqdn value")
 
@@ -123,10 +123,17 @@ def dns_records(self, hesiod_domain):
         records.append(DNSRecord(fqdn, ":".join(gl)))
 
         # ssh records
-        for ssh_key in self.ssh_keys:
-            fqdn = "%s.ssh.%s" % (self.username, hesiod_domain)
-            records.append(DNSRecord(fqdn, ssh_key))
+        if self.ssh_keys:
+            ssh_keys_count_fqdn = "%s.count.ssh.%s" % (self.username, hesiod_domain)
+            records.append(DNSRecord(ssh_keys_count_fqdn, str(len(self.ssh_keys))))
 
+            # Need to keep this around for backwards compatibility when only one ssh key worked
+            legacy_ssh_key_fqdn = "%s.ssh.%s" % (self.username, hesiod_domain)
+            records.append(DNSRecord(legacy_ssh_key_fqdn, self.ssh_keys[0]))
+
+            for _id, ssh_key in enumerate(self.ssh_keys):
+                ssh_key_fqdn = "%s.%s.ssh.%s" % (self.username, _id, hesiod_domain)
+                records.append(DNSRecord(ssh_key_fqdn, ssh_key))
         return records
 
     def passwd_line(self):