Add a DNSLookupError class (thanks Warren Kumari)

 * This patch makes the getTLSA method catch the DNSLookupError when no TLSA record is found

diff --git a/swede b/swede
index 4be39b51e1c7d060ffe765f8e3def42d785ece0e..2ebaf9bbe4d721e25b0a4d3260edb22f2ad156f2 100755 (executable)
--- a/swede
+++ b/swede
@@ -24,6 +24,7 @@ from binascii import a2b_hex, b2a_hex
 from hashlib import sha256, sha512
 from ipaddr import IPv4Address, IPv6Address
 
+
 def genTLSA(hostname, protocol, port, certificate, output='draft', usage=1, selector=0, mtype=1):
        """This function generates a TLSARecord object using the data passed in the parameters,
        it then validates the record and returns the RR as a string.
@@ -135,7 +136,7 @@ def getRecords(hostname, rrtype='A', secure=True):
                # If we are here the data was either secure or insecure data is accepted
                return result.data.raw
        else:
-               raise Exception('Error: Unsuccesful lookup or no data returned.')
+               raise DNSLookupError('Unsuccesful lookup or no data returned for rrtype %s.' % rrtype)
 
 def getHash(certificate, mtype):
        """Hashes the certificate based on the mtype.
@@ -170,6 +171,9 @@ def getTLSA(hostname, port=443, protocol='tcp', secure=True):
        except InsecureLookupException, e:
                print str(e)
                sys.exit(1)
+       except DNSLookupError, e:
+               print 'Unable to resolve %s: %s' % (hostname, str(e))
+               sys.exit(1)
        ret = []
        for record in records:
                hexdata = b2a_hex(record)
@@ -327,16 +331,13 @@ class AAAARecord:
 
 # Exceptions
 class RecordValidityException(Exception):
-       def __init__(self, value):
-               self.value = value
-       def __str__(self):
-               return self.value
+       pass
 
 class InsecureLookupException(Exception):
-       def __init__(self, value):
-               self.value = value
-       def __str__(self):
-               return self.value
+       pass
+
+class DNSLookupError(Exception):
+       pass
 
 if __name__ == '__main__':
        import argparse