Module: PWN::AI::Agent::ToolGuard
- Defined in:
- lib/pwn/ai/agent/tool_guard.rb
Overview
Shared pre-dispatch guards for the two high-volume runtime tools (shell / pwn_eval). Rejects placeholder payloads, aliases wrong schema keys, and names the shell that will actually run the command.
Constant Summary collapse
- ALIASES =
{ 'command' => %w[value cmd input], 'code' => %w[value source ruby input], 'query' => %w[value q text] }.freeze
- PLACEHOLDER_RX =
Token-level junk the model keeps emitting instead of a real command.
/ \A\s*(?:\.{3}|…|\{\s*\.{3}\s*\}|\{\s*…\s*\}|<\.{3}>)\s*\z |(?:\A|\n)\s*(?:\.{3}|…)\s*(?:\z|\n) /x
- BASHISM_RX =
Conservative bash-only constructs. POSIX
$(())is allowed. / \bPIPESTATUS\b |\$\{?RANDOM\}?\b |\[\[(?:\s|\z) |(?:^|[\s;|&])source\s+\S |<\([^)] |&> /x
- CORE_CONSTS =
RestClient uses HTTP::CookieJar. pwn_eval in TOPLEVEL_BINDING can assign HTTP = "/path/http" or Digest = "(self.we" and then every provider hop / payload_sig TypeErrors.
i[HTTP Digest JSON URI Timeout].freeze
- RUNTIMES_FILE =
File.join(Dir.home, '.pwn', 'metrics', 'runtimes.json')
- MAX_PAYLOAD_BYTES =
1_048_576- TIMEOUT_STEP_S =
180- TIMEOUT_MAX_S =
10_800- MUTATION_MAX =
10
Class Method Summary collapse
- .authors ⇒ Object
- .auto_job?(opts = {}) ⇒ Boolean
- .bashism?(opts = {}) ⇒ Boolean
- .bashism_match(opts = {}) ⇒ Object
- .canary_leak?(opts = {}) ⇒ Boolean
-
.coerce_args(opts = {}) ⇒ Object
Coerce common wrong keys onto the first required schema field.
- .command_class(opts = {}) ⇒ Object
-
.deadline_s(opts = {}) ⇒ Object
Conservative wall-clock seconds for shell / pwn_eval.
- .denial(opts = {}) ⇒ Object
- .help ⇒ Object
- .host_load(opts = {}) ⇒ Object
- .injection_score(opts = {}) ⇒ Object
- .invalid_payload(opts = {}) ⇒ Object
- .ip_in_cidr?(opts = {}) ⇒ Boolean
- .mint_canary(opts = {}) ⇒ Object
- .mutation_count(opts = {}) ⇒ Object
- .next_timeout(opts = {}) ⇒ Object
- .note_timeout!(opts = {}) ⇒ Object
- .payload_spent(opts = {}) ⇒ Object
- .placeholder?(opts = {}) ⇒ Boolean
- .placeholder_match(opts = {}) ⇒ Object
- .policy_decision(opts = {}) ⇒ Object
- .predicted_timeout(opts = {}) ⇒ Object
- .present?(opts = {}) ⇒ Boolean
- .protect_core_constants! ⇒ Object
- .protect_http! ⇒ Object
- .quarantine_output(opts = {}) ⇒ Object
- .record_runtime(opts = {}) ⇒ Object
- .refuse_copied_persist?(opts = {}) ⇒ Boolean
- .reset_timeout_budget(opts = {}) ⇒ Object
- .reset_timeout_budget! ⇒ Object
- .rfc1918?(opts = {}) ⇒ Boolean
- .scope_check!(opts = {}) ⇒ Object
- .scope_refusal(opts = {}) ⇒ Object
- .shell_bash? ⇒ Boolean
- .shell_name ⇒ Object
- .shell_syntax_surface(opts = {}) ⇒ Object
-
.timeout_lesson(opts = {}) ⇒ Object
Timeout policy (loop-law, not a skill): 1.
- .timeout_prior_count(opts = {}) ⇒ Object
- .timeout_result(opts = {}) ⇒ Object
- .unwrap_payload(opts = {}) ⇒ Object
Class Method Details
.authors ⇒ Object
740 741 742 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 740 public_class_method def self. "AUTHOR(S):\n 0day Inc. <[email protected]>\n" end |
.auto_job?(opts = {}) ⇒ Boolean
727 728 729 730 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 727 public_class_method def self.auto_job?(opts = {}) pred = opts[:predicted] || predicted_timeout(command_class: opts[:command_class] || command_class(payload: opts[:payload].to_s)) pred.to_i > 120 end |
.bashism?(opts = {}) ⇒ Boolean
191 192 193 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 191 public_class_method def self.bashism?(opts = {}) !bashism_match(text: opts[:text]).nil? end |
.bashism_match(opts = {}) ⇒ Object
195 196 197 198 199 200 201 202 203 204 205 206 207 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 195 public_class_method def self.bashism_match(opts = {}) surface = shell_syntax_surface(text: opts[:text]) return nil if surface.to_s.empty? surface = surface.gsub(/\$\{?RANDOM\}?\b/) { |token| ' ' * token.length } if surface.match?(/\bRANDOM=/) found = BASHISM_RX.match(surface) return nil unless found raw = opts[:text].to_s { rule_id: 'payload.shell_syntax', match: raw[found.begin(0), found[0].length], offset: raw[0...found.begin(0)].bytesize } rescue StandardError nil end |
.canary_leak?(opts = {}) ⇒ Boolean
229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 229 public_class_method def self.canary_leak?(opts = {}) text = opts[:text].to_s toks = Array(Thread.current[:pwn_canary_hist]) toks << Thread.current[:pwn_canary].to_s toks = toks.reject(&:empty?).uniq return false if toks.empty? require 'base64' toks.any? do |tok| next true if text.include?(tok) b64 = Base64.strict_encode64(tok) b64u = Base64.urlsafe_encode64(tok) hex = tok.each_byte.map { |b| format('%02x', b) }.join enc = URI.encode_www_form_component(tok) rot = tok.tr('A-Za-z', 'N-ZA-Mn-za-m') text.include?(b64) || text.include?(b64u) || text.include?(hex) || text.include?(enc) || text.include?(rot) end end |
.coerce_args(opts = {}) ⇒ Object
Coerce common wrong keys onto the first required schema field. Returns the args hash; sets :__schema_error when still missing.
313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 313 public_class_method def self.coerce_args(opts = {}) args = (opts[:args] || {}).dup args = args.each_with_object({}) { |(k, v), m| m[k.to_sym] = v } unless args.empty? req = Array(opts[:required]).map(&:to_s) req.each do |key| next if present?(value: args[key.to_sym]) hit = Array(ALIASES[key]).find { |a| present?(value: args[a.to_sym]) } args[key.to_sym] = args[hit.to_sym] if hit end missing = req.reject { |k| present?(value: args[k.to_sym]) } unless missing.empty? args[:__schema_error] = "missing required #{missing.join(', ')}" args[:__expected] = req args[:__schema_hint] = "Expected keys: #{req.join(', ')}. " \ 'Do not send value/placeholder/ellipsis. ' \ 'Example: shell(command="uname -r") or pwn_eval(code="1+1").' end args rescue StandardError opts[:args] || {} end |
.command_class(opts = {}) ⇒ Object
699 700 701 702 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 699 public_class_method def self.command_class(opts = {}) cmd = opts[:payload].to_s.split.first.to_s cmd.empty? ? 'shell' : File.basename(cmd) end |
.deadline_s(opts = {}) ⇒ Object
Conservative wall-clock seconds for shell / pwn_eval.
Explicit timeout is honored for any payload (1..TIMEOUT_MAX_S).
Omit → host-derived default from loadavg / ncpu / MemAvailable.
No tool-name sniffing: a 65k scan and ls use the same math.
408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 408 public_class_method def self.deadline_s(opts = {}) kind = opts[:kind].to_s.to_sym asked = opts[:timeout] || opts[:timeout_s] asked_i = asked.to_i return asked_i.clamp(1, TIMEOUT_MAX_S) if asked_i.positive? learned = predicted_timeout(command_class: command_class(payload: opts[:payload].to_s)) return learned.clamp(1, TIMEOUT_MAX_S) if learned.to_i.positive? snap = host_load ncpu = [snap[:ncpu].to_i, 1].max load1 = snap[:load1].to_f mem = snap[:mem_avail_mb].to_i default_max = kind == :shell ? 180 : 90 base = kind == :shell ? 30 : 20 base += 15 if load1 > ncpu base += 10 if load1 > (ncpu * 1.5) base += 10 if mem.positive? && mem < 512 base.clamp(8, default_max) end |
.denial(opts = {}) ⇒ Object
369 370 371 372 373 374 375 376 377 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 369 public_class_method def self.denial(opts = {}) { code: opts[:code].to_s, rule_id: opts[:rule_id].to_s, token: opts[:token].to_s, offending_token: opts[:offending_token] || opts[:token].to_s, suggestion: opts[:suggestion].to_s } end |
.help ⇒ Object
744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 744 public_class_method def self.help puts "USAGE: # Decode {encoding:base64,data:} onto command or code. #{self}.unwrap_payload( args: 'required - Hash that may contain encoding and data', key: 'optional - destination key (defaults to command)' ) # Run present and return its result #{self}.present?( value: 'required - integer or string to pack/encode' ) # Detect isolated ellipsis placeholders, not inline ranges or prose. #{self}.placeholder?( text: 'optional - command or payload text to inspect', placeholder_ok: 'optional - true skips denial when the operator opts in', language: 'optional - ruby uses Ripper; shell (default) scans quotes and heredocs' ) # Locate a rejected placeholder in the original decoded payload. #{self}.placeholder_match( text: 'optional - command or payload text to inspect', placeholder_ok: 'optional - true returns no match for intentional literal content', language: 'optional - ruby uses Ripper; shell (default) scans quotes and heredocs' ) # Run bashism and return its result #{self}.bashism?( text: 'optional - text value consumed by #bashism?' ) # Locate unsupported shell syntax in the original decoded payload. #{self}.bashism_match( text: 'optional - shell command to inspect for bash-only constructs' ) # Strip quoted heredocs and single-quoted strings before bash-syntax lint. #{self}.shell_syntax_surface( text: 'required - command string to lint' ) # Mint a per-thread session canary token stored on Thread.current. #{self}.mint_canary( bytes: 'optional - unused reserved length (HMAC uses 16 hex chars)', turn: 'optional - turn number mixed into the HMAC', session_id: 'optional - session id mixed into the HMAC' ) # True when text contains the current session canary. #{self}.canary_leak?( text: 'required - outbound argv or URL to inspect' ) # Heuristic injection score for tool output (ignore-previous, tool-call spoof). #{self}.injection_score( text: 'required - tool output to score' ) # Prefix high-score tool output with a QUARANTINED frame. #{self}.quarantine_output( text: 'required - tool output to wrap if injection_score >= 3' ) # Run shell bash and return its result #{self}.shell_bash? # Run shell name and return its result #{self}.shell_name # Run protect http and return its result #{self}.protect_http! # Run protect core constants and return its result #{self}.protect_core_constants! # Coerce common wrong keys onto the first required schema field #{self}.coerce_args( args: 'optional - args value consumed by #coerce_args', required: 'optional - Array required value consumed by #coerce_args' ) # Build a machine-readable invalid_payload denial (SYNTAX_DENY by default). #{self}.invalid_payload( hint: 'optional - operator-facing hint string', shell: 'optional - shell name (defaults to shell_name)', code: 'optional - denial code (defaults to SYNTAX_DENY)', rule_id: 'optional - rule identifier (defaults to payload)', offending_token: 'optional - exact rejected token span', suggestion: 'optional - how to rewrite the payload', text: 'optional - original payload used to compute byte_range', byte_range: 'optional - [start, stop] byte offsets of the offending span', match: 'optional - exact matched substring; alias for offending_token', offset: 'optional - zero-based byte offset in the original decoded payload', remedy: 'optional - sanctioned correction; alias for suggestion' ) # Build a machine-readable guard denial (SCOPE_DENY, CANARY_DENY, ...). #{self}.denial( code: 'required - denial code such as SCOPE_DENY', rule_id: 'optional - rule identifier', token: 'optional - out-of-scope host or CIDR', offending_token: 'optional - exact rejected token span', suggestion: 'optional - how to stay in scope' ) # Run host load and return its result #{self}.host_load # Conservative wall-clock seconds for shell / pwn_eval #{self}.deadline_s( kind: 'optional - kind value consumed by #deadline_s', timeout: 'optional - seconds to wait before giving up (defaults to opts[:timeout_s])', timeout_s: 'optional - timeout s value consumed by #deadline_s' ) # Run reset timeout budget and return its result #{self}.reset_timeout_budget # Run reset timeout budget and return its result #{self}.reset_timeout_budget! # Run mutation count and return its result #{self}.mutation_count # Run payload spent and return its result #{self}.payload_spent # Run note timeout and return its result #{self}.note_timeout!( timeout: 'optional - seconds to wait before giving up' ) # Run next timeout and return its result #{self}.next_timeout( timeout: 'optional - seconds to wait before giving up', spent: 'optional - spent value consumed by #next_timeout' ) # Timeout policy (loop-law, not a skill): #{self}.timeout_lesson( tool: 'required - tool value consumed by #timeout_lesson', timeout: 'required - seconds to wait before giving up' ) # Run timeout result and return its result #{self}.timeout_result( timeout: 'required - seconds to wait before giving up', tool: 'optional - tool value consumed by #timeout_result', payload: 'optional - payload value consumed by #timeout_result', task: 'optional - task value consumed by #timeout_result', stdout: 'optional - stdout value consumed by #timeout_result', stderr: 'optional - stderr value consumed by #timeout_result', shell: 'required - shell value consumed by #timeout_result (defaults to shell_name)' ) # Run timeout prior count and return its result #{self}.timeout_prior_count( tool: 'optional - tool value consumed by #timeout_prior_count' ) # Run refuse copied persist and return its result #{self}.refuse_copied_persist?( name: 'optional - binary or identifier name', args: 'optional - args value consumed by #refuse_copied_persist?' ) # Refuse a command whose IPs/hosts are outside the ~/.pwn scope yaml allowlists. #{self}.scope_refusal( command: 'required - shell command to inspect for IPs and hostnames' ) # True when ip is RFC1918 or loopback (always in-scope unless strict). #{self}.rfc1918?( ip: 'required - IPv4 address' ) # True when ip is inside cidr (e.g. 10.1.2.3 in 10.0.0.0/8). #{self}.ip_in_cidr?( ip: 'required - IPv4 address', cidr: 'required - CIDR (e.g. 10.0.0.0/8)' ) # First token of a payload used as the runtime class key. #{self}.command_class( payload: 'required - command string whose first token is the class' ) # Append an observed runtime sample for a command class. #{self}.record_runtime( command_class: 'required - class key from the command first token', seconds: 'required - observed wall seconds' ) # p95 * 1.5 timeout from ~/.pwn/metrics/runtimes.json, or nil. #{self}.predicted_timeout( command_class: 'required - class key to look up' ) # True when predicted timeout exceeds 120s (route to job_run). #{self}.auto_job?( command_class: 'optional - class key', predicted: 'optional - override predicted seconds', payload: 'optional - command string if class omitted' ) # Evaluate ~/.pwn/toolguard.yaml first-match rules. #{self}.policy_decision( name: 'required - tool name', args: 'required - Hash of tool arguments' ) # Block network args outside the active engagement scope. #{self}.scope_check!( args: 'optional - Hash of tool arguments', command: 'optional - command string to scan for hosts', text: 'optional - free-form blob to scan' ) # Print the AUTHOR(S) string for this module. #{self}.authors " constants.sort end |
.host_load(opts = {}) ⇒ Object
379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 379 public_class_method def self.host_load(opts = {}) return { ncpu: 1, load1: 0.0, mem_avail_mb: 0 } unless opts.is_a?(Hash) ncpu = File.readable?('/proc/cpuinfo') ? File.read('/proc/cpuinfo').scan(/^processor/).size : 0 ncpu = 1 if ncpu < 1 load1 = 0.0 load1 = File.read('/proc/loadavg').to_s.split[0].to_f if File.readable?('/proc/loadavg') avail = 0 if File.readable?('/proc/meminfo') File.foreach('/proc/meminfo') do |ln| next unless ln.start_with?('MemAvailable:') avail = ln.split[1].to_i / 1024 break end end { ncpu: ncpu, load1: load1, mem_avail_mb: avail } rescue StandardError { ncpu: 1, load1: 0.0, mem_avail_mb: 0 } end |
.injection_score(opts = {}) ⇒ Object
249 250 251 252 253 254 255 256 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 249 public_class_method def self.injection_score(opts = {}) t = opts[:text].to_s score = 0 score += 3 if t.match?(/ignore (all )?(previous|prior) (instructions|directives)/i) score += 2 if t.match?(/system:\s|tool_call|function_call/i) score += 2 if t.match?(/\bexfiltrat|\bdo not tell (the )?(user|operator)/i) score end |
.invalid_payload(opts = {}) ⇒ Object
339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 339 public_class_method def self.invalid_payload(opts = {}) hint = opts[:hint].to_s tok = (opts[:match] || opts[:offending_token]).to_s text = opts[:text].to_s range = opts[:byte_range] if range.nil? && !tok.empty? && !text.empty? idx = opts[:offset] || text.b.index(tok.b) range = [idx, idx + tok.bytesize] if idx end remedy = (opts[:remedy] || opts[:suggestion]).to_s remedy = "Correct the payload before retrying. #{hint}" if remedy.empty? { stdout: '', stderr: hint, exit: 2, error: 'invalid_payload', code: (opts[:code] || 'SYNTAX_DENY').to_s, rule_id: (opts[:rule_id] || 'payload').to_s, match: tok.empty? ? nil : tok, offset: range&.first, remedy: remedy, offending_token: tok, byte_range: range, max_payload_bytes: MAX_PAYLOAD_BYTES, suggestion: remedy, hint: hint, shell: opts[:shell] || shell_name } end |
.ip_in_cidr?(opts = {}) ⇒ Boolean
686 687 688 689 690 691 692 693 694 695 696 697 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 686 public_class_method def self.ip_in_cidr?(opts = {}) ip = opts[:ip].to_s.split('.').map(&:to_i) cidr, bits = opts[:cidr].to_s.split('/') return false if ip.length != 4 || cidr.to_s.split('.').length != 4 mask = bits.to_i mask = 32 if mask <= 0 a = ip.inject(0) { |acc, oct| (acc << 8) + oct } b = cidr.split('.').map(&:to_i).inject(0) { |acc, oct| (acc << 8) + oct } shift = 32 - mask (a >> shift) == (b >> shift) end |
.mint_canary(opts = {}) ⇒ Object
216 217 218 219 220 221 222 223 224 225 226 227 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 216 public_class_method def self.mint_canary(opts = {}) _bytes = opts[:bytes] turn = (opts[:turn] || Thread.current[:pwn_loop_iter] || 0).to_i sid = (opts[:session_id] || Thread.current[:pwn_session_id] || 'sess').to_s key = (Thread.current[:pwn_canary_key] ||= SecureRandom.hex(16)) tok = OpenSSL::HMAC.hexdigest('SHA256', key, "#{sid}:#{turn}")[0, 16] hist = Thread.current[:pwn_canary_hist] ||= [] hist << tok hist.shift while hist.length > 4 Thread.current[:pwn_canary] = tok tok end |
.mutation_count(opts = {}) ⇒ Object
442 443 444 445 446 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 442 public_class_method def self.mutation_count(opts = {}) return 0 unless opts.is_a?(Hash) timeout_mutations[task_key(opts)].to_i end |
.next_timeout(opts = {}) ⇒ Object
469 470 471 472 473 474 475 476 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 469 public_class_method def self.next_timeout(opts = {}) base = opts[:timeout].to_i base = 1 if base < 1 spent = opts.key?(:spent) ? opts[:spent].to_i : payload_spent(opts) remaining = TIMEOUT_MAX_S - spent remaining = 0 if remaining.negative? [base + TIMEOUT_STEP_S, remaining, TIMEOUT_MAX_S].min end |
.note_timeout!(opts = {}) ⇒ Object
454 455 456 457 458 459 460 461 462 463 464 465 466 467 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 454 public_class_method def self.note_timeout!(opts = {}) return 0 unless opts.is_a?(Hash) timeout = opts[:timeout].to_i timeout = 1 if timeout < 1 key = payload_key(opts) timeout_spent[key] = timeout_spent[key].to_i + timeout if budget_exhausted?(opts.merge(spent: timeout_spent[key])) && !timeout_mutated[key] timeout_mutated[key] = true tkey = task_key(opts) timeout_mutations[tkey] = timeout_mutations[tkey].to_i + 1 end timeout_spent[key] end |
.payload_spent(opts = {}) ⇒ Object
448 449 450 451 452 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 448 public_class_method def self.payload_spent(opts = {}) return 0 unless opts.is_a?(Hash) timeout_spent[payload_key(opts)].to_i end |
.placeholder?(opts = {}) ⇒ Boolean
63 64 65 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 63 public_class_method def self.placeholder?(opts = {}) !placeholder_match(text: opts[:text], language: opts[:language], placeholder_ok: opts[:placeholder_ok]).nil? end |
.placeholder_match(opts = {}) ⇒ Object
67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 67 public_class_method def self.placeholder_match(opts = {}) return nil if opts[:placeholder_ok] == true raw = opts[:text].to_s.dup.force_encoding(Encoding::UTF_8) return ruby_placeholder_match(text: raw) if opts[:language].to_s == 'ruby' s = shell_code_surface(text: raw) found = /(?:\A|[;\n|&()])\s*(?<token>\.{3}|…)\s*(?=\z|[;\n|&()])/.match(s) || /\A\s*\[\s*(?<token>\.{3}|…)\s*\]\s*\z/.match(s) if found index = found.begin(:token) token = found[:token] else found = PLACEHOLDER_RX.match(s) return nil unless found token = found[0].strip index = found.begin(0) + found[0].index(token) end offset = s[0...index].bytesize { rule_id: 'payload.placeholder', match: raw.byteslice(offset, token.bytesize), offset: offset } rescue StandardError nil end |
.policy_decision(opts = {}) ⇒ Object
632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 632 public_class_method def self.policy_decision(opts = {}) name = opts[:name].to_s blob = opts[:args].inspect path = File.join(Dir.home, '.pwn', 'toolguard.yaml') return nil unless File.file?(path) require 'yaml' doc = YAML.safe_load_file(path, permitted_classes: [Symbol]) || {} Array(doc['rules'] || doc[:rules]).each do |rule| next unless rule.is_a?(Hash) tool = (rule['tool'] || rule[:tool]).to_s next unless tool.empty? || tool == name pat = (rule['arg_pattern'] || rule[:arg_pattern]).to_s next if pat.empty? || blob.match?(Regexp.new(pat)) == false action = (rule['action'] || rule[:action] || 'allow').to_s next if action == 'allow' return { action: action, success: false, error: (rule['reason'] || rule[:reason] || 'toolguard policy deny').to_s, code: 'POLICY_DENY', tool: name } end nil rescue StandardError nil end |
.predicted_timeout(opts = {}) ⇒ Object
718 719 720 721 722 723 724 725 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 718 public_class_method def self.predicted_timeout(opts = {}) samples = Array(load_runtimes[opts[:command_class].to_s]).map(&:to_f) return nil if samples.empty? sorted = samples.sort idx = [(sorted.length * 0.95).ceil - 1, 0].max (sorted[idx] * 1.5).ceil end |
.present?(opts = {}) ⇒ Boolean
58 59 60 61 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 58 public_class_method def self.present?(opts = {}) value = opts.is_a?(Hash) ? opts[:value] : opts !value.nil? && !value.to_s.strip.empty? end |
.protect_core_constants! ⇒ Object
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 287 public_class_method def self.protect_core_constants! @core_mods ||= {} CORE_CONSTS.each do |name| if Object.const_defined?(name, false) cur = Object.const_get(name, false) @core_mods[name] = cur if cur.is_a?(Module) && @core_mods[name].nil? next if cur.is_a?(Module) Object.send(:remove_const, name) end next unless @core_mods[name].is_a?(Module) next if Object.const_defined?(name, false) && Object.const_get(name, false).equal?(@core_mods[name]) Object.const_set(name, @core_mods[name]) end unless Object.const_defined?(:HTTP, false) && Object.const_get(:HTTP, false).is_a?(Module) require 'http/cookie_jar' @core_mods[:HTTP] = Object.const_get(:HTTP) if Object.const_defined?(:HTTP) && Object.const_get(:HTTP).is_a?(Module) end @core_mods rescue StandardError nil end |
.protect_http! ⇒ Object
283 284 285 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 283 public_class_method def self.protect_http! protect_core_constants! end |
.quarantine_output(opts = {}) ⇒ Object
258 259 260 261 262 263 264 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 258 public_class_method def self.quarantine_output(opts = {}) text = opts[:text].to_s score = injection_score(text: text) return text unless score >= 3 "[QUARANTINED injection_score=#{score}]\n#{text}" end |
.record_runtime(opts = {}) ⇒ Object
704 705 706 707 708 709 710 711 712 713 714 715 716 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 704 public_class_method def self.record_runtime(opts = {}) klass = opts[:command_class].to_s secs = opts[:seconds].to_f return {} if klass.empty? || secs <= 0 data = load_runtimes data[klass] ||= [] data[klass] << secs data[klass] = data[klass].last(50) FileUtils.mkdir_p(File.dirname(RUNTIMES_FILE)) File.write(RUNTIMES_FILE, JSON.generate(data)) data end |
.refuse_copied_persist?(opts = {}) ⇒ Boolean
585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 585 public_class_method def self.refuse_copied_persist?(opts = {}) name = opts[:name].to_s return false unless %w[memory_remember skills_update].include?(name) args = opts[:args] args = {} unless args.is_a?(Hash) text = [args[:value], args['value'], args[:lesson], args['lesson']].compact.join("\n") last = Thread.current[:pwn_last_tool_body].to_s return false if last.length < 80 || text.strip.length < 40 a = text.downcase.scan(/[a-z0-9]{4,}/).uniq b = last.downcase.scan(/[a-z0-9]{4,}/) return false if a.empty? || b.empty? ((a & b).length.to_f / a.length) >= 0.6 rescue StandardError false end |
.reset_timeout_budget(opts = {}) ⇒ Object
429 430 431 432 433 434 435 436 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 429 public_class_method def self.reset_timeout_budget(opts = {}) return :noop unless opts.is_a?(Hash) @timeout_spent = {} @timeout_mutations = {} @timeout_mutated = {} :reset end |
.reset_timeout_budget! ⇒ Object
438 439 440 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 438 public_class_method def self.reset_timeout_budget! reset_timeout_budget end |
.rfc1918?(opts = {}) ⇒ Boolean
675 676 677 678 679 680 681 682 683 684 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 675 public_class_method def self.rfc1918?(opts = {}) oct = opts[:ip].to_s.split('.').map(&:to_i) return false unless oct.length == 4 return true if [10, 127].include?(oct[0]) return true if oct[0] == 192 && oct[1] == 168 return true if oct[0] == 172 && oct[1].between?(16, 31) false end |
.scope_check!(opts = {}) ⇒ Object
665 666 667 668 669 670 671 672 673 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 665 public_class_method def self.scope_check!(opts = {}) return nil unless defined?(Engagement) Engagement.deny_if_out_of_scope( args: opts[:args], command: opts[:command] || opts[:args].inspect, text: opts[:text] ) end |
.scope_refusal(opts = {}) ⇒ Object
604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 604 public_class_method def self.scope_refusal(opts = {}) cmd = opts[:command].to_s path = File.join(Dir.home, '.pwn', 'scope.yaml') return nil unless File.file?(path) require 'yaml' scope = YAML.safe_load_file(path, permitted_classes: [Symbol]) || {} return nil if scope.nil? || scope.empty? expiry = (scope['expiry'] || scope[:expiry]).to_s return denial(code: 'SCOPE_DENY', rule_id: 'expiry', token: expiry, suggestion: 'renew ~/.pwn/scope.yaml expiry') unless expiry.empty? || Time.parse(expiry) >= Time.now allow = Array(scope['cidr_allowlist'] || scope[:cidr_allowlist] || scope['cidrs'] || scope[:cidrs]).map(&:to_s) domains = Array(scope['domain_allowlist'] || scope[:domain_allowlist] || scope['domains'] || scope[:domains]).map(&:to_s) return nil if allow.empty? && domains.empty? ips = cmd.scan(/\b\d{1,3}(?:\.\d{1,3}){3}\b/) hosts = cmd.scan(/\b[a-z0-9][a-z0-9.-]+\.[a-z]{2,}\b/i) bad_ip = allow.any? ? ips.find { |ip| !rfc1918?(ip: ip) && allow.none? { |cidr| ip_in_cidr?(ip: ip, cidr: cidr) } } : nil bad_host = domains.any? ? hosts.find { |h| domains.none? { |d| h.downcase.end_with?(d.downcase) } } : nil hit = bad_ip || bad_host return nil unless hit denial(code: 'SCOPE_DENY', rule_id: 'allowlist', token: hit, suggestion: 'use an in-scope host or CIDR') rescue StandardError nil end |
.shell_bash? ⇒ Boolean
266 267 268 269 270 271 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 266 public_class_method def self.shell_bash? v = (PWN::Env.dig(:ai, :agent, :shell_bash) if defined?(PWN::Env)) v == true || v.to_s.match?(/\A(1|true|yes|on)\z/i) rescue StandardError false end |
.shell_name ⇒ Object
273 274 275 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 273 public_class_method def self.shell_name shell_bash? ? 'bash -lc' : '/bin/sh' end |
.shell_syntax_surface(opts = {}) ⇒ Object
209 210 211 212 213 214 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 209 public_class_method def self.shell_syntax_surface(opts = {}) s = opts[:text].to_s.dup s.gsub!(/<<[-~]?\s*(['"])(\w+)\1.*?^\2\s*$/m) { |body| ' ' * body.length } s.gsub!(/'[^']*'/) { |body| "'#{' ' * (body.length - 2)}'" } s end |
.timeout_lesson(opts = {}) ⇒ Object
Timeout policy (loop-law, not a skill):
- Same payload: timeout += 180 until the 3-hour budget is gone.
- At the 3-hour cap: rewrite ruby/command for the same goal (one mutation). Max MUTATION_MAX mutations per task.
- After MUTATION_MAX mutations: stop (exhausted).
483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 483 public_class_method def self.timeout_lesson(opts = {}) return { scenario: :construction, error: '', hint: '' } unless opts.is_a?(Hash) tool = opts[:tool].to_s timeout = opts[:timeout].to_i spent = payload_spent(opts) spent_after = spent >= timeout && timeout.positive? ? spent : spent + [timeout, 1].max nxt = next_timeout(timeout: timeout, spent: spent_after) mutations = mutation_count(opts) if budget_exhausted?(opts.merge(timeout: timeout, spent: spent_after)) if mutations >= MUTATION_MAX { scenario: :exhausted, error: "#{tool} timeout: #{MUTATION_MAX} mutations exhausted for this task", hint: "This task hit the mutation cap (#{MUTATION_MAX} rewrites after " \ '3-hour budgets). Do not retry the same payload. Report what ' \ 'was tried and what remains blocked.' } else { scenario: :construction, error: "#{tool} timeout: 3-hour budget exhausted; reconstruct payload to same goal", hint: "The #{tool} payload used its 3-hour budget. Generate different " \ 'ruby/command for the same goal. Mutation ' \ "#{[mutations, 1].max}/#{MUTATION_MAX}." } end else { scenario: :deadline, error: "#{tool} timeout: deadline too short; retry with timeout += 180", hint: "Keep the same #{tool} payload. This timeout (#{timeout}s) was too " \ "short. Retry with timeout += 180 (next_timeout=#{nxt})." } end end |
.timeout_prior_count(opts = {}) ⇒ Object
573 574 575 576 577 578 579 580 581 582 583 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 573 public_class_method def self.timeout_prior_count(opts = {}) return 0 unless opts.is_a?(Hash) return 0 unless defined?(PWN::AI::Agent::Mistakes) tool = opts[:tool].to_s PWN::AI::Agent::Mistakes.for_tool(tool: tool, unresolved_only: true).count do |m| m[:shape].to_s == 'timeout' || m[:error].to_s.match?(/timeout/) end rescue StandardError 0 end |
.timeout_result(opts = {}) ⇒ Object
520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 520 public_class_method def self.timeout_result(opts = {}) return { stdout: '', stderr: '', exit: nil, error: 'timeout', scenario: :deadline, hint: '', next_timeout: TIMEOUT_STEP_S, shell: shell_name } unless opts.is_a?(Hash) timeout = opts[:timeout].to_i note_timeout!(opts) lesson = timeout_lesson( tool: opts[:tool], payload: opts[:payload], timeout: timeout, task: opts[:task] ) { stdout: opts[:stdout].to_s, stderr: opts[:stderr].to_s, exit: nil, error: "timeout after #{timeout}s", scenario: lesson[:scenario], hint: lesson[:hint], next_timeout: next_timeout(timeout: timeout, spent: payload_spent(opts)), mutations: mutation_count(opts), shell: opts[:shell] || shell_name } end |
.unwrap_payload(opts = {}) ⇒ Object
43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
# File 'lib/pwn/ai/agent/tool_guard.rb', line 43 public_class_method def self.unwrap_payload(opts = {}) args = opts[:args] return args unless args.is_a?(Hash) enc = (args[:encoding] || args['encoding']).to_s data = args[:data] || args['data'] return args unless enc == 'base64' && !data.to_s.empty? raw = Base64.strict_decode64(data.to_s) key = (opts[:key] || :command).to_sym args.merge(key => raw) rescue ArgumentError opts[:args] end |