import os
import subprocess
import time
import sys
import argparse
import threading
import json
import signal
from groq import Groq
from colorama import init, Fore, Style
from pyfiglet import Figlet
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn

# Initialize colorama
init(autoreset=True, strip=False)

# Initialize rich console
console = Console(stderr=True, style="bold red")

# Define config path
CONFIG_PATH = os.path.expanduser("~/.pentest_ai_config.json")

# Supported tools
TOOLS = ["nmap", "sqlmap", "aircrack-ng", "airodump-ng", "metasploit", "nikto", "hashid", "hashcat", "john"]

# Global variables for command execution state
current_process = None
is_command_running = False
was_interrupted = False

# Signal handler for Ctrl+C
def signal_handler(sig, frame):
    global current_process, is_command_running, was_interrupted
    if is_command_running and current_process:
        current_process.terminate()
        try:
            current_process.wait(timeout=2)
        except subprocess.TimeoutExpired:
            current_process.kill()
        console.print("\n[bold yellow]Command stopped by user! 🛑[/]")
        was_interrupted = True
    else:
        # Exit immediately if no command is running
        console.print("\n[bold cyan]Exiting... 👋[/]")
        sys.exit(0)

# Register the signal handler
signal.signal(signal.SIGINT, signal_handler)

# Typewriter animation
def typewriter_print(text, delay=0.03):
    if isinstance(text, str):
        text = Text(text)
    for char in str(text):
        styled_char = Text(char, style=text.style)
        console.print(styled_char, end='', soft_wrap=True)
        sys.stdout.flush()
        time.sleep(delay)
    console.print()

# Display animated logo
def display_logo():
    f = Figlet(font='slant')
    logo = f.renderText("PenTest AI")
    color_styles = ["red", "yellow", "green", "cyan", "blue", "magenta"]
    for i, line in enumerate(logo.splitlines()):
        styled_line = Text(line, style=f"bold {color_styles[i % len(color_styles)]}")
        console.print(styled_line)
        time.sleep(0.2)
    typewriter_print(Text("AI Pentesting Chatbot 🔥️ Powered by AKM Korishee Apurbo ☠️", style="bold cyan"))

def load_api_key():
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH, 'r') as f:
            config = json.load(f)
            return config.get("api_key")
    return None

def save_api_key(api_key):
    config = {"api_key": api_key}
    with open(CONFIG_PATH, 'w') as f:
        json.dump(config, f)
    print(f"{Fore.GREEN}{Style.BRIGHT}API key saved to 👉️ {CONFIG_PATH}! 👈️{Style.RESET_ALL}")

def validate_api_key(api_key):
    try:
        client = Groq(api_key=api_key)
        client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=[{"role": "system", "content": "Test"}],
            max_tokens=1
        )
        return True
    except Exception:
        return False

def run_command_in_background(command):
    global current_process, is_command_running, was_interrupted
    try:
        is_command_running = True
        was_interrupted = False
        process = subprocess.Popen(
            command,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        current_process = process

        rotators = ['🌑️', '🌒️', '🌓️', '🌔️', '🌕️', '🌖️', '🌗️', '🌘️']
        spinners = ['⠦', '⠧', '⠇', '⠏', '⠋', '⠙', '⠹', '⠸']
        step = 0

        with Progress(
            TextColumn("{task.description}"),
            BarColumn(bar_width=40),
            TimeElapsedColumn(),
            transient=True
        ) as progress:
            task = progress.add_task(
                description=Text(f"[bright_red] {spinners[step % len(spinners)]} Running '{command}' {rotators[step % len(rotators)]}"),
                total=None
            )

            while process.poll() is None:
                step += 1
                progress.update(
                    task,
                    description=Text(f"[bright_red] {spinners[step % len(spinners)]} Running '{command}' {rotators[step % len(rotators)]}"))
                time.sleep(0.1)

            elapsed_time = progress.tasks[task].elapsed
            minutes, seconds = divmod(int(elapsed_time), 60)
            time_str = f"{minutes}:{seconds:02d}"
            progress.update(
                task,
                description=Text(f"[bold green]🔥️ Completed '{command}' {time_str}", style="#c20693"),
                completed=True
            )
            time.sleep(2)

        stdout, stderr = process.communicate()
        current_process = None
        is_command_running = False

        if was_interrupted:
            return None, "Command was interrupted by user."
        elif process.returncode == 0:
            return stdout.strip() if stdout else "Command completed successfully! ☠️", None
        else:
            error_msg = stderr.strip() or "Command failed with no specific error output."
            return None, error_msg
    except Exception as e:
        current_process = None
        is_command_running = False
        was_interrupted = False
        return None, f"Unexpected error running command: {str(e)} 😵"

def analyze_and_fix_error(client, command, error):
    system_prompt = (
        "You are 'Korishee,' an expert AI pentesting assistant running on Kali Linux 🌐. "
        "Given a failed command and its error message, analyze the issue and suggest a fix. "
        "Respond with ONLY these two elements:\n"
        "1. 'ANALYSIS: <detailed analysis of the error>'\n"
        "2. 'FIX_SUGGESTION: <suggested fix or command>' (use a concrete command if possible)\n"
        "Keep it concise, actionable, and engaging! 😎"
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"Command: {command}\nError: {error}"}
    ]
    try:
        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            max_tokens=200,
            temperature=0.5
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"ANALYSIS: Couldn’t analyze due to AI error: {str(e)}\nFIX_SUGGESTION: Please try again later or check manually! 😞"

def get_groq_response(client, user_input, chat_history=None):
    if chat_history is None:
        chat_history = []

    system_prompt = (
        "You are 'Korishee,' an expert AI pentesting assistant running on Kali Linux 🌐, created by AKM Korishee Apurbo. "
        "Guide the user interactively to build a pentesting command step-by-step. "
        "Respond with ONLY these three elements when generating a command:\n"
        "1. 'COMMAND: <full_command>' (no backticks, use placeholders like <hash> or <file> only if input is incomplete)\n"
        "2. 'EXPLANATION: <brief explanation>'\n"
        "3. 'QUESTION: <question>' with emojis 🤔 (omit if no further input needed)\n"
        "For other inputs:\n"
        "- If vague (e.g., 'hi'), respond with a greeting and 'QUESTION: <question>'\n"
        "- If asked for your name, say 'I’m Korishee, your pentesting sidekick! 🔥'\n"
        "- If asked for your creator, say 'I was crafted by the awesome AKM Korishee Apurbo! 🚀'\n"
        "Keep it fun, concise, and engaging! 😎"
    )

    messages = [{"role": "system", "content": system_prompt}] + chat_history + [{"role": "user", "content": user_input}]
    try:
        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            max_tokens=300,
            temperature=0.5
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        console.print(f"[bold red]Oops! AI Error: {str(e)} 😢[/]")
        return "Sorry, I hit a snag with the AI service! Try again later. 😞"

def main():
    parser = argparse.ArgumentParser(description="Pentest-AI: AI-powered pentesting tool")
    parser.add_argument("--api", type=str, help="Set the Groq API key (overrides config file if provided)")
    args = parser.parse_args()

    api_key = args.api if args.api else load_api_key()
    if not api_key:
        console.print("[bold red]No API key found! 😱 Please provide it via --api <key> to save it to config.[/]")
        api_key = input("Enter your Groq API key: ").strip()
        if validate_api_key(api_key):
            save_api_key(api_key)
        else:
            console.print("[bold red]Invalid API key! 😞 Exiting...[/]")
            sys.exit(1)
    elif not validate_api_key(api_key):
        console.print("[bold red]Stored API key is invalid! 😞 Please provide a new one via --api <key>.[/]")
        sys.exit(1)

    try:
        client = Groq(api_key=api_key)
    except Exception as e:
        console.print(f"[bold red]Failed to initialize Groq client: {str(e)} 😱[/]")
        sys.exit(1)

    display_logo()
    typewriter_print(Text("Welcome to Pentest AI! 🎉 Type 'exit' or use Ctrl+C to quit. 🚪", style="bold green"))

    chat_history = []
    pending_command = None

    while True:
        try:
            user_input = input(f"{Fore.YELLOW}{Style.BRIGHT}You 👤: {Style.RESET_ALL}")
        except KeyboardInterrupt:
            console.print("\n[bold cyan]Exiting... 👋[/]")
            sys.exit(0)
        except Exception as e:
            console.print(f"[bold red]Input error: {str(e)}. Falling back to basic input.[/]")
            user_input = input(f"{Fore.YELLOW}{Style.BRIGHT}You 👤: {Style.RESET_ALL}")

        if user_input.lower() == "exit":
            typewriter_print(Text("Goodbye! 👋 Catch you later! 😄", style="bold cyan"))
            break

        if pending_command:
            command = pending_command.replace("<hash>", user_input).replace("<file>", user_input).replace("<target_IP>", user_input)
            typewriter_print(Text(f"Korishee 🤖: COMMAND: {command} ⚡", style="bold magenta"))
            result, error = run_command_in_background(command)
            if result:
                console.print(Panel(
                    result,
                    title="Result 🌟",
                    border_style="bold green_yellow",
                    padding=(1, 2),
                    expand=False
                ))
            elif error and not was_interrupted:
                typewriter_print(Text(f"Error 🚨: {error}", style="bold red"))
                fix_response = analyze_and_fix_error(client, command, error)
                lines = fix_response.split("\n")
                for line in lines:
                    if line.startswith("ANALYSIS:"):
                        typewriter_print(Text(f"Analysis 🔍: {line.replace('ANALYSIS: ', '')}", style="bold yellow"))
                    elif line.startswith("FIX_SUGGESTION:"):
                        typewriter_print(Text(f"Fix Suggestion 💡: {line.replace('FIX_SUGGESTION: ', '')}", style="bold green"))
            elif was_interrupted:
                typewriter_print(Text("Skipping error analysis. 😊", style="bold yellow"))
            pending_command = None
            continue

        chat_history.append({"role": "user", "content": user_input})
        ai_response = get_groq_response(client, user_input, chat_history)
        
        command = None
        question_present = False
        response_text = ""
        lines = ai_response.split("\n")
        for line in lines:
            if line.startswith("COMMAND:"):
                command = line.replace("COMMAND: ", "").strip()
                typewriter_print(Text(f"Korishee 🤖: COMMAND: {command} ⚡", style="bold magenta"))
            elif line.startswith("EXPLANATION:"):
                typewriter_print(Text(f"Explanation 📚: {line.replace('EXPLANATION: ', '')}", style="cyan"))
            elif line.startswith("QUESTION:"):
                question_present = True
                response_text = line.replace("QUESTION: ", "").strip()
            else:
                response_text += line + "\n"

        if response_text and not question_present:
            typewriter_print(Text(f"Korishee 🤖: {response_text.strip()} 💬", style="bold magenta"))

        if question_present:
            typewriter_print(Text(f"Korishee asks ❓: {response_text}", style="bold green"))

        if command and all(x not in command for x in ["<", ">"]):
            result, error = run_command_in_background(command)
            if result:
                console.print(Panel(
                    result,
                    title="Result 🌟",
                    border_style="bold green_yellow",
                    padding=(1, 2),
                    expand=False
                ))
            elif error and not was_interrupted:
                typewriter_print(Text(f"Error 🚨: {error}", style="bold red"))
                fix_response = analyze_and_fix_error(client, command, error)
                lines = fix_response.split("\n")
                for line in lines:
                    if line.startswith("ANALYSIS:"):
                        typewriter_print(Text(f"Analysis 🔍: {line.replace('ANALYSIS: ', '')}", style="bold yellow"))
                    elif line.startswith("FIX_SUGGESTION:"):
                        typewriter_print(Text(f"Fix Suggestion 💡: {line.replace('FIX_SUGGESTION: ', '')}", style="bold green"))
            elif was_interrupted:
                typewriter_print(Text("Command was stopped, skipping error analysis. 😊", style="bold yellow"))
        elif command:
            pending_command = command

        chat_history.append({"role": "assistant", "content": ai_response})

if __name__ == "__main__":
    main()
