Module: PWN::Plugins::REPL

Defined in:
lib/pwn/plugins/repl.rb

Overview

This module contains methods related to the pwn REPL Driver.

Class Method Summary collapse

Class Method Details

.add_commands ⇒ Object 



66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
 
# File 'lib/pwn/plugins/repl.rb', line 66

def self.add_commands
  # 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
rescue StandardError => e
  raise e
end

.add_hooks(opts = {}) ⇒ Object

Supported Method Parameters

PWN::Plugins::REPL.add_hooks(

opts: 'required - Hash object passed in via pwn OptParser'

)



154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
 
# File 'lib/pwn/plugins/repl.rb', line 154

public_class_method def self.add_hooks(opts = {})
  # 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]
      is_encrypted = PWN::Plugins::Vault.file_encrypted?(file: yaml_config_path)

      if is_encrypted
        # TODO: Implement "something you know, something you have, && something you are?"
        decryption_file = opts[:decryption_file] ||= "#{Dir.home}/pwn.decryptor.yaml"
        yaml_decryptor = YAML.load_file(decryption_file, symbolize_names: true) if File.exist?(decryption_file)

        key = opts[:key] ||= yaml_decryptor[:key] ||= ENV.fetch('PWN_DECRYPTOR_KEY')
        key = PWN::Plugins::AuthenticationHelper.mask_password(prompt: 'Decryption Key') if key.nil?

        iv = opts[:iv] ||= yaml_decryptor[:iv] ||= ENV.fetch('PWN_DECRYPTOR_IV')
        iv = PWN::Plugins::AuthenticationHelper.mask_password(prompt: 'Decryption IV') if iv.nil?

        decrypted_yaml_config = PWN::Plugins::Vault.dump(
          file: yaml_config_path,
          key: key,
          iv: iv
        )
        yaml_config = YAML.load(decrypted_yaml_config, symbolize_names: true)
      else
        yaml_config = YAML.load_file(yaml_config_path, symbolize_names: true)
      end

      valid_ai_engines = %i[
        openai
        ollama
      ]
      ai_engine = yaml_config[:ai_engine].to_s.to_sym

      raise "ERROR: Unsupported AI Engine: #{ai_engine} in #{yaml_config_path}" unless valid_ai_engines.include?(ai_engine)

      pi.config.pwn_ai_engine = ai_engine
      Pry.config.pwn_ai_engine = ai_engine

      pi.config.pwn_ai_fqdn = yaml_config[ai_engine][:fqdn]
      Pry.config.pwn_ai_fqdn = yaml_config[ai_engine][:fqdn]

      pi.config.pwn_ai_key = yaml_config[ai_engine][:key]
      Pry.config.pwn_ai_key = yaml_config[ai_engine][:key]

      pi.config.pwn_ai_model = yaml_config[ai_engine][:model]
      Pry.config.pwn_ai_model = yaml_config[ai_engine][:model]

      true
    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_engine = pi.config.pwn_ai_engine.to_s.to_sym
      ai_key = pi.config.pwn_ai_key
      ai_key ||= ''
      if ai_key.empty?
        ai_key = PWN::Plugins::AuthenticationHelper.mask_password(
          prompt: 'pwn-ai Key'
        )
        pi.config.pwn_ai_key = ai_key
      end

      response_history = pi.config.pwn_ai_response_history
      speak_answer = pi.config.pwn_ai_speak
      model = pi.config.pwn_ai_model

      case ai_engine
      when :ollama
        fqdn = pi.config.pwn_ai_fqdn

        response = PWN::Plugins::Ollama.chat(
          fqdn: fqdn,
          token: ai_key,
          model: model,
          request: request.chomp,
          response_history: response_history,
          speak_answer: speak_answer
        )
      when :openai
        response = PWN::Plugins::OpenAI.chat(
          token: ai_key,
          request: request.chomp,
          response_history: response_history,
          speak_answer: speak_answer
        )
      else
        raise "ERROR: Unsupported AI Engine: #{ai_engine}"
      end

      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
rescue StandardError => e
  raise e
end

.authors ⇒ Object

Author(s)

0day Inc. <[email protected]>



331
332
333
334
335
 
# File 'lib/pwn/plugins/repl.rb', line 331

public_class_method def self.authors
  "AUTHOR(S):
    0day Inc. <[email protected]>
  "
end

.help ⇒ Object

Display Usage for this Module



339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
 
# File 'lib/pwn/plugins/repl.rb', line 339

public_class_method def self.help
  puts "USAGE:
    #{self}.refresh_ps1_proc(
      mode: 'required - :splat or nil'
    )

    #{self}.add_commands

    #{self}.add_hooks(
      opts: 'required - Hash object passed in via pwn OptParser'
    )

    #{self}.start(
      opts: 'required - Hash object passed in via pwn OptParser'
    )

    #{self}.authors
  "
end

.refresh_ps1_proc(opts = {}) ⇒ Object

Supported Method Parameters

PWN::Plugins::REPL.refresh_ps1_proc(

mode: 'required - :splat or nil'

)



16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
 
# File 'lib/pwn/plugins/repl.rb', line 16

public_class_method def self.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
      ai_engine = pi.config.pwn_ai_engine
      model = pi.config.pwn_ai_model
      pname = "pwn.ai:#{ai_engine}"
      pname = "pwn.ai:#{ai_engine}/#{model}" if model
      pname = "pwn.ai:#{ai_engine}/#{model}.SPEAK" if pi.config.pwn_ai_speak
      pi.config.prompt_name = pname

      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
rescue StandardError => e
  raise e
end

.start(opts = {}) ⇒ Object

Supported Method Parameters

PWN::Plugins::REPL.start(

opts: 'required - Hash object passed in via pwn OptParser'

)



309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
 
# File 'lib/pwn/plugins/repl.rb', line 309

public_class_method def self.start(opts = {})
  # Monkey Patch Pry, add commands, && hooks
  PWN::Plugins::MonkeyPatch.pry
  add_commands
  add_hooks(opts)

  # 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)
  ps1 = [arrow_ps1_proc, splat_ps1_proc]
  prompt = Pry::Prompt.new(:pwn, 'PWN Prototyping REPL', ps1)

  # Start PWN REPL
  Pry.start(self, prompt: prompt)
rescue StandardError => e
  raise e
end