aquatone-discover

#!/usr/bin/env ruby

require "aquatone"

options = {
  :fallback_nameservers => %w(8.8.8.8 8.8.4.4),
  :threads => 5,
  :ignore_private => true
}

OptionParser.new do |opts|
  opts.banner = "Usage: aquatone-discover OPTIONS"

  opts.on("-d", "--domain DOMAIN", "Domain name to assess") do |v|
    if !Aquatone::Validation.valid_domain_name?(v)
      puts "#{v} doesn't look like a valid domain name."
      exit 1
    end
    options[:domain] = v
  end

  opts.on("--nameservers NAMESERVERS", "Nameservers to use") do |v|
    ips = v.split(",").map(&:strip).uniq
    if ips.empty?
      puts "Nameservers can't be empty."
      exit 1
    end
    ips.each do |ip|
      if !Aquatone::Validation.valid_ip?(ip)
        puts "#{ip} is not a valid IP address."
        exit 1
      end
    end
    options[:nameservers] = ips
  end

  opts.on("--fallback-nameservers NAMESERVERS", "Nameservers to fall back to") do |v|
    ips = v.split(",").map(&:strip).uniq
    if ips.empty?
      puts "Fallback nameservers can't be empty."
      exit 1
    end
    ips.each do |ip|
      if !Aquatone::Validation.valid_ip?(ip)
        puts "#{ip} is not a valid IP address."
        exit 1
      end
    end
    options[:fallback_nameservers] = ips
  end

  opts.on("--[no-]ignore-private", "Ignore hosts resolving to private IP addresses") do |v|
    options[:ignore_private] = v
  end

  opts.on("--set-key KEY VALUE", "Save a key to key store") do |key|
    if ARGV.empty?
      puts "No key value given."
      exit 1
    end
    key   = key.to_s.downcase
    value = ARGV.join(" ").strip
    Aquatone::KeyStore.set(key, value)
    puts "Saved key #{key} with value #{value}."
    exit
  end

  opts.on("--list-collectors", "See information on collectors") do
    Aquatone::Collector.descendants.each do |collector|
      puts "Name............: #{collector.meta[:name]}"
      puts "Description.....: #{collector.meta[:description]}" if collector.meta[:description]
      puts "Author..........: #{collector.meta[:author]}"
      puts "Key Requirements: #{collector.meta.key?(:require_keys) ? collector.meta[:require_keys].join(', ') : 'None'}"
      puts "Key.............: #{collector.sluggified_name}\n\n"
      puts "--------------------------------------------------\n\n"
    end
    exit
  end

  opts.on("--only-collectors COLLECTORS", "Only run specified collectors") do |v|
    known_collectors = Aquatone::Collector.descendants.map(&:sluggified_name)
    collectors       = v.split(",").map(&:strip).uniq
    collectors.each do |collector|
      if !known_collectors.include?(collector)
        puts "Unknown collector key: #{collector}"
        exit 1
      end
    end
    options[:only_collectors] = collectors
  end

  opts.on("--disable-collectors COLLECTORS", "Disable specified collectors") do |v|
    known_collectors = Aquatone::Collector.descendants.map(&:sluggified_name)
    collectors       = v.split(",").map(&:strip).uniq
    collectors.each do |collector|
      if !known_collectors.include?(collector)
        puts "Unknown collector key: #{collector}"
        exit 1
      end
    end
    options[:disable_collectors] = collectors
  end

  opts.on("-t", "--threads THREADS", "Number of concurrent threads to use") do |v|
    options[:threads] = v.to_i
  end

  opts.on("-s", "--sleep SECONDS", "Seconds to sleep between lookups") do |v|
    if v.to_i < 1
      puts "Sleep can't be less than 1"
      exit 1
    end
    options[:sleep] = v.to_i
  end

  opts.on("-j", "--jitter PERCENTAGE", "Jitter factor for sleep intervals") do |v|
    if !v.to_i.between?(1, 100)
      puts "Jitter factor must be between 1 and 100"
      exit 1
    end
    options[:jitter] = v.to_f
  end

  Aquatone::Collector.descendants.each do |collector|
    collector.cli_options.each_pair do |option, description|
      opts.on("--#{option}", description) do |v|
        options[option.split(" ").first.gsub("-", "_").to_sym] = v
      end
    end
  end

  opts.on("-h", "--help", "Show help") do
    puts opts
    exit 0
  end
end.parse!

Aquatone::Commands::Discover.run(options)