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

require 'pwn'
require 'optparse'
require 'uri'
require 'cgi'
require 'htmlentities'

opts = {}
OptionParser.new do |options|
  options.banner = "USAGE:
    #{$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('-pPROXY', '--proxy=PROXY', '<Optional - Proxy SCHEME://ADDRESS:PORT>') do |p|
    opts[:proxy] = p
  end

  options.on('-tTHREADS', '--max-threads=THREADS', '<Optional # HTTP Requests to Run Simultaneously (Default 100)>') 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('-cCODES', '--only-report-response-codes=CODES', '<Optional - Comma-Delimited List of Response Codes to Only Include in Report)>') do |c|
    opts[:http_response_codes] = c
  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 `#{$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]

  rest_client_resp_hash = {}
  begin
    print '.'
    http_uri = "#{target_url}/#{wordlist_line}"
    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 = nil
    if http_request_headers
      headers = JSON.parse(
        http_request_headers,
        symbolize_names: true
      )
    end

    response = rest_client.execute(
      method: http_method,
      url: http_uri,
      headers: headers,
      verify_ssl: false
    )

    rest_client_resp_hash = {
      request_timestamp: Time.now.strftime('%Y-%m-%d_%H-%M-%S'),
      http_uri: http_uri,
      http_method: http_method,
      http_resp_code: response.code,
      http_resp_length: response.body.length,
      http_resp: "#{response.body[0..300]}..."
    }
  rescue Errno::ECONNREFUSED
    raise 'ERROR: Connection(s) Refused. Try lowering the --max-threads value.'
  rescue Errno::ECONNRESET,
         RestClient::Exceptions::ReadTimeout,
         RestClient::ServerBrokeConnection => e
    rest_client_resp_hash = {
      request_timestamp: Time.now.strftime('%Y-%m-%d_%H-%M-%S'),
      http_uri: http_uri,
      http_method: http_method,
      http_resp_code: e.class,
      http_resp_length: 'N/A',
      http_resp: e.class
    }
  rescue RestClient::ExceptionWithResponse => e
    rest_client_resp_hash = {
      request_timestamp: Time.now.strftime('%Y-%m-%d_%H-%M-%S'),
      http_uri: http_uri,
      http_method: http_method,
      http_resp_code: e.response.code,
      http_resp_length: e.response.body.length,
      http_resp: "#{e.response.body[0..300]}..."
    }
  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('/')

    retry
  rescue RestClient::TooManyRequests
    sleep 60
  end

  rest_client_resp_hash
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]
  proxy = opts[:proxy]

  max_threads = opts[:max_threads]
  max_threads ||= 100

  http_request_headers = opts[:http_request_headers]
  http_response_codes = opts[:http_response_codes]
  http_response_codes = http_response_codes.delete("\s").split(',') if 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]
    http_methods.each do |http_method|
      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
      )

      mutex.synchronize do
        if http_response_codes
          ret_http_resp_code = rest_client_resp_hash[:http_resp_code].to_s
          results_hash[:data].push(rest_client_resp_hash) if http_response_codes.include?(ret_http_resp_code)
        else
          results_hash[:data].push(rest_client_resp_hash)
        end
      end
    end
  end

  # Generate HTML Report
  print "#{$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 = Random.rand(1_025..65_535).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}/pwn_www_uri_buster.html"
    Dir.chdir(dir_path)
    system(
      'pwn_simple_http_server',
      '-i',
      listen_ip,
      '-p',
      listen_port
    )
  end
rescue SystemExit, Interrupt
  puts "\nGoodbye."
rescue StandardError => e
  raise e
end