enable ipv4 or ipv6 only options

diff --git a/swede b/swede
index 244cef94a5fddeba1f7b5026989a6db9a0352a36..bcf4d40c6905e892891968c0a0081a9d5ce1bf16 100755 (executable)
--- a/swede
+++ b/swede
@@ -25,6 +25,9 @@ from binascii import a2b_hex, b2a_hex
 from hashlib import sha256, sha512
 from ipaddr import IPv4Address, IPv6Address
 
+check_ipv4=True
+check_ipv6=True
+
 
 def genTLSA(hostname, protocol, port, certificate, output='generic', usage=1, selector=0, mtype=1):
        """This function generates a TLSARecord object using the data passed in the parameters,
@@ -59,6 +62,7 @@ def genTLSA(hostname, protocol, port, certificate, output='generic', usage=1, se
        return record.getRecord()
 
 def getA(hostname, secure=True):
+       if not check_ipv4: return []
        """Gets a list of A records for hostname, returns a list of ARecords"""
        try:
                records = getRecords(hostname, rrtype='A', secure=secure)
@@ -74,6 +78,7 @@ def getA(hostname, secure=True):
        return ret
 
 def getAAAA(hostname, secure=True):
+       if not check_ipv6: return []
        """Gets a list of A records for hostname, returns a list of AAAARecords"""
        try:
                records = getRecords(hostname, rrtype='AAAA', secure=secure)
@@ -389,8 +394,8 @@ if __name__ == '__main__':
        parser_create = subparsers.add_parser('create', help='Create a TLSA record')
        parser_create.set_defaults(function='create')
 
-       #parser.add_argument('-4', dest='ipv4', action='store_true',help='use ipv4 networking only')
-       #parser.add_argument('-6', dest='ipv6', action='store_true',help='use ipv6 networking only')
+       parser.add_argument('-4', dest='ipv4', action='store_true',help='use ipv4 networking only')
+       parser.add_argument('-6', dest='ipv6', action='store_true',help='use ipv6 networking only')
        parser.add_argument('--insecure', action='store_true', default=False, help='Allow use of non-dnssec secured answers')
        parser.add_argument('--resolvconf', metavar='/PATH/TO/RESOLV.CONF', action='store', default='', help='Use a recursive resolver from resolv.conf')
        parser.add_argument('-v', '--version', action='version', version='%(prog)s v0.2', help='show version and exit')
@@ -413,6 +418,15 @@ if __name__ == '__main__':
        parser_create.add_argument('--mtype', '-m', action='store', type=int, default=1, choices=[0,1,2], help='The Matching Type of the Certificate for Association. \'0\' for Exact match, \'1\' for SHA-256 hash, \'2\' for SHA-512 (default: %(default)s).')
 
        args = parser.parse_args()
+       import pprint
+       pprint.pprint(args)
+       if args.ipv4 == True and args.ipv6 == True: 
+               print "Cannot have only ipv4 and only ipv6 at the same time"
+               sys.exit()
+       elif args.ipv4 == True:
+               check_ipv6 = False
+       elif args.ipv6 == True:
+               check_ipv4 = False
 
        if args.host[-1] != '.':
                args.host += '.'
@@ -465,7 +479,13 @@ if __name__ == '__main__':
 
                        if not args.quiet:
                                print 'Attempting to verify the record with the TLS service...'
-                       addresses = getA(args.host, secure=secure) + getAAAA(args.host, secure=secure)
+                       if check_ipv4 and check_ipv6:
+                               addresses = getA(args.host, secure=secure) + getAAAA(args.host, secure=secure)
+                       elif check_ipv4:
+                               addresses = getA(args.host, secure=secure) 
+                       else:
+                               addresses = getAAAA(args.host, secure=secure)
+                               
                        for address in addresses:
                                if not args.quiet:
                                        print 'Got the following IP: %s' % str(address)
@@ -480,9 +500,12 @@ if __name__ == '__main__':
                                        sys.exit(1)
                                # Don't error when the verification fails in the SSL handshake
                                ctx.set_verify(SSL.verify_none, depth=9)
-                               if isinstance(address, AAAARecord):
+                               if  check_ipv6 and isinstance(address, AAAARecord):
                                        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
                                        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                               elif  check_ipv4 and isinstance(address, ARecord):
+                                       sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                                       sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                                else:
                                        sock = None
                                connection = SSL.Connection(ctx, sock=sock)
@@ -601,7 +624,13 @@ if __name__ == '__main__':
                                                sys.stdout.write('Port %s not numerical or within correct range (1 <= port <= 65535), try again (hit enter for default 443): ' % user_input)
                        # Get the address records for the host
                        try:
-                               addresses = getA(args.host, secure=secure) + getAAAA(args.host, secure=secure)
+                               if check_ipv4 and check_ipv6:
+                                       addresses = getA(args.host, secure=secure) + getAAAA(args.host, secure=secure)
+                               elif check_ipv4:
+                                       addresses = getA(args.host, secure=secure) 
+                               else:
+                                       addresses = getAAAA(args.host, secure=secure)
+
                        except InsecureLookupException, e:
                                print >> sys.stderr, str(e)
                                sys.exit(1)
@@ -611,9 +640,12 @@ if __name__ == '__main__':
                                # We do the certificate handling here, as M2Crypto keeps segfaulting when try to do stuff with the cert if we don't
                                ctx = SSL.Context()
                                ctx.set_verify(SSL.verify_none, depth=9)
-                               if isinstance(address, AAAARecord):
+                               if check_ipv6 and isinstance(address, AAAARecord):
                                        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
                                        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                               if check_ipv4 and isinstance(address, ARecord):
+                                       sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                                       sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                                else:
                                        sock = None
                                connection = SSL.Connection(ctx, sock=sock)