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



169
170
171
# File 'lib/pwn/ai/cli.rb', line 169

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

.help ⇒ Object



173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
# File 'lib/pwn/ai/cli.rb', line 173

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)


12
13
14
15
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
# File 'lib/pwn/ai/cli.rb', line 12

public_class_method def self.parse(opts = {})
  result = {}
  parser = OptionParser.new do |options|
    options.banner = 'Usage: pwn-ai [--analyze PATH | --replay ID | --rerun ID | --plan-only | --execute PATH | --resume RUN_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('--plan-only', 'Emit a YAML task DAG and write the mission ledger without running tools') { result[:plan_only] = true }
    options.on('--execute PATH', 'Run an approved YAML task DAG') { |v| result[:execute] = v }
    options.on('--resume RUN_ID', 'Resume a checkpointed DAG run, skipping completed steps') { |v| result[:resume] = v }
    options.on('--mission ID', 'Bind this plan or run to a durable mission') { |v| result[:mission] = v }
    options.on('--policy ACTION', %w[evaluate promote rollback], 'Offline policy: evaluate, promote, rollback (no AI session)') { |v| result[:policy] = v }
    options.on('--baseline PATH', 'Frozen baseline policy JSON') { |v| result[:baseline] = v }
    options.on('--candidate PATH', 'Frozen candidate policy JSON') { |v| result[:candidate] = v }
    options.on('--reports PATH', 'JSON array from --policy evaluate; replayed before promotion') { |v| result[:reports] = v }
    options.on('--receipt PATH', 'Saved successful promotion JSON for rollback') { |v| result[:receipt] = v }
    options.on('--live-policy PATH', 'Explicit existing policy target; no default') { |v| result[:live_path] = v }
    options.on('--approve-policy-change', 'Operator explicitly approves promotion or rollback') { result[:enabled] = true }
    options.on('--policy-writers-stopped', 'Operator attests ALL policy writers are stopped') { result[:quiescent] = true }
    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, --plan-only, --execute, --resume' if i[analyze replay rerun plan_only execute resume].count { |key| result.key?(key) } > 1
  raise OptionParser::InvalidArgument, '--ai cannot accompany --replay or --rerun' if result[:ai] && (result[:replay] || result[:rerun])
  raise OptionParser::InvalidArgument, '--plan-only requires --ai' if result[:plan_only] && !result[:ai]

  validate_policy_options(parsed: result)

  result[:help_text] = parser.to_s
  result
end

.run(opts = {}) ⇒ Object



48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
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
# File 'lib/pwn/ai/cli.rb', line 48

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'
  return run_policy(parsed: parsed, output: output) if parsed[:policy]

  if parsed[:plan_only]
    prompt = parsed[:ai] == '-' ? (opts[:input] || $stdin).read : parsed[:ai]
    raise ArgumentError, '--ai requires a non-empty prompt' if prompt.to_s.strip.empty?

    dag = PWN::AI::Agent::Mission.plan!(request: prompt, id: parsed[:mission])
    output.puts(YAML.dump(JSON.parse(JSON.generate(dag))))
    return 0
  end
  if parsed[:resume] || parsed[:execute]
    PWN::Config.refresh_env(**parsed.slice(:pwn_env_path, :pwn_dec_path))
    begin
      report = if parsed[:resume]
                 PWN::AI::Agent::TaskDAG.resume(run_id: parsed[:resume], operator: true, mission_id: parsed[:mission])
               else
                 PWN::AI::Agent::TaskDAG.execute(path: parsed[:execute], approved: true, operator: true, mission_id: parsed[:mission])
               end
    rescue ArgumentError => e
      output.puts(e.message)
      return 1
    end
    bind_mission!(mission_id: parsed[:mission], report: report)
    output.puts(JSON.generate(report))
    return report[:ok] == false ? 1 : 0
  end
  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