greydgl--pentestgpt
Gelei Deng
3295953495
* chore: 🤖 update to poetry * chore: 🤖 major poetry update and lint fix ✅ Closes: #200
112 行
3.8 KiB
Python
112 行
3.8 KiB
Python
import dataclasses
|
|
import json
|
|
import os
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import boto3
|
|
import loguru
|
|
import tiktoken
|
|
from tenacity import *
|
|
|
|
from pentestgpt.utils.llm_api import LLMAPI
|
|
|
|
logger = loguru.logger
|
|
logger.remove()
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Message:
|
|
ask_id: str = None
|
|
ask: dict = None
|
|
answer: dict = None
|
|
answer_id: str = None
|
|
request_start_timestamp: float = None
|
|
request_end_timestamp: float = None
|
|
time_escaped: float = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Conversation:
|
|
conversation_id: str = None
|
|
message_list: List[Message] = dataclasses.field(default_factory=list)
|
|
|
|
def __hash__(self):
|
|
return hash(self.conversation_id)
|
|
|
|
def __eq__(self, other):
|
|
if not isinstance(other, Conversation):
|
|
return False
|
|
return self.conversation_id == other.conversation_id
|
|
|
|
|
|
class TitanAPI(LLMAPI):
|
|
def __init__(self, config_class, use_langfuse_logging=False):
|
|
self.name = str(config_class.model)
|
|
self.model = config_class.model
|
|
self.log_dir = config_class.log_dir
|
|
self.history_length = 5 # maintain 5 messages in the history. (5 chat memory)
|
|
self.conversation_dict: Dict[str, Conversation] = {}
|
|
self.error_waiting_time = 3 # wait for 3 seconds
|
|
self._bedrock_connection()
|
|
|
|
logger.add(sink=os.path.join(self.log_dir, "titan.log"), level="WARNING")
|
|
|
|
def _bedrock_connection(self):
|
|
self.bedrock = boto3.client(
|
|
"bedrock",
|
|
"us-west-2",
|
|
endpoint_url="https://bedrock.us-west-2.amazonaws.com",
|
|
)
|
|
|
|
def _chat_completion(
|
|
self, history: List, model="amazon.titan-tg1-large", temperature=0.5
|
|
) -> str:
|
|
"""
|
|
:param history: a list of strings
|
|
:return: a string
|
|
"""
|
|
body = json.dumps(history)
|
|
modelId = model
|
|
accept = "application/json"
|
|
contentType = "application/json"
|
|
try:
|
|
print("body: ", body)
|
|
print("modelId: ", modelId)
|
|
response = self.bedrock.invoke_model(
|
|
body=body, modelId=modelId, accept=accept, contentType=contentType
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_string = response_body.get("results")[0].get("outputText")
|
|
return response_string
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error: {e}")
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from module_import import TitanConfigClass
|
|
|
|
bedrock = boto3.client(
|
|
"bedrock", "us-west-2", endpoint_url="https://bedrock.us-west-2.amazonaws.com"
|
|
)
|
|
output_text = bedrock.list_foundation_models()
|
|
config_class = TitanConfigClass()
|
|
config_class.log_dir = "logs"
|
|
titan = TitanAPI(config_class)
|
|
|
|
# test is below
|
|
# 1. create a new conversation
|
|
result, conversation_id = titan.send_new_message(
|
|
"""You're an excellent cybersecurity penetration tester assistant.
|
|
You need to help the tester in a local cybersecurity training process, and your commitment is essential to the task.
|
|
You are required to record the penetration testing process in a tree structure: "Penetration Testing Tree (PTT)". It is structured as follows:
|
|
(1) The tasks are in layered structure, i.e., 1, 1.1, 1.1.1, etc. Each task is one operation in penetration testing; task 1.1 should be a sub-task of task 1.
|
|
(2) Each task has a completion status: to-do, completed, or not applicable.
|
|
(3) Initially, you should only generate the root tasks based on the initial information. In most cases, it should be reconnaissance tasks. You don't generate tasks for unknown ports/services. You can expand the PTT later.
|
|
|
|
You shall not provide any comments/information but the PTT. You will be provided with task info and start the testing soon. Reply Yes if you understand the task."""
|
|
)
|
|
print("Answer 1")
|
|
print(result)
|