This repository has been archived on 2025-04-28. You can view files and clone it, but cannot push or open issues or pull requests.
gensec-westgate-djw2/hw6/tools.py
David Westgate f7eae67744 added tools
2024-05-22 15:59:39 -07:00

74 lines
2.7 KiB
Python

import subprocess
import requests
import tempfile
import zipfile
import io
import os
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain.agents import AgentExecutor, create_react_agent, load_tools
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain.tools import tool
from langchain import hub
from langchain_community.tools import ShellTool
from dotenv import load_dotenv
from time import time
shell_tool = ShellTool()
# class EnableMonitorMode(BaseModel):
# query: str = Field(description="should be a request to enable monitor mode of the wireless adapter")
# def get_wireless_interface():
# result = subprocess.run(['iw', 'dev'], stdout=subprocess.PIPE)
# output = result.stdout.decode('utf-8')
# for line in output.splitlines():
# if 'Interface' in line:
# return line.split()[1]
# @tool("enable monitor mode", args_schema=EnableMonitorMode ,return_direct=False)
# def enable_monitor_mode(interface_name: str) -> str:
# """Can enable monitor mode"""
# if interface_name.find('mon'):
# res = "Montior mode already enabled"
# else:
# res = shell_tool.run({"commands": ["sudo airmon-ng check kill", f"sudo airmon-ng start {interface_name}"]})
# return res
class ChangeMonitorMode(BaseModel):
query: str = Field(description="Should be command line parameters to 'airmon-ng' to change the state of a given wireless iterface mode")
# def get_wireless_interface():
# result = subprocess.run(['iw', 'dev'], stdout=subprocess.PIPE)
# output = result.stdout.decode('utf-8')
# for line in output.splitlines():
# if 'Interface' in line:
# return line.split()[1]
@tool("Change the state of the wireless adapter mode with airmon-ng", args_schema=ChangeMonitorMode ,return_direct=False)
def change_adapter_mode(params: str) -> str:
"""Can pass parameters to airmon-ng to change the mode of the wireless adapter"""
# if interface_name.find('mon'):
# res = "Montior mode already enabled"
# else:
res = shell_tool.run({"commands": [f"sudo airmon-ng {params}"]})
return res
class Iwconfig(BaseModel):
params: str = Field(description="should be command line parameters to 'iwconfig', if needed")
@tool("Get interface information", args_schema=Iwconfig ,return_direct=False)
def get_wireless_interface(params: str) -> str:
"""Return wireless interface information via iwconfig"""
res = shell_tool.run({"commands": [f"iwconfig {params}"]})
return res