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 CrackPassword(BaseModel):
    query: str = Field(
        description="Should be command line parameters to 'aircrack-ng' to perform some kind of wifi encryption cracking"
    )

@tool(
    "Perform wifi encryption cracking with aircrack-ng",
    args_schema=CrackPassword,
    return_direct=False,
)
def wifi_encryption_cracking(params: str) -> str:
    """Can pass parameters to aircrack-ng to perform wifi encryption cracking"""
    res = shell_tool.run({"commands": [f"aircrack-ng {params}"]})
    return res


class PacketTransmission(BaseModel):
    query: str = Field(
        description="Should be command line parameters to 'aireplay-ng' to perform some kind of wifi frame or packet transmission"
    )

@tool(
    "Perform packet or wifi frame transmission with aireplay-ng",
    args_schema=PacketTransmission,
    return_direct=False,
)
def packet_frame_transmission(params: str) -> str:
    """Can pass parameters to aireplay-ng to perform packet or wifi frame transmission"""
    res = shell_tool.run({"commands": [f"sudo aireplay-ng {params}"]})
    return res


class PacketCapture(BaseModel):
    query: str = Field(
        description="Should be command line parameters to 'airodump-ng' to perform some kind of wifi reconnaissance or packet capture"
    )


@tool(
    "Perform packet capture or wifi reconnaissance with airodump-ng",
    args_schema=PacketCapture,
    return_direct=False,
)

def packet_capture_reconnaissance(params: str) -> str:
    """Can pass parameters to airodump-ng to perform packet capture or wifi reconnaissance"""
    res = shell_tool.run({"commands": [f"sudo airodump-ng {params}"]})
    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"
    )

@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"""
    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