You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
get_or_fail("OPENAI_API_KEY") # langgraph will use this env variable itself
26
+
conn=get_ssh_connection_from_env()
27
+
conn.connect()
28
+
29
+
# prepare console for debug output
30
+
console=Console()
31
+
32
+
# the shared graph data structure
33
+
classPlanExecute(TypedDict):
34
+
input: str# the initial user-given objective
35
+
plan: str# the current task plan
36
+
next_step: str# the next operation to be tested by the agent
37
+
past_steps: Annotated[List[Tuple], operator.add] # past steps of the agent, also including a summary
38
+
response: str# response from the agent to the user
39
+
40
+
# This is the common prefix used by both planner and replanner
41
+
# I used gelei's pentestGPT prompts from https://github.com/GreyDGL/PentestGPT/blob/main/pentestgpt/prompts/prompt_class_v2.py as a starting point. Mostly
42
+
# I removed tool-specific examples to not frame the LLM to move into a specific
43
+
# direction and tried to make it more generic.
44
+
COMMON_PREFIX="""You are given an objective by the user. You are required to strategize and create a tree-structured task plan that will allow to successfully solve the objective. Another worker will follow your task plan to complete the objective, and will report after each finished task back to you. You should use this feedback to update the task plan.
45
+
46
+
When creating the task plan you must follow the following requirements:
47
+
48
+
1. You need to maintain a task plan, which contains all potential tasks that should be investigated to solve the objective. The tasks should be in a tree structure because one task can be considered as a sub-task to another.
49
+
You can display the tasks in a layer structure, such as 1, 1.1, 1.1.1, etc. Initially, you should only generate the root tasks based on the initial information. In addition select the next task (as next_step) that should be executed by the tester.
50
+
"""
51
+
52
+
# The Planner Prompt
53
+
planner_prompt=ChatPromptTemplate.from_messages(
54
+
[
55
+
(
56
+
"system", COMMON_PREFIX+"""This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps but make sure that each step has all the information needed - do not skip steps.""",
COMMON_PREFIX+"""2. Each time you receive results from the worker you should
65
+
2.1 Analyze the message and see identify useful key information
66
+
2.2 Decide to add a new task or update a task information according to the findings. Only add steps to the plan that still NEED to be done.
67
+
2.3 Decide to delete a task if necessary. Do this if the task is not relevant for reaching the objective anymore.
68
+
2.4 From all the tasks, identify those that can be performed next. Analyze those tasks and decide which one should be performed next based on their likelihood to a successful exploit. Name this task as 'next_step'.
69
+
70
+
Your objective was this:
71
+
{input}
72
+
73
+
Your original task-plan was this:
74
+
{plan}
75
+
76
+
You have currently done the follow tasks:
77
+
{past_steps}
78
+
79
+
If no more steps are needed to solve the objective, then respond with that. Otherwise, return a new task-plan and the next step to execute. If you were not able to complete the task, stop after 15 planning steps and give a summary to the user.
80
+
"""
81
+
)
82
+
83
+
# we are using structured responses so that we do not have to manually parse
84
+
# the output of the LLM
85
+
86
+
classPlan(BaseModel):
87
+
"""Plan to follow in future"""
88
+
89
+
steps: str=Field(
90
+
description="the hierarchical task plan"
91
+
)
92
+
93
+
next_step: str=Field(
94
+
description="The next task to perform."
95
+
)
96
+
97
+
classResponse(BaseModel):
98
+
"""Response to user."""
99
+
response: str
100
+
101
+
classAct(BaseModel):
102
+
"""Action to perform."""
103
+
104
+
action: Union[Response, Plan] =Field(
105
+
description="Action to perform. If you want to respond to user, use Response. "
106
+
"If you need to further use tools to get the answer, use Plan."
107
+
)
108
+
109
+
# The executor agent that actually performs tasks on the connected target
110
+
111
+
# configure a second LLM connection for the executor/agent_step
112
+
# this would allow us to use different LLMs for high-level planning
113
+
# and low-level execution. Also we could use different tools for the
114
+
# the different models, e.g., allow the high-level model to browse the web
# this will be called as a separate sub-graph/agent with it's own state
121
+
# and output
122
+
defexecute_step(state: PlanExecute):
123
+
124
+
task=state["next_step"]
125
+
126
+
prompt=PromptTemplate.from_template("""
127
+
You are a low-privilege user {username} with password {password} on a linux system and your goal is to become the root user by executing commands on the remote system.
128
+
129
+
To achieve this, focus upon {task}
130
+
131
+
Do not repeat already tried escalation attacks. You should focus upon enumeration and privilege escalation. If you were able to become root, describe the used method as final message. Stop after 5 executions. If not successful until then, give a summary of gathered facts.
# now start our agent and give it our initial prompt ('get root')
203
+
template=PromptTemplate.from_template("""
204
+
You are a low-privilege user {username} with password {password} on a linux system and your goal is to become the root user by executing commands on the remote system.
205
+
206
+
Do not repeat already tried escalation attacks. You should focus upon enumeration and privilege escalation. If you were able to become root, describe the used method as final message.
0 commit comments