#!/usr/bin/env ruby

# Path setting slight of hand
$:.unshift File.join(File.dirname(__FILE__), "../lib")

require 'ssh_scan'
require 'optparse'
require 'netaddr'

#Default options
options = {
  :targets => [],
  :port => 22,
  :policy => File.expand_path("../../policies/mozilla_modern.yml", __FILE__),
  :unit_test => false
}
opt_parser = OptionParser.new do |opts|
  opts.banner = "ssh_scan v#{SSHScan::VERSION} (https://github.com/mozilla/ssh_scan)\n\n" +
                "Usage: ssh_scan [options]"

  opts.on("-t", "--target [IP/Hostname]", Array,
          "IP/Hostname (IPv4/IPv6/FQDNs)") do |ips|
    ips.each do |ip|
      if ip.fqdn?
        options[:targets] += [ip]
      else
        options[:targets] += NetAddr::CIDR.create(ip).enumerate
      end
    end
  end

  opts.on("-f", "--file [FilePath]",
          "File Path of the file containing IPs") do |file|
    txt = open(file)
    options[:targets] = txt.read.chomp.split(',')
  end

  opts.on("-o", "--outputFile [FilePath]",
          "Writing JSON documents to disk") do |file|
    $stdout.reopen(file, "w")
  end

  opts.on("-p", "--port [PORT]",
          "Port (Default: 22)") do |port|
    options[:port] = port.to_i
  end

  opts.on("-P", "--policy [FILE]",
          "Policy file (Default: Mozilla Modern)") do |policy|
    options[:policy] = policy
  end

  opts.on("-u", "--unit-test [FILE]",
          "Throw appropriate exit codes based on compliance status") do
    options[:unit_test] = true
  end

  opts.on("-v", "--version",
          "Display just version info") do
    puts SSHScan::VERSION
    exit
  end

  opts.on_tail("-h", "--help", "Show this message") do
    puts opts
    puts "\nExamples:"
    puts "\n  ssh_scan -t 192.168.1.1"
    puts "  ssh_scan -t server.example.com"
    puts "  ssh_scan -t ::1"
    puts "  ssh_scan -f filePath"
    puts "  ssh_scan -o filePath"
    puts "  ssh_scan -t 192.168.1.1 -p 22222"
    puts "  ssh_scan -t 192.168.1.1 -P custom_policy.yml"
    puts "  ssh_scan -t 192.168.1.1 --unit-test -P custom_policy.yml"
    puts ""
    exit
  end
end

opt_parser.parse!

if options[:targets].nil?
  puts opt_parser.help
  puts "\nReason: no target specified"
  exit
end

options[:targets].each do |target|
  unless target.ip_addr? || target.fqdn?
    puts opt_parser.help
    puts "\nReason: #{options[:targets]} is not a valid target"
    exit
  end
end

unless (0..65535).include?(options[:port])
  puts opt_parser.help
  puts "\nReason: port supplied is not within acceptable range"
  exit
end

unless File.exists?(options[:policy])
  puts opt_parser.help
  puts "\nReason: policy file supplied is not a file"
  exit
end

options[:policy_file] = SSHScan::Policy.from_file(options[:policy])

# Perform scan and get results
scan_engine = SSHScan::ScanEngine.new()
results = scan_engine.scan(options)

puts JSON.pretty_generate(results)

results.each do |result|
  if result["compliance"] && result["compliance"][:compliant] == false
    exit 1 #non-zero means a false
  else
    exit 0 #non-zero means pass
  end
end