#!/usr/bin/python

# dane is a tool to generate and verify HASTLS and TLSA records
# By Paul Wouters <paul@xelerance.com> and Christopher Olah <chris@xelerance.com>
# Copyright 2011 by Xelerance http://www.xelerance.com/
# License: GNU GENERAL PUBLIC LICENSE Version 2 or later
#
# https://datatracker.ietf.org/wg/dane/charter/
# https://datatracker.ietf.org/doc/draft-ietf-dane-protocol/ 

import sys
import binascii
import ssl, socket
import hashlib
import warnings

try:
	import ldnsx
except:
	import daneldnsx as ldnsx

try:
	import sys
	import argparse
except ImportError , e:
	module = str(e)[16:]
	print >> sys.stderr, "dane requires the python module " + module
	if module in ["argparse", "ipcalc", "ldns"]:
		print >> sys.stderr, "Fedora/CentOS: yum install " \
                   + {"argparse":"python-argparse", "ipcalc": "python-ipcalc", "ldns": "ldns-python"}
		print >> sys.stderr, "Debian/Ubuntu: apt-get install " \
                   + {"argparse":"python-argparse", "ipcalc": "python-ipcalc", "ldns": "python-ldns"}
		print >> sys.stderr, "openSUSE: zypper in " \
                   + {"argparse":"python-argparse", "ipcalc": "python-ipcalc", "ldns": "python-ldns"}
	sys.exit(1)

def create_txt(hostname, pubkey):
	""" Kaminsky / Gilmore type TLS pubkey in TXT RRtype -- testing version"""
	warnings.warn("create_txt is untested...")
	hashobj = hashlib.sha1() #Are there other valid hashes?
	hashobj.update(pubkey) #Should we try and get the pubkey ourselves, like in create_TLSA?
	result = hashobj.hexdigest().upper()
	return "%s IN TXT \"v=key1 ha=sha1 h=%s\"" % (hostname, result)

def create_hastls(hostname, fallback_default, services):
	""" Creates a HASTLS RRtype """

def checkExistingTLSA(hostname, certtype, reftype):
	""" check if a TLSA record already exists, to compare and notify if update is needed 
	    UNTESTED!!!! """
	
	warnings.warn("checkExistingTLSA is untested...")
	res = ldnsx.resolver() #Get the appropriate resource records...
	pkt = res.query(hostname, "TYPE65468", tries = 3)
	pres_TLSAs=set(pkt.answer(rr_type="TYPE65468"))
	pres_TLSAs = map(lambda rr: str(rr).strip(), pres_TLSAs)
	
	TLSAs = set(create_TLSAs(hostname, certtype,reftype).split('\n'))
	
	return TLSAs == pres_TLSAs # are things the same as before?
	

def checkExistingHASTLS(address):
	""" check if a HASTLS record already exists, to compare and notify if update is needed """


def create_tlsa(hostname, certtype, reftype):
	"""Creates a TLSA RRtype"""
	""" get all A/AAAA records for hostname """
	if secure:
		pkt = ldnsx.secure_query(hostname, "ANY", flex=True)
	else:
		pkt = ldnsx.query(hostname, "ANY")
	records = pkt.answer( rr_type = { "ipv4":"A",  "ipv6":"AAAA",  "both":"A|AAAA" }[transport] )

	drafts, rfcs = [], []

	for rr in records:
		draft = genTLSA(hostname, rr.ip(), certtype, reftype, draft=True)
		rfc = genTLSA(hostname, rr.ip(), certtype, reftype, draft=False)
		if draft and not (draft in drafts):
			drafts.append(draft.strip())
		if rfc and not (rfc in rfcs):
			rfcs.append(rfc.strip())
	ret = ""
	if not fmt == "rfc":
		ret += "\n".join(drafts)
	if not fmt == "draft":
		ret += "\n".join(rfcs)
	return ret

def genTLSA(hostname, address, certtype, reftype, draft=True):
	try:
		# errors will be thrown already before we get here
		dercert = get_cert(hostname, address)
	except:
		return
	if not dercert:
		return

	if certtype != 1:
		raise Exception("Only EE-cert supported right now")
	certhex = hashCert(reftype, dercert)
	if draft:
		# octet length is half of the string length; remember certtype and reftype are part of the length so +2
		return "_443._tcp.%s IN TYPE65468 \# %s 0%s0%s%s"%(hostname, len(certhex)/2+2, certtype, reftype, certhex )
	else:
		return "_443._tcp.%s IN TLSA %s %s %s"%(hostname, certtype, reftype, certhex)

def get_cert(hostname, address):
	# We don't use ssl.get_server_certificate because it does not support IPv6, and it converts DER to PEM, which
	# we would just have to convert back to DER using ssl.PEM_cert_to_DER_cert()
	try:
		# kinda ugly kludge
		if ":" in address:
			conn = socket.socket(socket.AF_INET6)
		else:
			conn = socket.socket(socket.AF_INET)
		conn.connect((address, 443))
	except socket.error , e:
		#raise Exception("%s (%s): %s"%(hostname, address, str(e)))
		print >> sys.stderr, "%s (%s): %s"%(hostname, address, str(e))
		return
	try:
		sock = ssl.wrap_socket(conn)
	except ssl.SSLError , e:
		#raise Exception("%s (%s): %s"%(hostname, address, str(e)))
		print >> sys.stderr, "%s (%s): %s"%(hostname, address, str(e))
		return
	try:
		dercert = sock.getpeercert(True)
	except AttributeError , e:
		#print >> sys.stderr, "%s (%s): %s"%(hostname, address, str(e))
		return
	sock.close()
	conn.close()
	return dercert

# take PEM encoded EE-cert and DER encode it, then sha256 it
def hashCert(reftype,certblob):
	if reftype == 0:
		return binascii.b2a_hex(certblob).upper()
	elif reftype == 1:
		hashobj = hashlib.sha256()
		hashobj.update(certblob)
	elif reftype == 2:
		hashobj = hashlib.sha512()
		hashobj.update(certblob)
	else:
		return 0
	return hashobj.hexdigest().upper()

secure, transport, quiet, fmt = True, "both", False, "draft"

# create the parser
parser = argparse.ArgumentParser(description='Create TLS related DNS records for hosts or an entire zone. version 1.2.1')

# AXFR
parser.add_argument('-n', '--nameserver', metavar="nameserver", action='append', help='nameserver to query')
parser.add_argument('--axfr', action='store_true', help='use AXFR (all A/AAAA records will be scanned)')

# IETF status related, currently --draft is the default
parser.add_argument('--draft', action='store_true',help='output in draft private rrtype (65468/65469) format (default)')
parser.add_argument('--rfc', action='store_true',help='output in rfc (TLSA/HASTLS) rrtype format')

# TLSA related	
parser.add_argument('--tlsa', action='store_true',help='generate TLSA record (default:yes)')
parser.add_argument('--eecert', action='store_true',help='use EEcert for TLSA record (default)')
parser.add_argument('--cacert', action='store_true',help='use CAcert for TLSA record (not supported yet)')
parser.add_argument('--pubkey', action='store_true',help='use pubkey for TLSA record (not supported yet)')
parser.add_argument('--txt', action='store_true',help='generate Kaminsky style TXT record (not supported yet)')

parser.add_argument('--sha256', action='store_true',help='use SHA256 for the TLSA cert type')
parser.add_argument('--sha512', action='store_true',help='use SHA512 for the TLSA cert type')
parser.add_argument('--full', action='store_true',help='use full certificate for the TLSA cert type')

# allow non-dnssec answers
parser.add_argument('--insecure', action='store_true',help='allow use of non-dnssec answers to find SSL hosts')

# limit networking to ipv4 or ipv6 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('-q', '--quiet', action='store_true',help='suppress warnings and errors')
parser.add_argument('-v', '--version', action='store_true',help='show version and exit')

# finally, the host list
parser.add_argument('hosts', metavar="hostname", nargs='+')

args = parser.parse_args(sys.argv[1:])

if args.version:
	sys.exit("dane: version 1.2.2")
if not args.rfc:
	args.draft = True

if args.cacert:
	sys.exit("TLSA CAcert type record not yet supported")
if args.pubkey:
	sys.exit("TLSA Pubkey type record not yet supported")

if args.sha512:
	reftype=2
elif args.full:
	reftype=0
else:
	reftype=1

if args.quiet:
	quiet = True

if args.tlsa:
	args.eecert = True

if args.insecure:
	secure = False

if args.ipv4 and not args.ipv6:
	transport = "ipv4"
if args.ipv6 and not args.ipv4:
	transport = "ipv6"

# filter the non-options arguments for an "@argument" and convert it for the axfr option.
filterHost= []
if not args.nameserver:
	args.nameserver = []
for host in args.hosts:
	if host[0] == "@":
		args.nameserver.append(host[1:])
		args.hosts.remove(host)
		args.axfr=True
	
if args.rfc and args.draft:
	fmt = "both"
elif args.rfc:
	fmt = "rfc"
else:
	fmt = "draft"

if not args.hosts:
	sys.exit("Host are needed.")
#	main("--help")

for host in args.hosts:
	if host[-1] != ".":
		host += "."
if not args.axfr:
	for host in args.hosts:
		print create_tlsa(host,1,reftype)
if args.axfr:
	# Try and AXFR it
	if len(args.nameserver) == 0:
		sys.exit("nameserver needed. syntax: -n nameserver")
	resolver = ldnsx.resolver(args.nameserver[0], dnssec=True)
	for host in args.hosts:
		ipv4 = []
		ipv6 = []
		for rr in resolver.AXFR(host):
			if rr.rr_type() == "A"    and rr.owner() not in ipv4:
				ipv4.append(rr.owner())
			if rr.rr_type() == "AAAA" and rr.owner() not in ipv6:
				ipv6.append(rr.owner())
		if transport != "ipv6":
			for host in ipv4:
				print create_tlsa(host,1,reftype)
		if transport != "ipv4":
			for host in ipv6:
				print create_tlsa(host,1,reftype)