Module: PWN::AI::CLI

Defined in:
lib/pwn/ai/cli.rb

Overview

Explicit standalone entrypoint; help and parsing never load a vault.

Class Method Summary collapse

Class Method Details

.authors ⇒ Object 



69
70
71
 
# File 'lib/pwn/ai/cli.rb', line 69

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

.help ⇒ Object 



73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
 
# File 'lib/pwn/ai/cli.rb', line 73

public_class_method def self.help
  puts parse(argv: ['--help'])[:help_text]
  puts "
    # Parse CLI arguments without loading a vault.
    #{self}.parse(argv: 'optional - argument array; defaults to []')
    # Execute the requested CLI action; replay/rerun never load credentials.
    #{self}.run(
      argv: 'optional - argument array; defaults to ARGV',
      input: 'optional - prompt input IO; defaults to stdin',
      output: 'optional - result IO; defaults to stdout',
      error: 'optional - ingestion diagnostics IO; defaults to stderr'
    )
    # Display author information.
    #{self}.authors
  "
end

.parse(opts = {}) ⇒ Object

Raises:

  • (OptionParser::InvalidArgument) 


10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
 
# File 'lib/pwn/ai/cli.rb', line 10

public_class_method def self.parse(opts = {})
  result = {}
  parser = OptionParser.new do |options|
    options.banner = 'Usage: pwn-ai [--analyze PATH | --replay ID | --rerun ID] [--ai PROMPT]'
    options.on('--analyze PATH', 'Ingest evidence before starting the AI session') { |v| result[:analyze] = v }
    options.on('--replay ID', 'Render saved trace without executing tools') { |v| result[:replay] = v }
    options.on('--rerun ID', 'Re-execute saved tools in a fresh context (side effects possible)') { |v| result[:rerun] = v }
    options.on('--ai PROMPT', 'One-shot request; - reads standard input') { |v| result[:ai] = v }
    options.on('--pwn-env PATH', 'Use the specified encrypted configuration') { |v| result[:pwn_env_path] = v }
    options.on('--pwn-dec PATH', 'Use the specified decryptor') { |v| result[:pwn_dec_path] = v }
    options.on('-h', '--help', 'Show help without loading configuration') { result[:help] = true }
  end
  remaining = parser.parse!(Array(opts[:argv]).dup)
  raise OptionParser::InvalidArgument, "unexpected arguments: #{remaining.join(' ')}" unless remaining.empty?
  raise OptionParser::InvalidArgument, 'choose only one of --analyze, --replay, --rerun' if i[analyze replay rerun].count { |key| result.key?(key) } > 1
  raise OptionParser::InvalidArgument, '--ai cannot accompany --replay or --rerun' if result[:ai] && (result[:replay] || result[:rerun])

  result[:help_text] = parser.to_s
  result
end

.run(opts = {}) ⇒ Object 



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
62
63
64
65
66
67
 
# File 'lib/pwn/ai/cli.rb', line 31

public_class_method def self.run(opts = {})
  parsed = parse(argv: opts[:argv] || ARGV)
  output = opts[:output] || $stdout
  if parsed[:help]
    output.puts(parsed[:help_text])
    return 0
  end

  require 'pwn'
  if parsed[:replay] || parsed[:rerun]
    require 'pwn/session_trace'
    if parsed[:replay]
      PWN::SessionTrace.replay(session_id: parsed[:replay], io: output)
    else
      report = PWN::SessionTrace.rerun(session_id: parsed[:rerun], environment: :bubblewrap)
      output.puts(JSON.generate(report))
      return 1 if report[:results].any? { |result| result['exit_status'] != 0 }
    end
    return 0
  end

  PWN::Config.refresh_env(**parsed.slice(:pwn_env_path, :pwn_dec_path))
  session = PWN::Sessions.create(title: 'pwn-ai', source: 'pwn-ai-cli')
  if parsed[:analyze]
    report = PWN::AI::Context.ingest(path: parsed[:analyze], session_id: session[:id], format: 'binary')
    (opts[:error] || $stderr).puts(JSON.generate(report))
  end
  if parsed[:ai]
    prompt = parsed[:ai] == '-' ? (opts[:input] || $stdin).read : parsed[:ai]
    raise ArgumentError, '--ai requires a non-empty prompt' if prompt.strip.empty?

    output.puts(PWN::AI::Agent::Loop.run(request: prompt, session_id: session[:id], enabled_toolsets: PWN::Env.dig(:ai, :agent, :toolsets)))
  else
    PWN::Plugins::REPL.start(ai_session_id: session[:id])
  end
  0
end