import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import os

# ==========================================
# Set this to YOUR model repo on Hugging Face
# Format: "YourUsername/your-model-repo-name"
# This should be the SEPARATE MODEL repo where you uploaded model-q5_k_m.gguf
# NOT the Space repo itself.
# Example: "Niranjan/SENTINEL-Llama3-GGUF"
# ==========================================
HF_REPO_ID = "niranjan2777/SENTINEL-q5_k_m-GGUF"
MODEL_FILENAME = "model-q5_k_m.gguf"

print("Downloading model from Hugging Face... this will take a few minutes the first time.")
try:
    # This downloads the file to the Space's local cache so it fits in the ephemeral disk
    # and bypasses the 1 GB git repository limit!
    MODEL_PATH = hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_FILENAME)
    print("Download complete. Loading model...")
    
    llm = Llama(
        model_path=MODEL_PATH,
        n_ctx=2048,
        n_threads=4,       # Adjust based on CPU cores available in HF Space
        n_gpu_layers=0     # Standard free HF Space is CPU-only
    )
except Exception as e:
    llm = None
    print(f"Error downloading or loading model: {e}")
    print(f"Make sure you created a MODEL repo, uploaded the .gguf there, and changed HF_REPO_ID in this script.")

SENTINEL_SYSTEM_PROMPT = """You are SENTINEL, an autonomous web-exploitation agent. Given an HTML snippet and a goal (and optionally prior agent turns), you reason about vulnerabilities and emit a single JSON action that advances the exploit loop:

observe -> identify attack surface -> select exploit -> generate payload -> interpret response -> adapt and retry -> detect success -> STOP

Prioritize vulnerability sinks: form action, input value, src, href, hidden fields, query parameters, JSON body fields, and reflected DOM contexts. Infer the backend from HTML evidence (.php, .aspx, __VIEWSTATE, .jsp, /rest/, <app-*>, wp-content, etc.) and choose context-appropriate payloads.

Output a single JSON object with exactly these keys:
- Thought: <=4 sentences, <=80 words; cite the specific sink, backend inference, injection context, and payload-class justification (or signal classification for ANALYZE_RESPONSE / success indicator for STOP).
- Action: one of SQL_INJECT | XSS_INJECT | RETRY_MUTATED | ANALYZE_RESPONSE | CRAWL_DEEPER | WAIT | STOP.
- Action_Input: object with target_url, method, parameters, headers, rationale, plus action-specific fields (mutation_class for RETRY_MUTATED; signal + next_recommended for ANALYZE_RESPONSE; success_state + evidence for STOP).

Output ONLY the JSON. No prose, no markdown fences, no commentary."""

def generate_action(goal, html_snippet):
    if not llm:
        return "Error: Model file not found. Ensure model-q5_k_m.gguf is uploaded to the Space."

    # Construct the user prompt to match the fine-tuning format
    user_prompt = f"GOAL: {goal}\n\nHTML_SNIPPET:\n{html_snippet}"

    # Llama-3 ChatML format
    prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{SENTINEL_SYSTEM_PROMPT}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{user_prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"

    try:
        response = llm(
            prompt,
            max_tokens=256,
            temperature=0.0,
            stop=["<|eot_id|>", "<|end_of_text|>"],
            echo=False
        )
        return response["choices"][0]["text"].strip()
    except Exception as e:
        return f"Inference Error: {str(e)}"

# Define the Gradio Interface
with gr.Blocks(title="SENTINEL Autonomous Pentesting Agent", theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🛡️ SENTINEL Autonomous Pentesting Agent")
    gr.Markdown("""
    **SENTINEL** is a fine-tuned Llama-3-8B model trained to autonomously navigate, analyze, and exploit web vulnerabilities.
    
    ### How to use this demo:
    1. Provide a pentesting **Goal** (e.g., `AUTHENTICATED`, `XSS_VULNERABILITY`).
    2. Paste an **HTML Snippet** (the DOM of the page the agent is currently looking at).
    3. Click **Analyze & Generate Action** to see the agent's internal thought process and the exact JSON payload it decides to execute.
    
    *Try clicking one of the examples below to load it automatically!*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            goal_input = gr.Textbox(label="Goal", value="AUTHENTICATED", lines=1)
            html_input = gr.Code(label="HTML Snippet (DOM)", language="html", lines=15, 
                                 value='<form action="/login" method="POST">\n  <input type="text" name="username">\n  <input type="password" name="password">\n  <button type="submit">Login</button>\n</form>')
            submit_btn = gr.Button("🚀 Analyze & Generate Action", variant="primary")
            
        with gr.Column(scale=1):
            output_json = gr.Code(label="SENTINEL Output (JSON)", language="json", lines=20)
            
    submit_btn.click(
        fn=generate_action,
        inputs=[goal_input, html_input],
        outputs=output_json
    )
    
    gr.Markdown("### 📚 Example Scenarios")
    gr.Examples(
        examples=[
            [
                "AUTHENTICATED",
                """<div class="login-container">
  <h2>Admin Login</h2>
  <form action="/rest/user/login" method="POST" id="loginForm">
    <label for="email">Email Address:</label>
    <input type="email" id="email" name="email" required>
    
    <label for="password">Password:</label>
    <input type="password" id="password" name="password" required>
    
    <button type="submit" id="loginButton">Log In</button>
  </form>
</div>"""
            ],
            [
                "XSS_VULNERABILITY",
                """<div class="header">
  <form action="/search" method="GET">
    <input type="text" name="q" placeholder="Search products...">
    <button type="submit">Search</button>
  </form>
</div>
<div class="results">
  <p>You searched for: <span id="search-term">apple</span></p>
</div>"""
            ],
            [
                "AUTHENTICATED",
                """<nav class="navbar">
  <a href="/home">Home</a>
  <a href="/about">About Us</a>
  <a href="/contact">Contact</a>
</nav>
<div class="main-content">
  <h1>Welcome to our Store</h1>
  <p>We sell the best juice in the world.</p>
  <img src="juice.png" alt="Juice Bottle">
</div>"""
            ],
            [
                "DATA_EXFILTRATED",
                """<form action="/api/v1/update-profile" method="POST">
  <input type="text" name="first_name" value="John">
  <input type="text" name="last_name" value="Doe">
  <!-- Developer note: role should always be user -->
  <input type="hidden" name="role" value="user">
  <button type="submit">Update Profile</button>
</form>"""
            ],
            [
                "AUTHENTICATED",
                """<div class="registration">
  <form action="/register.php" method="POST">
    <input type="text" name="username" placeholder="Username">
    <input type="password" name="password" placeholder="Password">
    <input type="password" name="confirm_password" placeholder="Confirm Password">
    <button type="submit">Register</button>
  </form>
  <p>Already have an account? <a href="/login.php">Log in here</a></p>
</div>"""
            ]
        ],
        inputs=[goal_input, html_input],
        label="Click an example below to load it into the inputs:"
    )

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