#!/usr/bin/env ruby

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

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

#Default options
options = {
  :sockets => [],
  :policy => File.expand_path("../../policies/mozilla_modern.yml", __FILE__),
  :unit_test => false,
  :timeout => 2,
  :threads => 5,
}

target_parser = SSHScan::TargetParser.new()

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/Range/Hostname]", Array,
          "IP/Ranges/Hostname to scan") do |ips|
    ips.each do |ip|
      options[:sockets] += target_parser.enumerateIPRange(ip)
    end
  end

  opts.on("-f", "--file [FilePath]",
          "File Path of the file containing IP/Range/Hostnames to scan") do |file|
    unless File.exists?(file)
      puts "\nReason: input file supplied is not a file"
      exit
    end
    File.open(file).each do |line|
      line.chomp.split(',').each do |socket|
        ip, port = socket.chomp.split(':')
        port = port.nil? ? 22 : port
        options[:sockets] += target_parser.enumerateIPRange(ip, port)
      end
    end
  end

  opts.on("-T", "--timeout [seconds]",
          "Timeout per connect after which ssh_scan gives up on the host") do |timeout|
    options[:timeout] = timeout.to_i
  end

  opts.on("-O", "--from_json [FilePath]",
          "File to read JSON output from") do |file|
    unless File.exists?(file)
      puts "\nReason: Invalid file"
      exit
    end
    file = open(file)
    json = file.read
    parsed_json = JSON.parse(json)
    parsed_json.each do |host|
      options[:sockets] += target_parser.enumerateIPRange(host['ip'], host['port'])
    end
  end

  opts.on("-o", "--output [FilePath]",
          "File to write JSON output to") do |file|
    $stdout.reopen(file, "w")
  end

  opts.on("-p", "--port [PORT]",
          "Port (Default: 22)") do |port|
    socket = options[:sockets].shift
    ip = socket.chomp.split(':').shift
    options[:sockets] += target_parser.enumerateIPRange(ip, port)
  end

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

  opts.on("--threads [NUMBER]",
          "Number of worker threads (Default: 5)") do |threads|
    options[:threads] = threads.to_i
  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 -t ::1 -T 5"
    puts "  ssh_scan -f hosts.txt"
    puts "  ssh_scan -o output.json"
    puts "  ssh_scan -O output.json -o rescan_output.json"
    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[:sockets].nil?
  puts opt_parser.help
  puts "\nReason: no target specified"
  exit 1
end

options[:sockets].each do |socket|
  ip, port = socket.chomp.split(':')
  unless ip.ip_addr? || ip.fqdn?
    puts opt_parser.help
    puts "\nReason: #{socket} is not a valid target"
    exit 1
  end
end

options[:sockets].each do |socket|
  ip, port = socket.chomp.split(':')
  unless (0..65535).include?(port.to_i)
    puts opt_parser.help
    puts "\nReason: port supplied is not within acceptable range"
    exit 1
  end
end

unless File.exists?(options[:policy])
  puts opt_parser.help
  puts "\nReason: policy file supplied is not a file"
  exit 1
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)

if options[:unit_test] == true
  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
end