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):
    params: str = Field(
        description="Should be command line parameters to 'aircrack-ng' to perform some kind of wifi encryption cracking"
    )
    def get_wordlists():
        directory = "wordlists"
        # Check if the directory exists
        if not os.path.exists(directory):
            raise FileNotFoundError(f"The directory {directory} does not exist.")

        # List all files in the directory
        files = [
            file
            for file in os.listdir(directory)
            if os.path.isfile(os.path.join(directory, file))
        ]

        # Check if the list is empty
        if not files:
            raise Exception(f"No files found in the directory {directory}.")

        # Return the first file, for the sake of simplicity. TODO: accomodate the possibility of multiple word list files
        return files[0]


@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):
    params: 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):
#     params: 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 timeout 15s airodump-ng {params}"]})
#     return res

class IwScan(BaseModel):
    interface: str = Field(
        description="Should be a wireless interface name, used as a paramater to 'iw' to scan for wifi networks"
    )

    network: str = Field(
        description="Should be the name or SSID of the wifi network you are interested in"
    )

@tool(
    "Perform wifi scanning with iw",
    args_schema=IwScan,
    return_direct=False,
)
def wifi_network_reconnissance(interface: str, network: str) -> str:
    """Can pass a wireless interface name and wifi network name to return technical information about a wifi network"""
    res = shell_tool.run({"commands": [f'sudo iw {interface} scan | grep -B 9 -A 206 "{network}"']})
    return res


class ChangeMonitorMode(BaseModel):
    params: 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 none are needed, this should be left blank"
    )

@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