#!/usr/bin/env ruby

require 'resolv'
require 'ipaddr'

$debug = 0

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

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

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

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

  spf_record
end

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

  spf_record = get_spf_record(domain)

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

  ipv4 = []
  ipv6 = []
  includes = []
  redirects = []
  a_domains = []
  others = []

  skip = 0
  skip = 1 if spf_record.nil? || spf_record.class != String
  skip = 1 if spf_record.class == String && spf_record.empty?
  skip = 1 if spf_record.class == String && !spf_record.start_with?("v=spf1")

  puts("collect_spf_record_details skips for: #{domain}, #{skip}") if $debug > 2

  unless skip == 1
    puts("collect_spf_record_details 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("collect_spf_record_details 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 mean MX or A or AAAA records yet to be implemented

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

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

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

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

  redirects.concat(res["redirects"])

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

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

  redirects
end

def collect_all_includes_recursively(domain)
  puts("collect_all_includes_recursively: #{domain}") if $debug > 2
  includes = []
  res = collect_spf_record_details(domain)
  includes.concat(res["includes"]) if res["includes"].size > 0

  includes.each do |r|
    record = collect_spf_record_details(r.to_s)
    includes.concat(record["includes"]) if record["includes"].size > 0
  end

  includes
end

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

  ipv4 = []

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

  redirects = collect_all_redirects_recursively(domain)

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

  puts("includes recursion started: #{domain}") if $debug > 2

  includes = collect_all_includes_recursively(domain)

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

  base_dom = collect_spf_record_details(domain)

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

  puts("Working on redirects for: #{domain}") if $debug > 2
  redirects.each do |redirect|
    sub_ipv4 = get_all_spf_ip4(redirect)
    sub_ipv4.each { |cidr| ipv4 << cidr }
  end
  puts("Finished working on redirects for: #{domain}") if $debug > 2

  puts("Working on includes for: #{domain}") if $debug > 2
  includes.each do |included|
    sub_ipv4 = get_all_spf_ip4(included)
    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.select { |ip| valid_cidr?(ip) || valid_ip_address?(ip) }

  new_ipv4
end

puts get_all_spf_ip4(domain)