#!/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'
require 'logger'

#Default options
options = {
  "sockets" => [],
  "policy" => File.join(File.dirname(__FILE__),"../config/policies/mozilla_modern.yml"),
  "unit_test" => false,
  "timeout" => 2,
  "threads" => 5,
  "verbosity" => nil,
  "logger" => Logger.new(STDERR),
  "fingerprint_database" => File.join(File.dirname(__FILE__),"../data/fingerprints.yml")
}

# Reorder arguments before parsing
def reorder_args!(order, opt_parser)
  old_args = opt_parser.default_argv
  new_args = []
  len = opt_parser.default_argv.length
  order.each do |next_keyset|
    i = 0
    (0...len).each do
      if next_keyset.include?(opt_parser.default_argv[i])
        new_args << old_args.delete_at(i) << old_args.delete_at(i)
      else
        i += 1
      end
    end
  end
  new_args += old_args
  opt_parser.default_argv = new_args
end

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 |sockets|
    sockets.each do |socket|
      ip, port = socket.chomp.split(':')
      options["sockets"] += target_parser.enumerateIPRange(ip, port)
    end
  end

  opts.on("-f", "--file [FilePath]",
          "File Path of the file containing IP/Range/Hostnames to \
scan") do |file|
    unless File.exist?(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(':')
        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("-L", "--logger [Log File Path]",
          "Enable logger") do |log_file|
    if log_file.nil?
      options["logger"] = Logger.new(STDERR)
    else
      options["logger"] = Logger.new $stdout.reopen(log_file, "w")
    end
  end

  opts.on("-O", "--from_json [FilePath]",
          "File to read JSON output from") do |file|
    unless File.exist?(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]", Array,
          "Port (Default: 22)") do |ports|
    temp = []
    options["sockets"].each do |socket|
      ports.each do |port|
        ip, old_port = socket.chomp.split(':')
        if !old_port.nil?
          puts "Specifying port simultaneously with -t and -p is not\
 allowed. Please fix this and try again"
          exit 1
        end
        temp += target_parser.enumerateIPRange(ip, port)
      end
    end
    options["sockets"] = temp
  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("--fingerprint-db [FILE]",
          "File location of fingerprint database (Default: \
./fingerprints.db)") do |fingerprint_db|
    options["fingerprint_database"] = fingerprint_db
  end

  opts.on("--suppress-update-status", "Do not check for updates") do
    options["suppress_update_status"] = true
  end

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

  opts.on("-V", "--verbosity [STD_LOGGING_LEVEL]") do |verb|
    case verb
    when "INFO"
      options["logger"].level ==  Logger::INFO
    when "WARN"
      options["logger"].level ==  Logger::WARN
    when "DEBUG"
      options["logger"].level ==  Logger::DEBUG
    when "ERROR"
      options["logger"].level ==  Logger::ERROR
    when "FATAL"
      options["logger"].level ==  Logger::FATAL
    else
      options["logger"].fatal("Unrecognized logging verbosity level #{verb}")
      exit 1
    end
  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 22222 -L output.log -V INFO"
    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

reorder_args!([["-t", "--target"], ["-p", "--port"]], opt_parser)
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 = socket.chomp.split(':')[0]
  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|
  port = socket.chomp.split(':')[1]
  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.exist?(options["policy"])
  puts opt_parser.help
  puts "\nReason: policy file supplied is not a file #{options["policy"]}"
  exit 1
end

# Check to see if we're running the latest released version
#if !options["suppress_update_status"]
#  update = SSHScan::Update.new
#  if update.newer_gem_available?
#    options["logger"].warn(
#      "You're NOT using the latest version of ssh_scan, try 'gem update \
#ssh_scan' to get the latest"
#    )
#  else
#    if update.errors.any?
#      update.errors.each do |error|
#        options["logger"].error(error)
#      end
#    else
#      options["logger"].info(
#        "You're using the latest version of ssh_scan #{SSHScan::VERSION}"
#      )
#    end
#  end
#end

# Limit scope of fingerprints DB to (per scan)
if options["fingerprint_database"] && File.exists?(options["fingerprint_database"])
  File.unlink(options["fingerprint_database"])
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.compliant == false
      exit 1 #non-zero means a false
    else
      exit 0 #non-zero means pass
    end
  end
end