import gradio as gr
import subprocess
import pyshark
from smolagents import Tool

# HFModelDownloadsTool Definition
class HFModelDownloadsTool(Tool):
    name = "model_download_counter"
    description = """
    This is a tool that returns the most downloaded model of a given task on the Hugging Face Hub.
    It returns the name of the checkpoint."""
    inputs = {
        "task": {
            "type": "string",
            "description": "the task category (such as text-classification, depth-estimation, etc)",
        }
    }
    output_type = "string"

    def forward(self, task: str):
        from huggingface_hub import list_models

        model = next(iter(list_models(filter=task, sort="downloads", direction=-1)))
        return model.id

# Instantiate the tool
model_downloads_tool = HFModelDownloadsTool()

# Function to integrate HFModelDownloadsTool into Gradio
def get_most_downloaded_model(task):
    if not task:
        return "Error: Task cannot be empty."
    try:
        return model_downloads_tool.forward(task)
    except Exception as e:
        return f"Error: {str(e)}"

# Other functions (Nmap, Nikto, Hydra, PCAP) remain the same
def run_nmap(target):
    if not target:
        return "Error: Target cannot be empty."
    try:
        result = subprocess.run(["nmap", target], capture_output=True, text=True)
        return result.stdout if result.returncode == 0 else "Error running Nmap scan."
    except Exception as e:
        return f"Error: {str(e)}"

def run_nikto(target):
    if not target:
        return "Error: Target cannot be empty."
    try:
        result = subprocess.run(["nikto", "-h", target], capture_output=True, text=True)
        return result.stdout if result.returncode == 0 else "Error running Nikto scan."
    except Exception as e:
        return f"Error: {str(e)}"

def run_hydra(target, service, wordlist):
    if not target or not service or not wordlist:
        return "Error: Target, service, and wordlist cannot be empty."
    try:
        result = subprocess.run(
            ["hydra", "-l", "admin", "-P", wordlist, f"{service}://{target}"], 
            capture_output=True, text=True
        )
        return result.stdout if result.returncode == 0 else "Error running Hydra attack."
    except Exception as e:
        return f"Error: {str(e)}"

def analyze_pcap(file_path):
    if not file_path:
        return "Error: Please upload a valid PCAP file."
    try:
        capture = pyshark.FileCapture(file_path['name'])
        summary = "\n".join([str(pkt) for pkt in capture])
        return f"PCAP Analysis Completed. Summary:\n{summary}"
    except Exception as e:
        return f"Error analyzing PCAP file: {str(e)}"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Cybersecurity Scanning Tool with Hugging Face Integration")

    # Nmap Scan
    with gr.Row():
        nmap_target = gr.Textbox(label="Enter Target IP for Nmap Scan")
        nmap_button = gr.Button("Run Nmap Scan")
    nmap_result = gr.Textbox(label="Nmap Scan Results", interactive=False)
    nmap_button.click(run_nmap, inputs=nmap_target, outputs=nmap_result)

    # Nikto Scan
    with gr.Row():
        nikto_target = gr.Textbox(label="Enter Web Server URL for Nikto Scan")
        nikto_button = gr.Button("Run Nikto Scan")
    nikto_result = gr.Textbox(label="Nikto Scan Results", interactive=False)
    nikto_button.click(run_nikto, inputs=nikto_target, outputs=nikto_result)

    # Hydra Attack
    with gr.Row():
        hydra_target = gr.Textbox(label="Enter Target IP for Hydra Attack")
        hydra_service = gr.Textbox(label="Enter Service (e.g., ssh, ftp)")
        hydra_wordlist = gr.Textbox(label="Enter Path to Wordlist")
        hydra_button = gr.Button("Run Hydra Attack")
    hydra_result = gr.Textbox(label="Hydra Attack Results", interactive=False)
    hydra_button.click(run_hydra, inputs=[hydra_target, hydra_service, hydra_wordlist], outputs=hydra_result)

    # PCAP Analysis (Wireshark)
    with gr.Row():
        pcap_file = gr.File(label="Upload PCAP File")
        pcap_button = gr.Button("Analyze PCAP File")
    pcap_result = gr.Textbox(label="PCAP Analysis Results", interactive=False)
    pcap_button.click(analyze_pcap, inputs=pcap_file, outputs=pcap_result)

    # Hugging Face Most Downloaded Model Tool
    with gr.Row():
        hf_task_input = gr.Textbox(label="Enter Task Category (e.g., text-classification)")
        hf_button = gr.Button("Get Most Downloaded Model")
    hf_result = gr.Textbox(label="Most Downloaded Model", interactive=False)
    hf_button.click(get_most_downloaded_model, inputs=hf_task_input, outputs=hf_result)

if __name__ == "__main__":
    demo.launch()