#!/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.expand_path("../../policies/mozilla_modern.yml", __FILE__),
  :unit_test => false,
  :timeout => 2,
  :threads => 5,
  :verbosity => nil,
  :logger => Logger.new(STDERR),
  :fingerprint_database => "./fingerprints.db"
}

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.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(':')
        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.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]", 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",
          "Set the logger level (Accepted Params: INFO, DEBUG, WARN, ERROR, FATAL)") do |verbosity|
    options[:logger].level = case options[:verbosity]
      when "INFO" then Logger::INFO
      when "DEBUG" then Logger::DEBUG
      when "WARN" then Logger::WARN
      when "ERROR" then Logger::ERROR
      when "FATAL" then Logger::FATAL
      else
        puts "Can't convert #{options[:verbosity]} to any of the Logger level constants"
        exit
    end
  end

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

  opts.on("-l", "--listen", "Listen and serve API requests") do
    SSHScan::API.run!
    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

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

# 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.size > 0
      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

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