-
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathswitch-to-react.py
More file actions
48 lines (36 loc) · 1.67 KB
/
switch-to-react.py
File metadata and controls
48 lines (36 loc) · 1.67 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
from dotenv import load_dotenv
from rich.console import Console
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from helper.common import get_or_fail
from helper.ui import print_event_stream
from tools.ssh import SshExecuteTool, SshTestCredentialsTool,get_ssh_connection_from_env
# setup configuration from environment variables
load_dotenv()
conn = get_ssh_connection_from_env()
get_or_fail("OPENAI_API_KEY") # langgraph will use this env variable itself
# connect to the target system over SSH
conn.connect()
# initialize the ChatOpenAI model and register the tool (ssh connection)
llm = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [SshExecuteTool(conn), SshTestCredentialsTool(conn)]
# configure our agent
agent_executor = create_react_agent(llm, tools)
# now start our agent and give it our initial prompt ('get root')
template = PromptTemplate.from_template("""
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.
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.
""").format(username=conn.username, password=conn.password)
if __name__ == '__main__':
console = Console()
events = agent_executor.stream(
{
"messages": [
("user", template),
]
},
stream_mode="values",
)
# output all the events that we're getting from the agent
print_event_stream(console, events)