diff --git a/web/reNgine/definitions.py b/web/reNgine/definitions.py index d4abe300..73b71328 100644 --- a/web/reNgine/definitions.py +++ b/web/reNgine/definitions.py @@ -509,8 +509,8 @@ VULNERABILITY_DESCRIPTION_SYSTEM_MESSAGE = """ ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT = """ You are a highly skilled penetration tester who has recently completed a reconnaissance on a target. - As a penetration tester, I've conducted a thorough reconnaissance on a specific subdomain. - Based on my reconnaissance you will be given with a + As a penetration tester, you've conducted a thorough reconnaissance on a specific subdomain. + Based on the reconnaissance you will be given with a - Subdomain Name - Subdomain Page Title - Open Ports if any detected diff --git a/web/reNgine/gpt.py b/web/reNgine/gpt.py index 389da1fc..0567efe5 100644 --- a/web/reNgine/gpt.py +++ b/web/reNgine/gpt.py @@ -93,27 +93,37 @@ class GPTVulnerabilityReportGenerator: class GPTAttackSuggestionGenerator: def __init__(self): - self.api_key = get_open_ai_key() - self.model_name = 'gpt-3.5-turbo' - if not self.api_key: - self.ollama = Ollama(base_url='http://ollama:11434', model="llama2-uncensored") + selected_model = OllamaSettings.objects.first() + self.model_name = selected_model.selected_model if selected_model else 'gpt-3.5-turbo' + self.use_ollama = selected_model.use_ollama if selected_model else False + self.openai_api_key = None + self.ollama = None - def get_attack_suggestion(self, input): + def get_attack_suggestion(self, user_input): ''' - input (str): input for gpt + user_input (str): input for gpt ''' - if not self.api_key: - prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + input + if self.use_ollama: + prompt = ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT + "\nUser: " + user_input + self.ollama = Ollama( + base_url=OLLAMA_INSTANCE, + model=self.model_name + ) response_content = self.ollama(prompt) else: - openai.api_key = self.api_key - print(input) + openai_api_key = get_open_ai_key() + if not openai_api_key: + return { + 'status': False, + 'error': 'OpenAI API Key not set' + } try: + openai.api_key = openai_api_key gpt_response = openai.ChatCompletion.create( model=self.model_name, messages=[ {'role': 'system', 'content': ATTACK_SUGGESTION_GPT_SYSTEM_PROMPT}, - {'role': 'user', 'content': input} + {'role': 'user', 'content': user_input} ] ) response_content = gpt_response['choices'][0]['message']['content'] @@ -121,11 +131,11 @@ class GPTAttackSuggestionGenerator: return { 'status': False, 'error': str(e), - 'input': input + 'input': user_input } return { 'status': True, 'description': response_content, - 'input': input + 'input': user_input } \ No newline at end of file