Update app.py

app.py

@@ -1,52 +1,121 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import subprocess
+import pyshark
+from smolagents import Tool
 
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# 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
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": system_message}]
-    
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+        model = next(iter(list_models(filter=task, sort="downloads", direction=-1)))
+        return model.id
 
-    messages.append({"role": "user", "content": message})
-    
-    response = ""
-    
+# 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:
-        for message in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            token = message.choices[0].delta.content
-            response += token
-            yield response
+        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:
-        yield f"Error: {str(e)}"
-
-
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+        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()