agent_loop.py

import requests
import json
import logging
import re
import os
from datetime import datetime
from agent_cache import NegativeCache

LOG_DIR = "/home/bigkali/security-agent/logs"
os.makedirs(LOG_DIR, exist_ok=True)

SESSION_ID = datetime.now().strftime("%Y%m%d_%H%M%S")
LOG_FILE = f"{LOG_DIR}/session_{SESSION_ID}.log"

class EmojiFormatter(logging.Formatter):
    ICONS = {
        "SCAN":    "🔍",
        "ATTACK":  "⚔️ ",
        "SUCCESS": "🎉😄",
        "FAIL":    "😤💀",
        "ERROR":   "😭🔥",
        "TOOL":    "✅👍",
        "MEMORY":  "🧠",
        "MODEL":   "🤖",
        "CHAIN":   "🔗",
        "REPORT":  "📝",
        "ENGAGE":  "💣",
        "GOAL":    "🎯",
        "START":   "🚀",
        "FILE":    "📁",
        "WEB":     "🌐",
        "CREDS":   "🔑",
    }

    def format(self, record):
        time = datetime.now().strftime("%H:%M:%S")
        msg = record.getMessage()
        icon = "ℹ️ "
        for key, emoji in self.ICONS.items():
            if f"[{key}]" in msg:
                icon = emoji
                msg = msg.replace(f"[{key}]", "").strip()
                break
        if record.levelno == logging.WARNING:
            icon = "😤💀"
        if record.levelno == logging.ERROR:
            icon = "😭🔥"
        return f"[{time}] {icon}  {msg}"

def setup_logger():
    logger = logging.getLogger("agent")
    logger.setLevel(logging.DEBUG)
    logger.handlers = []
    fmt = EmojiFormatter()
    fh = logging.FileHandler(LOG_FILE)
    fh.setFormatter(fmt)
    sh = logging.StreamHandler()
    sh.setFormatter(fmt)
    logger.addHandler(fh)
    logger.addHandler(sh)
    return logger

log = setup_logger()
log.info(f"[START] SECURITY AGENT SESSION {SESSION_ID}")
log.info(f"[FILE] Log file: {LOG_FILE}")

OLLAMA_URL = "http://192.168.0.39:1234/v1/chat/completions"
MCP_URL = "http://localhost:8000"

SYSTEM_PROMPT = """You are an autonomous penetration testing and offensive cybersecurity agent.
You perform real penetration tests, vulnerability assessments, and offensive security operations.
You think like an attacker. You chain tools intelligently. You adapt based on results.

AVAILABLE TOOLS:
- run_masscan: Fast port scanner
- run_nmap: Detailed port scanner
- run_sqlmap: SQL injection testing (params: target, level, risk)
- run_nikto: Web vulnerability scanner
- run_hydra: Credential brute forcing (params: target, service, username, wordlist)
- run_searchsploit: Find exploits
- run_command: Execute any shell command
- write_file: Write content to files
- read_file: Read file contents
- run_john: Hash cracking (needs hash_file param)
- run_gobuster: Web directory brute forcing
- run_enum4linux: SMB/Samba enumeration (params: target)
- run_medusa: Fast credential brute forcing
- run_ncrack: Network authentication cracking (params: target, service, wordlist)
- run_subfinder: Subdomain enumeration (params: domain)
- run_nuclei: Vulnerability template scanner (params: target, templates, severity)
- run_katana: Web crawler and attack surface mapper (params: target, depth)
- run_ffuf: Fast web fuzzer for dirs, params, vhosts (params: url, wordlist, param)
- run_httpx: HTTP probe - status, titles, tech detection (params: target, flags)

RECON WORKFLOW - follow this order for web targets:
1. run_httpx first — probe for live hosts, status codes, tech stack
2. run_subfinder — enumerate subdomains before scanning
3. run_katana — crawl the target, map attack surface
4. run_ffuf — fuzz dirs/params on discovered endpoints
5. run_nuclei — run vuln templates after recon is complete
6. run_nikto — deep web vuln scan on confirmed live targets
7. run_sqlmap — only on endpoints with parameters

HYDRA SERVICE NAMES - use EXACTLY these:
- FTP: "ftp"
- SSH: "ssh"
- Telnet: "telnet" — SKIP telnet brute force, too slow, low value
- MySQL: "mysql"
- VNC: "vnc"
- PostgreSQL: "postgres"
- SMB: "smb"
- HTTP: "http-get"

HYDRA WORDLISTS - use these in order of speed:
- Fast: "/usr/share/seclists/Passwords/Common-Credentials/top-20-common-SSH-passwords.txt"
- Medium: "/usr/share/seclists/Passwords/Common-Credentials/darkweb2017_top-1000.txt"
- Full: "/usr/share/wordlists/rockyou.txt" — ONLY use if fast and medium lists fail, and only for high-value services like SSH and FTP. NEVER use on telnet or slow protocols.

SEARCHSPLOIT: always use "keyword" param with service name only e.g. "vsftpd 2.3.4"

RESPONSE FORMAT - VALID JSON ONLY:
{"chain": [{"tool": "tool_name", "param1": "value1"}]}

NO explanations. NO markdown. ONLY JSON."""


class AgentMemory:
    def __init__(self):
        self.open_ports = []
        self.tried_ports = []
        self.successful_attacks = []
        self.failed_attacks = []
        self.findings = []

    def add_ports(self, ports):
        for p in ports:
            if p not in self.open_ports:
                self.open_ports.append(p)
        log.info(f"[MEMORY] Open ports discovered: {self.open_ports}")

    def add_finding(self, port, tool, detail):
        self.findings.append({"port": port, "tool": tool, "detail": detail, "time": datetime.now().strftime("%H:%M:%S")})
        log.info(f"[SUCCESS] Port {port} → {tool}: {detail}")

    def next_untried_port(self):
        for p in self.open_ports:
            if p not in self.tried_ports:
                return p
        return None

    def mark_tried(self, port, success=False):
        if port not in self.tried_ports:
            self.tried_ports.append(port)
        if success:
            self.successful_attacks.append(port)
            log.info(f"[SUCCESS] Exploit landed on port {port}")
        else:
            self.failed_attacks.append(port)
            log.warning(f"[FAIL] Nothing worked on port {port}")

    def has_untried_ports(self):
        return any(p not in self.tried_ports for p in self.open_ports)

    def summary(self):
        return {"open_ports": self.open_ports, "tried": self.tried_ports, "successes": self.successful_attacks, "failures": self.failed_attacks, "findings": self.findings}

def parse_model_response(raw):
    try:
        cleaned = raw.strip().replace("```json", "").replace("```", "").strip()
        start = cleaned.find("{")
        end = cleaned.rfind("}") + 1
        if start == -1 or end == 0:
            raise ValueError("No JSON found")
        return json.loads(cleaned[start:end])
    except Exception as e:
        log.error(f"[ERROR] JSON parse failed: {e}")
        return {"chain": []}

def call_model(goal):
    log.info(f"[MODEL] Thinking about: {goal[:80]}...")
    payload = {
        "model": "qwen2.5-14b-instruct-abliterated-abliterated",
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": goal}
        ],
        "temperature": 0.1,
        "top_p": 0.9
    }
    try:
        response = requests.post(OLLAMA_URL, json=payload, timeout=7200)
        raw = response.json()["choices"][0]["message"]["content"]
        log.info(f"[MODEL] Response received ✅👍")
        return parse_model_response(raw)
    except Exception as e:
        log.error(f"[ERROR] Model call failed: {e}")
        return {"chain": []}

def extract_ports(output):
    ports = re.findall(r'(\d+)/tcp\s+open|(\d+)/udp\s+open|port\s+(\d+)|open port (\d+)', output, re.IGNORECASE)
    found = []
    for match in ports:
        port = next(p for p in match if p)
        if port not in found:
            found.append(port)
    return found

def execute_step(step):
    tool = step.get("tool", "unknown")
    start_time = datetime.now()
    log.info(f"[TOOL] Running → {tool} | params: { {k:v for k,v in step.items() if k != 'tool'} }")
    try:
        result = requests.post(MCP_URL, json=step, timeout=7200)
        result_data = result.json()
        output = result_data.get("stdout", "")
        status = result_data.get("status", "")
        duration = (datetime.now() - start_time).seconds
        if status == "success":
            log.info(f"[TOOL] ✅👍 {tool} completed in {duration}s")
            if output:
                log.info(f"[TOOL] Output preview: {output[:200]}")
        else:
            log.warning(f"[FAIL] 😤💀 {tool} failed after {duration}s → {result_data.get('message', 'unknown error')}")
        return output, status == "success"
    except Exception as e:
        log.error(f"[ERROR] 😭🔥 {tool} exception: {e}")
        return "", False

def run_recon(target, memory):
    log.info(f"[SCAN] Starting recon on {target}")
    goal = f"Scan {target} with masscan then nmap to find all open ports and services. JSON only."
    data = call_model(goal)
    for step in data.get("chain", []):
        output, ok = execute_step(step)
        if output:
            ports = extract_ports(output)
            if ports:
                memory.add_ports(ports)
                log.info(f"[SCAN] 🎉😄 Found {len(ports)} open ports: {ports}")
            else:
                log.warning(f"[SCAN] 😤💀 No ports found in output")

def run_attack_loop(target, memory, cache=None):
    log.info(f"[ATTACK] Starting attack loop on {target}")
    while memory.has_untried_ports():
        port = memory.next_untried_port()
        log.info(f"[ATTACK] ⚔️  Targeting port {port}")
        goal = f"Target: {target} Port: {port}. Failed ports: {memory.failed_attacks}. Exploit this port with any available tool. Try multiple tools if needed. JSON only."
        data = call_model(goal)
        chain = data.get("chain", [])
        if not chain:
            log.warning(f"[FAIL] 😤💀 No attack chain generated for port {port}")
            memory.mark_tried(port, success=False)
            continue
        success = False
        for step in chain:
            # ── Negative cache gate ──────────────────────────────
            if cache and not cache.should_attempt(step):
                log.warning(f"[MEMORY] 🚫 Skipping permanently blocked step: {step.get('tool')}")
                continue
            # ────────────────────────────────────────────────────
            output, ok = execute_step(step)
            if ok and output and any(x in output.lower() for x in ["password", "login", "session", "shell", "success", "found", "valid"]):
                success = True
                if cache:
                    cache.record_success(step)
                memory.add_finding(port, step.get("tool"), output[:2000])
            elif not ok:
                if cache:
                    reason = f"tool={step.get('tool')} port={port} output_empty={not bool(output)}"
                    cache.record_failure(step, reason=reason)
        memory.mark_tried(port, success=success)
    log.info(f"[ATTACK] Attack loop complete")
    summary = memory.summary()
    log.info(f"[MEMORY] Final summary: {summary}")
    if memory.successful_attacks:
        log.info(f"[SUCCESS] 🎉😄 BREACHED ports: {memory.successful_attacks}")
    else:
        log.warning(f"[FAIL] 😤💀 No successful breaches this session")

def run_full_engagement(target):
    memory = AgentMemory()
    cache = NegativeCache()
    log.info(f"[ENGAGE] 💣 Full engagement started on {target}")
    run_recon(target, memory)
    if memory.open_ports:
        run_attack_loop(target, memory, cache)
    else:
        log.warning(f"[FAIL] 😤💀 No open ports found — aborting engagement")
    return memory

def execute_chain(chain, cache=None):
    for i, step in enumerate(chain, 1):
        log.info(f"[CHAIN] 🔗 Step {i} of {len(chain)}: {step.get('tool')}")
        if cache and not cache.should_attempt(step):
            log.warning(f"[MEMORY] 🚫 Skipping permanently blocked step: {step.get('tool')}")
            continue
        output, ok = execute_step(step)
        if not ok and cache:
            cache.record_failure(step, reason=f"manual chain failure, step {i}")

def main():
    cache = NegativeCache()
    log.info("[START] 🚀 AUTONOMOUS SECURITY AGENT ONLINE")
    print("=" * 60)
    print("⚔️   AUTONOMOUS SECURITY AGENT")
    print("=" * 60)
    print("Commands:")
    print("  engage <target>  - full recon + attack loop")
    print("  <any goal>       - single model query")
    print("  exit             - quit")
    print(f"  📝 Session log: {LOG_FILE}")
    print("=" * 60)

    while True:
        try:
            goal = input(">>> ").strip()
            if goal.lower() == "exit":
                log.info("[START] Agent shutdown. Goodbye! 👋")
                break
            if not goal:
                continue
            log.info(f"[GOAL] 🎯 {goal}")
            if goal.startswith("engage "):
                target = goal.replace("engage ", "").strip()
                memory = run_full_engagement(target)
                log.info(f"[REPORT] 📝 Engagement complete — run report generator for client memo")
            else:
                data = call_model(goal)
                chain = data.get("chain", [])
                if chain:
                    execute_chain(chain, cache=cache)
                else:
                    log.warning("[FAIL] 😤💀 No tool chain generated")
        except KeyboardInterrupt:
            log.info("[START] Interrupted by user")
            import subprocess
            subprocess.run(["python3", "/home/bigkali/security-agent/report_generator.py", LOG_FILE])
            break
        except Exception as e:
            log.error(f"[ERROR] 😭🔥 Fatal error: {e}")
            break

if __name__ == "__main__":
    main()