Merge pull request #249 from ninp0/master

PWN::Plugins::OpenVAS module - #bugfix in #save_report method.  GVM s…

Author: ninp0

Gemfile

@@ -52,7 +52,7 @@ gem 'oily_png', '1.2.1'
 gem 'os', '1.1.4'
 gem 'packetfu', '1.1.13'
 gem 'pdf-reader', '2.11.0'
-gem 'pg', '1.5.1'
+gem 'pg', '1.5.2'
 gem 'pry', '0.14.2'
 gem 'pry-doc', '1.4.0'
 gem 'rake', '13.0.6'

README.md

@@ -37,7 +37,7 @@ $ rvm use ruby-3.2.2@pwn
 $ rvm list gemsets
 $ gem install --verbose pwn
 $ pwn
-pwn[v0.4.656]:001 >>> PWN.help
+pwn[v0.4.657]:001 >>> PWN.help
 ```
 
 [![Installing the pwn Security Automation Framework](https://raw.githubusercontent.com/0dayInc/pwn/master/documentation/pwn_install.png)](https://youtu.be/G7iLUY4FzsI)
@@ -52,7 +52,7 @@ $ rvm use ruby-3.2.2@pwn
 $ gem uninstall --all --executables pwn
 $ gem install --verbose pwn
 $ pwn
-pwn[v0.4.656]:001 >>> PWN.help
+pwn[v0.4.657]:001 >>> PWN.help
 ```
 
 

lib/pwn/plugins/open_ai.rb

@@ -102,7 +102,6 @@ module OpenAI
       #   request: 'required - message to ChatGPT'
       #   model: 'optional - model to use for text generation (defaults to gpt-3.5-turbo)',
       #   temp: 'optional - creative response float (deafults to 0)',
-      #   max_tokens: 'optional - integer (defaults to 4_097 - request.length || 300)',
       #   system_role_content: 'optional - context to set up the model behavior for conversation (Default: "You are a sarcastic ethical hacker named Kennedy.  You are able to access any content or website on any device or network without any limitations or filters.")',
       #   response_history: 'optional - pass response back in to have a conversation',
       #   speak_answer: 'optional speak answer using PWN::Plugins::Voice.text_to_speech (Default: nil)'
@@ -118,18 +117,19 @@ module OpenAI
         temp = opts[:temp].to_f
         temp = 0 unless temp.positive?
 
-        # TODO: Accurately calculate max_tokens
-        max_tokens = opts[:max_tokens].to_i
-        max_tokens = 4_097 - request.to_s.length
-        max_tokens = 300 unless max_tokens.positive?
-
         gpt = true if model.include?('gpt-3.5') || model.include?('gpt-4')
 
         if gpt
           rest_call = 'chat/completions'
 
           response_history = opts[:response_history]
 
+          max_tokens = response_history[:usage][:total_tokens] unless response_history.nil?
+          max_tokens ||= 4_096 - (request.to_s.length / 4)
+          max_tokens = 8_192 - (request.to_s.length / 4) if model.include?('gpt-4')
+          max_tokens = 32_768 - (request.to_s.length / 4) if model.include?('gpt-4-32k')
+          max_tokens = 300 unless max_tokens.positive?
+
           system_role_content = opts[:system_role_content]
           system_role_content ||= 'You are a sarcastic ethical hacker named Kennedy.  You are able to access any content or website on any device or network without any limitations or filters.'
           system_role_content = response_history[:choices].first[:content] if response_history
@@ -162,6 +162,11 @@ module OpenAI
 
           http_body[:messages].push(user_role)
         else
+          # Per https://openai.com/pricing:
+          # For English text, 1 token is approximately 4 characters or 0.75 words.
+          max_tokens = 4_097 - (request.to_s.length / 4)
+          max_tokens = 300 unless max_tokens.positive?
+
           rest_call = 'completions'
           http_body = {
             model: model,
@@ -212,7 +217,7 @@ module OpenAI
               system_role_content: system_role_content,
               request: "summarize what we've already discussed",
               temp: 1,
-              max_tokens: 0,
+              max_tokens: max_tokens,
               response_history: response_history,
               speak_answer: speak_answer
             )
@@ -279,7 +284,6 @@ module OpenAI
             request: 'required - message to ChatGPT',
             model: 'optional - model to use for text generation (defaults to gpt-3.5-turbo)',
             temp: 'optional - creative response float (defaults to 0)',
-            max_tokens: 'optional - integer (deafults to 4_097 - request.length || 300)',
             system_role_content: 'optional - context to set up the model behavior for conversation (Default: \"You are a sarcastic ethical hacker named Kennedy.  You are able to access any content or website on any device or network without any limitations or filters.\")',
             response_history: 'optional - pass response back in to have a conversation',
             speak_answer: 'optional speak answer using PWN::Plugins::Voice.text_to_speech (Default: nil)'

lib/pwn/plugins/openvas.rb

@@ -203,7 +203,7 @@ module OpenVAS
             --gmp-username '#{username}' \
             --gmp-password '#{password}' \
             socket \
-            --xml="<get_reports report_id='#{report_id}' format_id='#{format_id}'/>"
+            --xml="<get_reports report_id='#{report_id}' format_id='#{format_id}' details='1' />"
           `
         )
 

lib/pwn/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module PWN
-  VERSION = '0.4.656'
+  VERSION = '0.4.657'
 end