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

require 'time'
require 'pwn'
require 'optparse'
require 'uri'
require 'cgi'
require 'htmlentities'
require 'faker'
require 'json'

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

  options.on('-uURL', '--target-url=URL', '<Required - Target URL)>') do |t|
    opts[:target_url] = t
  end

  options.on('-wFILE', '--word-list=FILE', '<Required - Wordlist File to Use>') do |w|
    opts[:wordlist] = w
  end

  options.on('-aPATTERN', '--append=PATTERN', '<Optional - Append Pattern to Each Entry in the Wordlist (e.g .inc, .BAK, .git, etc>') do |a|
    opts[:append_pattern] = a
  end

  options.on('-TFLOAT', '--max-timeout=FLOAT', '<Optional - HTTP Request Timeout (Default 6.0 Seconds)>') do |t|
    opts[:max_timeout] = t
  end

  options.on('-TFLOAT', '--max-retries=FLOAT', '<Optional - HTTP Request Retries (Default 3)>') do |r|
    opts[:max_retries] = r
  end

  options.on('-pPROXY', '--proxy=PROXY', '<Optional - Proxy SCHEME://ADDRESS:PORT || tor>') do |p|
    opts[:proxy] = p
  end

  options.on('-tTHREADS', '--max-threads=THREADS', '<Optional # HTTP Requests to Run Simultaneously (Default 25)>') do |t|
    opts[:max_threads] = t
  end

  options.on('-rHEADERS', '--request-headers=HEADERS', '<Optional - HTTP Request Headers in JSON Format)>') do |h|
    opts[:http_request_headers] = h
  end

  options.on('-ICODES', '--include-response-codes=CODES', '<Optional - Comma-Delimited List of Response Codes to Include in Report)>') do |i|
    opts[:include_http_response_codes] = i
  end

  options.on('-ECODES', '--exclude-response-codes=CODES', '<Optional - Comma-Delimited List of Response Codes to Exclude in Report)>') do |e|
    opts[:exclude_http_response_codes] = e
  end

  options.on('-dDIR', '--dir-path=DIR', '<Optional - Report Output Directory (Defaults to ".")>') do |w|
    opts[:wordlist] = w
  end

  options.on('-nREPORTNAME', '--report-name=REPORTNAME', '<Optional Report Name (Defaults to, "<TARGET_URL_HOST>-<File.basename(wordlist)>-<Time.now.strftime("%Y-%m-%d_%H-%M-%S")>")>') do |n|
    opts[:report_name] = n
  end

  options.on('-s', '--[no-]start-reporting-server', '<Optional - Start Simple HTTP Server for Reporting>') do |s|
    opts[:start_reporting_server] = s
  end
end.parse!

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

def request_path(opts = {})
  target_url = opts[:target_url]
  proxy = opts[:proxy]
  http_request_headers = opts[:http_request_headers]
  wordlist_line = opts[:wordlist_line]
  http_method = opts[:http_method]
  max_timeout = opts[:max_timeout]
  timeout = opts[:max_timeout] ||= max_timeout / 21.0
  max_retries = opts[:max_retries]

  rest_client_resp_hash = {}

  # request_count && timeout values to begin with
  request_count = 1

  if proxy
    browser_obj = PWN::Plugins::TransparentBrowser.open(
      browser_type: :rest,
      proxy: proxy
    )
  else
    browser_obj = PWN::Plugins::TransparentBrowser.open(
      browser_type: :rest
    )
  end
  rest_client = browser_obj[:browser]::Request

  headers = {
    user_agent: Faker::Internet.user_agent
  }

  http_uri = "#{target_url}/#{wordlist_line}"

  if http_request_headers
    headers = JSON.parse(
      http_request_headers,
      symbolize_names: true
    )
  end

  timestamp_fmt = '%Y-%m-%d %H:%M:%S.%L'

  begin
    request_timestamp = Time.now.strftime(timestamp_fmt)
    response = rest_client.execute(
      method: http_method,
      url: http_uri,
      headers: headers,
      verify_ssl: false,
      timeout: timeout
    )
    response_timestamp = Time.now.strftime(timestamp_fmt)

    # duration in milliseconds
    duration = Time.parse(response_timestamp) - Time.parse(request_timestamp)

    rest_client_resp_hash = {
      request_timestamp: request_timestamp,
      response_timestamp: response_timestamp,
      duration: duration,
      http_uri: http_uri,
      http_method: http_method,
      http_resp_code: response.code,
      http_resp_length: response.body.length,
      http_resp_headers: JSON.pretty_generate(response.headers),
      http_resp: "#{response.body[0..300]}..."
    }
  rescue Errno::ECONNREFUSED,
         Errno::ECONNRESET,
         NoMethodError,
         OpenSSL::SSL::SSLError,
         RestClient::Exceptions::ReadTimeout,
         RestClient::Exceptions::OpenTimeout,
         RestClient::ServerBrokeConnection,
         SOCKSError => e

    timeout += 0.1
    retry if timeout < max_timeout

    # May be best to switch Tor channel if SOCKSError is rescued
    response_timestamp = Time.now.strftime(timestamp_fmt)
    duration = Time.parse(response_timestamp) - Time.parse(request_timestamp)
    rest_client_resp_hash = {
      request_timestamp: request_timestamp,
      response_timestamp: response_timestamp,
      duration: duration,
      http_uri: http_uri,
      http_method: http_method,
      http_resp_code: e.class,
      http_resp_length: 'N/A',
      http_resp_headers: 'N/A',
      http_resp: "ERROR: #{e.message}"
    }
  rescue RestClient::ExceptionWithResponse => e
    response_timestamp = Time.now.strftime(timestamp_fmt)
    duration = Time.parse(response_timestamp) - Time.parse(request_timestamp)

    if e.respond_to?(:response)
      rest_client_resp_hash = {
        request_timestamp: request_timestamp,
        response_timestamp: response_timestamp,
        duration: duration,
        http_uri: http_uri,
        http_method: http_method,
        http_resp_code: e.response.code,
        http_resp_length: e.response.body.length,
        http_resp_headers: JSON.pretty_generate(e.response.headers),
        http_resp: "#{e.response.body[0..300]}..."
      }
    else
      resp_client_resp_hash = {
        request_timestamp: request_timestamp,
        response_timestamp: response_timestamp,
        duration: duration,
        http_uri: http_uri,
        http_method: http_method,
        http_resp_code: 'N/A',
        http_resp_length: 'N/A',
        http_resp_headers: 'N/A',
        http_resp: 'N/A'
      }
    end
  rescue URI::InvalidURIError
    url_encoded_wordlist_arr = []
    wordlist_line.split('/').each do |path|
      url_encoded_wordlist_arr.push(CGI.escape(path))
    end
    wordlist_line = url_encoded_wordlist_arr.join('/')
  rescue RestClient::TooManyRequests
    request_count += 1
    sleep 60

    retry if request_count < max_retries
  rescue SystemExit, Interrupt
    puts "\n#{File.basename($PROGRAM_NAME)}.#{__method__} Goodbye."
  end

  rest_client_resp_hash
ensure
  browser_obj = PWN::Plugins::TransparentBrowser.close(
    browser_obj: browser_obj
  )
end

begin
  pwn_provider = 'ruby-gem'
  # pwn_provider = ENV.fetch('PWN_PROVIDER') if ENV.keys.select { |s| s == 'PWN_PROVIDER' }.any?
  pwn_provider = ENV.fetch('PWN_PROVIDER') if ENV.keys.any? { |s| s == 'PWN_PROVIDER' }
  $stdout.sync = true

  target_url = opts[:target_url]
  parsed_target_url = URI.parse(target_url)

  wordlist = opts[:wordlist]
  raise "ERROR: #{wordlist} Does Not Exist." unless File.exist?(wordlist)

  append_pattern = opts[:append_pattern]
  max_timeout = opts[:max_timeout] ||= 6
  if max_timeout
    max_timeout = max_timeout.to_f
    raise 'ERROR: --max-timeout must be a positive float.' unless max_timeout.positive?
  end

  max_retries = opts[:max_retries] ||= 3
  if max_retries
    max_retries = max_retries.to_i
    raise 'ERROR: --max-retries must be a positive integer.' unless max_retries.positive?
  end

  proxy = opts[:proxy]
  max_threads = opts[:max_threads] ||= 25
  if max_threads
    max_threads = max_threads.to_i
    raise 'ERROR: --max-threads must be a positive integer.' unless max_threads.positive?
  end

  http_request_headers = opts[:http_request_headers]
  include_http_response_codes = opts[:include_http_response_codes]
  include_http_response_codes = include_http_response_codes.to_s.delete("\s").split(',') if include_http_response_codes

  exclude_http_response_codes = opts[:exclude_http_response_codes]
  exclude_http_response_codes = exclude_http_response_codes.to_s.delete("\s").split(',') if exclude_http_response_codes

  raise 'ERROR: Flags --include-response-codes and --exclude-response-codes cannot be used together.' if include_http_response_codes && exclude_http_response_codes

  dir_path = opts[:dir_path]
  dir_path ||= '.'

  report_name = opts[:report_name]
  report_name ||= "#{parsed_target_url.host}-#{File.basename(wordlist)}-#{Time.now.strftime('%Y-%m-%d_%H-%M-%S')}"

  start_reporting_server = opts[:start_reporting_server]

  mutex = Mutex.new

  results_hash = {
    report_name: HTMLEntities.new.encode(
      report_name.to_s.scrub.strip.chomp
    ),
    data: []
  }

  wordlist_arr = File.readlines(wordlist)
  PWN::Plugins::ThreadPool.fill(
    enumerable_array: wordlist_arr,
    max_threads: max_threads
  ) do |this_wl_line|
    wordlist_line = this_wl_line.to_s.scrub.strip.chomp

    next if wordlist_line.match?(/^#/)

    http_methods = %i[
      DELETE
      GET
      HEAD
      OPTIONS
      PATCH
      POST
      PUT
      TRACE
    ].shuffle

    timeout = nil
    http_methods.each do |http_method|
      # TODO: Implement HTTP response timeout
      prev_rest_client_resp_hash = results_hash[:data].last
      rest_client_resp_hash = request_path(
        target_url: target_url,
        proxy: proxy,
        http_request_headers: http_request_headers,
        wordlist_line: "#{wordlist_line}#{append_pattern}",
        http_method: http_method,
        max_timeout: max_timeout,
        timeout: timeout,
        max_retries: max_retries
      )

      mutex.synchronize do
        ret_http_resp_code = rest_client_resp_hash[:http_resp_code].to_s
        # Better waY to implement this?
        if include_http_response_codes.is_a?(Array)
          # If include_http_response_codes is an array, only include
          # the response into the results_hash if the response code
          # is in the include_http_response_codes array
          results_hash[:data].push(rest_client_resp_hash) if include_http_response_codes.include?(ret_http_resp_code) &&
                                                             prev_rest_client_resp_hash != rest_client_resp_hash &&
                                                             rest_client_resp_hash.any?
        else
          # If exclude_http_response_codes is an array, only include
          # the response into the results_hash if the response code
          # is not in the exclude_http_response_codes array
          results_hash[:data].push(rest_client_resp_hash) unless exclude_http_response_codes.is_a?(Array) &&
                                                                 exclude_http_response_codes.include?(ret_http_resp_code) &&
                                                                 prev_rest_client_resp_hash != rest_client_resp_hash &&
                                                                 rest_client_resp_hash.any?
        end
        # Finally puts the last line of the results_hash to STDOUT
        puts results_hash[:data].last if results_hash[:data].last.any?
        timeout = results_hash[:data].last[:duration].to_f if results_hash[:data].last.any?
      end
    end
  end

  # Generate HTML Report
  print "#{File.basename($PROGRAM_NAME)} Generating Report..."
  PWN::Reports::URIBuster.generate(
    dir_path: dir_path,
    results_hash: results_hash
  )
  puts 'complete.'

  # Start Simple HTTP Server (If Requested)
  if start_reporting_server
    listen_port = PWN::Plugins::Sock.get_random_unused_port.to_s

    if pwn_provider == 'docker'
      listen_ip = '0.0.0.0'
    else
      listen_ip = '127.0.0.1'
    end

    puts "For Scan Results Navigate to: http://127.0.0.1:#{listen_port}/#{report_name}.html"
    Dir.chdir(dir_path)
    system(
      'pwn_simple_http_server',
      '-i',
      listen_ip,
      '-p',
      listen_port
    )
  end
# rescue SystemExit, Interrupt
#   puts "\nGoodbye."
rescue Interrupt
  puts "\n#{File.basename($PROGRAM_NAME)} Goodbye."
rescue StandardError => e
  puts e.backtrace
  raise e
end