start work on app and tools
This commit is contained in:
parent
051b1157bb
commit
427e7a1ffd
63
final/app.py
63
final/app.py
@ -0,0 +1,63 @@
|
|||||||
|
import os
|
||||||
|
from langchain_openai import ChatOpenAI
|
||||||
|
from langchain.agents import AgentExecutor, create_react_agent
|
||||||
|
from langchain.tools import tool
|
||||||
|
from tools import (
|
||||||
|
get_wireless_interface,
|
||||||
|
change_adapter_mode,
|
||||||
|
wifi_encryption_cracking,
|
||||||
|
deauth_and_capture,
|
||||||
|
reconnaissance,
|
||||||
|
)
|
||||||
|
from langchain import hub
|
||||||
|
from langchain_community.tools import ShellTool
|
||||||
|
from langsmith import Client
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
|
||||||
|
"""
|
||||||
|
Main rag tool-chain application. Simply load the tools and prompt the user.
|
||||||
|
"""
|
||||||
|
|
||||||
|
os.environ["LANGCHAIN_TRACING_V2"] = "true"
|
||||||
|
os.environ["LANGCHAIN_PROJECT"] = f"LangSmith Introduction"
|
||||||
|
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
|
||||||
|
|
||||||
|
client = Client()
|
||||||
|
shell_tool = ShellTool()
|
||||||
|
llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
|
||||||
|
tools = []
|
||||||
|
tools.extend(
|
||||||
|
[
|
||||||
|
|
||||||
|
]
|
||||||
|
)
|
||||||
|
base_prompt = hub.pull("langchain-ai/react-agent-template")
|
||||||
|
prompt = base_prompt.partial(
|
||||||
|
instructions="""
|
||||||
|
You are a packet analysis assistant. Use any combination of the tools provided to best serve the users request.
|
||||||
|
If the request cannot be served with the tools provided, state why and offer advice on how the user could solve the problem.
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
agent = create_react_agent(llm, tools, prompt)
|
||||||
|
agent_executor = AgentExecutor(
|
||||||
|
agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
print(
|
||||||
|
"I am a packet analysis an assistant. I can perform various tasks related to packet capture files."
|
||||||
|
)
|
||||||
|
print(f"I am configured with the following tools")
|
||||||
|
|
||||||
|
for tool in tools:
|
||||||
|
print(f" Tool: {tool.name} = {tool.description}")
|
||||||
|
|
||||||
|
while True:
|
||||||
|
line = input("llm>> ")
|
||||||
|
if line:
|
||||||
|
result = agent_executor.invoke({"input": line})
|
||||||
|
print(result["output"])
|
||||||
|
else:
|
||||||
|
break
|
72
final/tools.py
Normal file
72
final/tools.py
Normal file
@ -0,0 +1,72 @@
|
|||||||
|
import json
|
||||||
|
import threading
|
||||||
|
import random
|
||||||
|
from langchain_core.pydantic_v1 import BaseModel, Field
|
||||||
|
from langchain.tools import tool
|
||||||
|
from langchain_community.tools import ShellTool
|
||||||
|
from langchain.chains import LLMChain
|
||||||
|
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
|
||||||
|
from scapy import *;
|
||||||
|
"""
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
shell_tool = ShellTool()
|
||||||
|
|
||||||
|
|
||||||
|
"""
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
class TCPDump(BaseModel):
|
||||||
|
params: str = Field(
|
||||||
|
description="""A string of all parameters to the `tcpdump` command, including arguments and flags"""
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@tool(
|
||||||
|
"Perform packet analysis using tcpdump",
|
||||||
|
args_schema=TCPDump,
|
||||||
|
return_direct=True,
|
||||||
|
)
|
||||||
|
def tcp_dump(params: str) -> str:
|
||||||
|
"""Must pass all parameters to `tcpdump` including arguments and flags to perform packet analysis"""
|
||||||
|
res = shell_tool.run({"commands": [f"tcpdump {params}"]})
|
||||||
|
return res
|
||||||
|
|
||||||
|
|
||||||
|
class TShark(BaseModel):
|
||||||
|
params: str = Field(
|
||||||
|
description="""A string of all parameters to the `tshark` command, including arguments and flags"""
|
||||||
|
)
|
||||||
|
@tool(
|
||||||
|
"Perform packet analysis using tshark",
|
||||||
|
args_schema=TShark,
|
||||||
|
return_direct=True,
|
||||||
|
)
|
||||||
|
def tshark(params: str) -> str:
|
||||||
|
"""Must pass all parameters to `tshark` including arguments and flags to perform packet analysis"""
|
||||||
|
res = shell_tool.run({"commands": [f"tcpdump {params}"]})
|
||||||
|
return res
|
||||||
|
|
||||||
|
|
||||||
|
# TODO: Scapy tool
|
||||||
|
|
||||||
|
# @tool(
|
||||||
|
# "Perform wifi encryption cracking with aircrack-ng",
|
||||||
|
# args_schema=CrackPassword,
|
||||||
|
# return_direct=True,
|
||||||
|
# )
|
||||||
|
# def wifi_encryption_cracking(json_params: str) -> str:
|
||||||
|
# """Must pass bssid and capfile parameters (as a string containing a json object) to aircrack-ng to perform wifi encryption cracking"""
|
||||||
|
# json_params = json_params.replace("\\", "")
|
||||||
|
# json_obj = json.loads(json_params)
|
||||||
|
# bssid = json_obj["bssid"]
|
||||||
|
# cap = json_obj["capfile"]
|
||||||
|
|
||||||
|
# res = shell_tool.run(
|
||||||
|
# {
|
||||||
|
# "commands": [f"aircrack-ng --bssid {bssid} -w wordlist/rockyou.txt {cap}"]
|
||||||
|
# } # TODO: Abstrace out wordlist - allow custom/multiple wordlist files as long as they sit in the directory
|
||||||
|
# )
|
||||||
|
# return res
|
Reference in New Issue
Block a user