#!/usr/bin/env ruby

#require 'spf/query'
require 'ipaddr'

domain = ARGV[0]

$debug = 0

def getSpfRecord(domain)
	puts("getSpfRecord domain: #{domain}") if $debug > 2

#	spf_record = SPF::Query::Record.query(domain).to_s

	command = "dig +short TXT #{domain} | grep -E '^\"v=spf1'"
	output = `#{command}`

	spf_record = output.strip.gsub("\"","")

	puts("getSpfRecord spf_record for: #{domain}, #{spf_record}") if $debug > 2

	return spf_record
end


def collectSpfRecordDetails(domain)
	puts("collectSpfRecordDetails: #{domain}") if $debug > 2

	spf_record = getSpfRecord(domain)

	puts("collectSpfRecordDetails spf_record for: #{domain}, #{spf_record}") if $debug > 2

	ipv4 = []
 	ipv6 = []
	includes = []
	redirects = []
	a_domains = []
	others = []
	
	skip = 0
	skip = 1 if spf_record.nil? or spf_record.class != String
	skip = 1 if spf_record.class == String and spf_record.empty?
	skip = 1 if spf_record.class == String and not spf_record.start_with?("v=spf1")
	
	puts("collectSpfRecordDetails skips for: #{domain}, #{skip}") if $debug > 2

	if skip == 0
		puts("collectSpfRecordDetails domain: #{domain}") if $debug > 1
	
		redirect_regex = /redirect\=([a-z\.\-\_0-9]+)/
		redirects_res = spf_record.scan(redirect_regex)
		redirects_res.each { |r| redirects << r[0] }

		puts("collectSpfRecordDetails redirects for: #{domain}, #{redirects}") if $debug > 2


		include_regex = /include\:([a-z\.\-\_0-9]+)/
		includes_res = spf_record.scan(include_regex)
		includes_res.each { |i| includes << i[0] }

		a_domains_regex = /a\:([a-z\.\-\_0-9]+)/
		a_domains_res = spf_record.scan(a_domains_regex)
		a_domains_res.each { |a| a_domains << a[0] }

		ip4_regex = /ip4\:([0-9\.\/]+)/
		ipv4_res = spf_record.scan(ip4_regex)
		ipv4_res.each { |cidr| ipv4 << cidr[0] }

		ip6_regex = /ip6\:([a-f0-9\:\/]+)/
		ipv6_res = spf_record.scan(ip6_regex)
		ipv6_res.each { |cidr| ipv6 << cidr[0] }
	end
	# others means mx or a or aaaa records yet to be implemented

	return {"ipv4" => ipv4, "ipv6" => ipv6, "includes" => includes, "redirects" => redirects, "others" => others, "a" => a_domains }

end

def collectAllRedirectsRecursivly(domain)
	puts("collectAllRedirectsRecursivly: #{domain}") if $debug > 2

	redirects = []
	res = collectSpfRecordDetails(domain)
	puts("collectAllRedirectsRecursivly res for: #{domain}, #{res}") if $debug > 2

	puts("collectAllRedirectsRecursivly redirects for: #{domain}, #{redirects}") if $debug > 2

	if res["redirects"].size > 0
		res["redirects"].each { |r| redirects << r}
	end

	puts("collectAllRedirectsRecursivly redirects for: #{domain}, #{redirects}") if $debug > 2

	redirects.each do |r|
		record = collectSpfRecordDetails(r.to_s)
	        if record["redirects"].size > 0
        	        record["redirects"].each { |r| redirects << r}
	        end
	end
	puts("collectAllRedirectsRecursivly redirects: #{redirects}") if $debug > 2

	return redirects
end

def collectAllIncludesRecursivly(domain)
	puts("collectAllIncludesRecursivly: #{domain}") if $debug > 2
	includes = []
	res = collectSpfRecordDetails(domain)
	if res["includes"].size > 0
		res["includes"].each { |i| includes << i}
	end

	includes.each do |r|
	        record = collectSpfRecordDetails(r.to_s)
	        if record["includes"].size > 0
        		record["includes"].each { |i| includes << i}
	        end
	end

	return includes
end

def getAllSpfIp4(domain)
	puts("getAllSpfIp: #{domain}") if $debug > 2

	ipv4 = []

	puts("redirects recurse started: #{domain}") if $debug > 2

	redirects = collectAllRedirectsRecursivly(domain)

	puts("redirects recurse ended: #{domain}") if $debug > 2

	puts("indluess recurse started: #{domain}") if $debug > 2

	includes = collectAllIncludesRecursivly(domain)

	puts("includes recurse ended: #{domain}") if $debug > 2


	base_dom = collectSpfRecordDetails(domain)

	
	puts("redirects found: #{redirects}") if $debug > 2

	puts("includes found: #{includes}") if $debug > 2


	puts("Starting to work on redirects for: #{domain}") if $debug > 2

	redirects.each do |redirect|
		sub_ipv4 = getAllSpfIp4(redirect)
		puts("Working in redirect: #{redirect}") if $debug > 2
		puts("Working in redirect: #{redirect}, #{sub_ipv4}") if $debug > 2

		sub_ipv4.each { |cidr| ipv4 << cidr }

	end
	puts("Finished working on redirects for: #{domain}") if $debug > 2

	puts("Starting to work on includes for: #{domain}") if $debug > 2

	includes.each do |included|
		sub_ipv4 = getAllSpfIp4(included)
		puts("Working on include: #{included}") if $debug > 2
		puts("Working in include: #{included}, #{sub_ipv4}") if $debug > 2

		sub_ipv4.each { |cidr| ipv4 << cidr }
	end

	puts("Finished working on includes for: #{domain}") if $debug > 2



	base_dom["ipv4"].each { |cidr| ipv4 << cidr }

	# filter ipv4 list
	new_ipv4 = []
	ipv4.each do |ip|
		if valid_cidr?(ip) or valid_ip_address?(ip)
			new_ipv4 << ip
		end
	end

	return new_ipv4
end


def valid_cidr?(input)
  IPAddr.new(input)
  true
rescue IPAddr::InvalidAddressError
  false
end

def valid_ip_address?(input)
  IPAddr.new(input)
  true
rescue IPAddr::InvalidAddressError
  false
end

puts getAllSpfIp4(domain)