#!/usr/bin/env ruby
# frozen_string_literal: false

require 'pwn'
require 'optparse'
require 'yaml'
require 'json'

opts = {}
OptionParser.new do |options|
  options.banner = "USAGE:
    #{$PROGRAM_NAME} [opts]
  "

  options.on('-cYAML', '--config-yaml=YAML', '<Required - YAML config containing api_key from Shodan.io>') do |y|
    opts[:yaml] = y
  end

  options.on('-qFILE', '--query-file=FILE', '<Required - File containing one Shodan.io query string per line>') do |q|
    opts[:query_file] = q
  end

  options.on('-oFILE', '--output-results-file=FILE', '<Optional - Defaults to /tmp/shodan-results-Time.now.strftime("%y-%m-%d.%H:%M:%S").txt>') do |o|
    opts[:output_results_file] = o
  end
end.parse!

if opts.empty?
  puts `#{$PROGRAM_NAME} --help`
  exit 1
end

begin
  yaml_file = opts[:yaml].to_s
  raise "ERROR: #{yaml_file} does not exist." unless File.exist?(yaml_file)

  yaml = YAML.load_file(yaml_file, symbolize_names: true)

  api_key = yaml[:api_key]

  query_file = opts[:query_file].to_s
  raise "ERROR: #{query_file} does not exist." unless File.exist?(query_file)

  queries = File.readlines(query_file)

  timestamp = Time.now.strftime('%Y-%m-%d.%H:%M:%S')
  query_results_file = opts[:output_results_file]
  query_results_file ||= "/tmp/shodan-results-#{timestamp}"

  raw_query_results_file = "#{query_results_file}-RAW.json"
  File.open(raw_query_results_file, 'w') do |r|
    File.open(query_results_file, 'w') do |f|
      queries.each do |query_line|
        query = query_line.chomp
        print "QUERY: '#{query}'"
        f.puts("QUERY: '#{query}'")
        search_results = PWN::Plugins::Shodan.search(
          api_key: api_key,
          query: query
        )
        puts " >>> Matches: #{search_results[:total]}"

        raw_results = {}
        raw_results[:query] = query
        raw_results[:results] = search_results
        r.puts raw_results.to_json

        search_results[:matches].select do |m|
          f.puts "ORG: #{m[:org]} | PUBIP: #{m[:ip_str]} #{'*' * 36}"
          f.puts "Product: #{m[:product]}"
          f.puts "TCP Port: #{m[:port]}"
          f.puts "Data: #{m[:data]}\n\n\n"
        end
      end
    end
  end
rescue SystemExit, Interrupt
  puts "\nGoodbye."
end