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

require 'optparse'
require 'pwn'
require 'pry'
require 'tty-prompt'
require 'yaml'

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

  options.on('-cPATH', '--yaml-config=PATH', '<Optional - OpenAI YAML File>') do |p|
    opts[:yaml_config_path] = p
  end
end.parse!

begin
  def cleanup_pids(opts = {})
    pids_arr = opts[:pids_arr]

    pids_arr.each do |pid_line|
      pid = pid_line[2].to_i
      Process.kill('TERM', pid)
    rescue Errno::ESRCH
      next
    end
  end

  def refresh_ps1_proc(opts = {})
    mode = opts[:mode]

    proc do |_target_self, _nest_level, pi|
      pi.config.pwn_repl_line += 1
      line_pad = format(
        '%0.3d',
        pi.config.pwn_repl_line
      )

      pi.config.prompt_name = :pwn
      name = "\001\e[1m\002\001\e[31m\002#{pi.config.prompt_name}\001\e[0m\002"
      version = "\001\e[36m\002v#{PWN::VERSION}\001\e[0m\002"
      line_count = "\001\e[34m\002#{line_pad}\001\e[0m\002"
      dchars = "\001\e[32m\002>>>\001\e[0m\002"
      dchars = "\001\e[33m\002***\001\e[0m\002" if mode == :splat

      if pi.config.pwn_asm
        pi.config.prompt_name = 'pwn.asm'
        name = "\001\e[1m\002\001\e[37m\002#{pi.config.prompt_name}\001\e[0m\002"
        dchars = "\001\e[32m\002>>>\001\e[33m\002"
        dchars = "\001\e[33m\002***\001\e[33m\002" if mode == :splat
      end

      if pi.config.pwn_ai
        pi.config.prompt_name = 'pwn.ai'
        pi.config.prompt_name = 'pwn.ai.SPEAKING' if pi.config.pwn_ai_speak
        name = "\001\e[1m\002\001\e[33m\002#{pi.config.prompt_name}\001\e[0m\002"
        dchars = "\001\e[32m\002>>>\001\e[33m\002"
        dchars = "\001\e[33m\002***\001\e[33m\002" if mode == :splat
        if pi.config.pwn_ai_debug
          dchars = "\001\e[32m\002(DEBUG) >>>\001\e[33m\002"
          dchars = "\001\e[33m\002(DEBUG) ***\001\e[33m\002" if mode == :splat
        end
      end

      "#{name}[#{version}]:#{line_count} #{dchars} ".to_s.scrub
    end
  end

  # Pry Monkey Patches \_(oo)_/
  class Pry
    # Overwrite Pry::History.push method in History class to get duplicate history entries
    # in order to properly replay automation in this prototyping driver
    class History
      def push(line)
        return line if line.empty? || invalid_readline_line?(line)

        begin
          last_line = @history[-1]
        rescue IndexError
          last_line = nil
        end

        @history << line
        @history_line_count += 1
        @saver.call(line) if !should_ignore?(line) &&
                             Pry.config.history_save

        line
      end
      alias << push
    end

    def handle_line(line, options)
      if line.nil?
        config.control_d_handler.call(self)
        return
      end

      ensure_correct_encoding!(line)
      Pry.history << line unless options[:generated]

      @suppress_output = false
      inject_sticky_locals!
      begin
        unless process_command_safely(line)
          @eval_string += "#{line.chomp}\n" if !line.empty? || !@eval_string.empty?
        end
      rescue RescuableException => e
        self.last_exception = e
        result = e

        Pry.critical_section do
          show_result(result)
        end
        return
      end

      # This hook is supposed to be executed after each line of ruby code
      # has been read (regardless of whether eval_string is yet a complete expression)
      exec_hook :after_read, eval_string, self

      begin
        complete_expr = true if config.pwn_ai || config.pwn_asm
        complete_expr = Pry::Code.complete_expression?(@eval_string) unless config.pwn_ai || config.pwn_asm
      rescue SyntaxError => e
        output.puts e.message.gsub(/^.*syntax error, */, "SyntaxError: ")
        reset_eval_string
      end

      if complete_expr
        @suppress_output = true if @eval_string =~ /;\Z/ ||
                                   @eval_string.empty? ||
                                   @eval_string =~ /\A *#.*\n\z/ ||
                                   config.pwn_ai ||
                                   config.pwn_asm

        # A bug in jruby makes java.lang.Exception not rescued by
        # `rescue Pry::RescuableException` clause.
        #
        # * https://github.com/pry/pry/issues/854
        # * https://jira.codehaus.org/browse/JRUBY-7100
        #
        # Until that gets fixed upstream, treat java.lang.Exception
        # as an additional exception to be rescued explicitly.
        #
        # This workaround has a side effect: java exceptions specified
        # in `Pry.config.unrescued_exceptions` are ignored.
        jruby_exceptions = []
        jruby_exceptions << Java::JavaLang::Exception if Helpers::Platform.jruby?

        begin
          # Reset eval string, in case we're evaluating Ruby that does something
          # like open a nested REPL on this instance.
          eval_string = @eval_string
          reset_eval_string

          result = evaluate_ruby(eval_string) unless config.pwn_ai ||
                                                     config.pwn_asm

          result = eval_string if config.pwn_ai ||
                                  config.pwn_asm
        rescue RescuableException, *jruby_exceptions => e
          # Eliminate following warning:
          # warning: singleton on non-persistent Java type X
          # (http://wiki.jruby.org/Persistence)
          if Helpers::Platform.jruby? && e.class.respond_to?('__persistent__')
            e.class.__persistent__ = true
          end
          self.last_exception = e
          result = e
        end

        Pry.critical_section do
          show_result(result)
        end
      end

      throw(:breakout) if current_binding.nil?
    end

    # Ensure the return value in pwn_ai mode reflects the input
    def evaluate_ruby(code)
      # if config.pwn_ai || config.pwn_asm
      #   result = message = code.to_s
      #   return
      # end
      inject_sticky_locals!
      exec_hook :before_eval, code, self

      result = current_binding.eval(code, Pry.eval_path, Pry.current_line)
      set_last_result(result, code)
    ensure
      update_input_history(code)
      exec_hook :after_eval, result, self
    end
  end

  # Define Custom REPL Commands
  Pry::Commands.create_command 'welcome-banner' do
    description 'Display the random welcome banner, including basic usage.'

    def process
      puts PWN::Banner.welcome
    end
  end

  Pry::Commands.create_command 'toggle-pager' do
    description 'Toggle less on returned objects surpassing the terminal.'

    def process
      pi = pry_instance
      pi.config.pager ? pi.config.pager = false : pi.config.pager = true
    end
  end

  #  class PWNCompleter < Pry::InputCompleter
  #    def call(input)
  #    end
  #  end

  Pry::Commands.create_command 'pwn-asm' do
    description 'Initiate pwn.asm shell.'

    def process
      pi = pry_instance
      pi.config.pwn_asm = true
      pi.custom_completions = proc do
        prompt = TTY::Prompt.new
        [pi.input.line_buffer]
        # prompt.select(pi.input.line_buffer)
      end
    end
  end

  Pry::Commands.create_command 'pwn-ai' do
    description 'Initiate pwn.ai chat interface.'

    def process
      pi = pry_instance
      pi.config.pwn_ai = true
      pi.config.color = false if pi.config.pwn_ai
      pi.config.color = true unless pi.config.pwn_ai
    end
  end

  Pry::Commands.create_command 'toggle-pwn-ai-debug' do
    description 'Display the response_history object while using pwn.ai'

    def process
      pi = pry_instance
      pi.config.pwn_ai_debug ? pi.config.pwn_ai_debug = false : pi.config.pwn_ai_debug = true
    end
  end

  Pry::Commands.create_command 'toggle-pwn-ai-speaks' do
    description 'Use speech capabilities within pwn.ai to speak answers.'

    def process
      pi = pry_instance
      pi.config.pwn_ai_speak ? pi.config.pwn_ai_speak = false : pi.config.pwn_ai_speak = true
    end
  end

  Pry::Commands.create_command 'back' do
    description 'Jump back to pwn REPL when in pwn-asm || pwn-ai.'

    def process
      pi = pry_instance
      pi.config.pwn_asm = false if pi.config.pwn_asm
      pi.config.pwn_ai = false if pi.config.pwn_ai
      pi.config.pwn_ai_debug = false if pi.config.pwn_ai_debug
      pi.config.pwn_ai_speak = false if pi.config.pwn_ai_speak
      pi.config.completer = Pry::InputCompleter
    end
  end

  # Define REPL Hooks
  # Welcome Banner Hook
  Pry.config.hooks.add_hook(:before_session, :welcome) do |output, _binding, _pi|
    output.puts PWN::Banner.welcome
  end

  # pwn.ai Hooks
  Pry.config.hooks.add_hook(:before_session, :init_opts) do |_output, _binding, pi|
    if opts[:yaml_config_path] && File.exist?(opts[:yaml_config_path])
      yaml_config_path = opts[:yaml_config_path]
      yaml_config = YAML.load_file(yaml_config_path, symbolize_names: true)
      pi.config.pwn_ai_key = yaml_config[:ai_key]
      Pry.config.pwn_ai_key = pi.config.pwn_ai_key
    end
  end

  Pry.config.hooks.add_hook(:after_read, :pwn_asm_hook) do |request, pi|
    if pi.config.pwn_asm && !request.chomp.empty?
      request = pi.input.line_buffer

      # Analyze request to determine if it should be processed as opcodes or asm.
      straight_hex = /^[a-fA-F0-9\s]+$/
      hex_esc_strings = /\\x[\da-fA-F]{2}/
      hex_comma_delim_w_dbl_qt = /"(?:[0-9a-fA-F]{2})",?/
      hex_comma_delim_w_sng_qt = /'(?:[0-9a-fA-F]{2})',?/
      hex_byte_array_as_str = /^\[\s*(?:"[0-9a-fA-F]{2}",\s*)*"[0-9a-fA-F]{2}"\s*\]$/

      if request.match?(straight_hex) ||
         request.match?(hex_esc_strings) ||
         request.match?(hex_comma_delim_w_dbl_qt) ||
         request.match?(hex_comma_delim_w_sng_qt) ||
         request.match?(hex_byte_array_as_str)

        response = PWN::Plugins::Assembly.opcodes_to_asm(
          opcodes: request,
          opcodes_always_strings_obj: true
        )
      else
        response = PWN::Plugins::Assembly.asm_to_opcodes(asm: request)
      end
      puts "\001\e[31m\002#{response}\001\e[0m\002"
    end
  end

  Pry.config.hooks.add_hook(:after_read, :pwn_ai_hook) do |request, pi|
    if pi.config.pwn_ai && !request.chomp.empty?
      request = pi.input.line_buffer.to_s
      debug = pi.config.pwn_ai_debug
      ai_key = pi.config.pwn_ai_key
      ai_key ||= ''
      if ai_key.empty?
        ai_key = PWN::Plugins::AuthenticationHelper.mask_password(
          prompt: 'OpenAI API Key'
        )
        pi.config.pwn_ai_key = ai_key
      end

      response_history = pi.config.pwn_ai_response_history
      speak_answer = pi.config.pwn_ai_speak
      response = PWN::Plugins::OpenAI.chat(
        token: ai_key,
        request: request.chomp,
        temp: 1,
        response_history: response_history,
        speak_answer: speak_answer
      )
      last_response = response[:choices].last[:content]
      puts "\n\001\e[32m\002#{last_response}\001\e[0m\002\n\n"

      response_history = {
        id: response[:id],
        object: response[:object],
        model: response[:model],
        usage: response[:usage]
      }
      response_history[:choices] ||= response[:choices]

      if debug
        puts 'DEBUG: response_history => '
        pp response_history
        puts "\nresponse_history[:choices] Length: #{response_history[:choices].length}\n" unless response_history.nil?
      end
      pi.config.pwn_ai_response_history = response_history
    end
  end

  # Define PS1 Prompt
  Pry.config.pwn_repl_line = 0
  Pry.config.prompt_name = :pwn
  arrow_ps1_proc = refresh_ps1_proc
  splat_ps1_proc = refresh_ps1_proc(mode: :splat)
  prompt_ps1 = [arrow_ps1_proc, splat_ps1_proc]
  prompt = Pry::Prompt.new(
    :pwn,
    'PWN Prototyping REPL',
    prompt_ps1
  )

  # Start PWN REPL
  pwn_pid = Process.pid
  Pry.start(
    self,
    prompt: prompt
  )
rescue StandardError => e
  raise e
ensure
  proc_list_arr = PWN::Plugins::PS.list

  kid_pids_arr = proc_list_arr.select { |proc_line| proc_line[3] == pwn_pid.to_s }
  # pp kid_pids_arr

  grandkid_pids_arr = []
  kid_pids_arr.each do |kid_pid|
    gk_arr = proc_list_arr.select { |proc_line| proc_line[3] == kid_pid[2] }
    gk_arr.each { |gk| grandkid_pids_arr.push(gk) }
  end
  # pp grandkid_pids_arr

  cleanup_pids(pids_arr: grandkid_pids_arr)
  cleanup_pids(pids_arr: kid_pids_arr)
end