73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
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
|