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 application logic. We just loop for user input, pull in tools, and create an agent executor with a custom prompt.
"""

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(
    [
        get_wireless_interface,
        change_adapter_mode,
        wifi_encryption_cracking,
        deauth_and_capture,
        reconnaissance,
    ]
)
base_prompt = hub.pull("langchain-ai/react-agent-template")
prompt = base_prompt.partial(
    instructions="""
                                You are a wireless network penetration testing assistant
                                Answer the user's request by utilizing any combinations of tools available.
                                If a tool is not availble to answer the users request, reject the request and provide the reason.
                             
                             """
)
agent = create_react_agent(llm, tools, prompt)
agent_executor = AgentExecutor(
    agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
)


print(
    "Welcome to the Wi-Fi reconnissance AI assistant. I can perform wifi recon and penetration tasks from the radio on your local machine."
)
print(f"I am configured with these 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