tools.py

"""Tool registry for Nova's agentic loop."""

from __future__ import annotations

import json
import re
import subprocess
import sys
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

MEMORY_PATH = Path(__file__).parent / "memory.json"

import os

import requests
from bs4 import BeautifulSoup
from dotenv import load_dotenv

load_dotenv(Path(__file__).parent / ".env")


@dataclass
class Tool:
    name: str
    description: str
    parameters: dict  # JSON Schema object
    fn: Callable[..., str]  # must return a plain string result
    requires_confirmation: bool = False  # prompt user before executing

    def to_ollama(self) -> dict:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


TOOL_REGISTRY: list[Tool] = []


def register(tool: Tool) -> Tool:
    TOOL_REGISTRY.append(tool)
    return tool


def get_tool(name: str) -> Tool | None:
    return next((t for t in TOOL_REGISTRY if t.name == name), None)


def ollama_tools() -> list[dict] | None:
    """Return tools in Ollama's expected format, or None if registry is empty."""
    return [t.to_ollama() for t in TOOL_REGISTRY] or None


def load_memory() -> list[str]:
    if MEMORY_PATH.exists():
        try:
            return json.loads(MEMORY_PATH.read_text())
        except Exception:
            return []
    return []


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------

def _remember_fact(fact: str) -> str:
    facts = load_memory()
    if fact in facts:
        return "Already remembered."
    facts.append(fact)
    MEMORY_PATH.write_text(json.dumps(facts, indent=2))
    return f"Remembered: {fact}"


register(Tool(
    name="remember_fact",
    description="Save a fact to persistent memory so it's available in future sessions. Use when the user explicitly asks you to remember something.",
    parameters={
        "type": "object",
        "properties": {
            "fact": {"type": "string", "description": "The fact to remember"},
        },
        "required": ["fact"],
    },
    fn=_remember_fact,
))


# Patterns that are hard-blocked regardless of user confirmation.
# Focused on catastrophic / irreversible operations only — not exhaustive.
_BLOCKED = [re.compile(p, re.IGNORECASE) for p in [
    r"rm\s+.*-[a-z]*r[a-z]*f[a-z]*\s+(\/|~|\/\*|~\/?\*?)\s*$",  # rm -rf / or ~
    r"rm\s+.*--no-preserve-root",
    r"\bdd\b.+\bof=/dev/",             # dd to raw device
    r"\bmkfs\b",                        # format filesystem
    r":\(\)\s*\{.*\|",                  # fork bomb
    r"(curl|wget)\s+.+\|\s*(ba)?sh",   # pipe to shell
    r">\s*/dev/(sd|nvme|vd|hd)[a-z]",  # write directly to disk
]]

def _run_shell(command: str) -> str:
    # NOTE: prompt injection risk — web search results or other tool output fed
    # back to the model could contain instructions to run malicious commands.
    # User confirmation (requires_confirmation=True) is the primary mitigation.
    if any(p.search(command) for p in _BLOCKED):
        return "Error: command blocked by safety filter."
    try:
        proc = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=30,
            cwd=Path.home(),
        )
        output = (proc.stdout + proc.stderr).strip()
        if len(output) > 4000:
            output = output[:4000] + "\n[output truncated]"
        return output or "(no output)"
    except subprocess.TimeoutExpired:
        return "Error: command timed out after 30 seconds."
    except Exception as e:
        return f"Error: {e}"


register(Tool(
    name="run_shell",
    description="Execute a shell command on the user's local machine and return the real output. Use this whenever the user asks about current system state — running processes, disk usage, files, logs, network, hardware, etc. Do not suggest commands for the user to run; execute them directly using this tool.",
    parameters={
        "type": "object",
        "properties": {
            "command": {"type": "string", "description": "Shell command to run"},
        },
        "required": ["command"],
    },
    fn=_run_shell,
    requires_confirmation=True,
))


# Patterns blocked in Python exec — prevent shell escape and destructive fs ops.
_BLOCKED_PYTHON = [re.compile(p, re.IGNORECASE) for p in [
    r"os\.system\s*\(",
    r"os\.popen\s*\(",
    r"subprocess\.(run|Popen|call|check_output|getoutput)\s*\(",
    r"shutil\.rmtree\s*\(",
    r"os\.(remove|unlink|rmdir)\s*\(",
]]


_CODE_FENCE = re.compile(r"^```[a-z]*\n?|```$", re.MULTILINE)


def _run_python(code: str) -> str:
    # Strip markdown code fences the model may have included.
    code = _CODE_FENCE.sub("", code).strip()
    # NOTE: same prompt injection risk as run_shell — user confirmation is the
    # primary mitigation. Blocklist prevents the most obvious shell-escape paths.
    if any(p.search(code) for p in _BLOCKED_PYTHON):
        return "Error: code blocked by safety filter."
    tmp = None
    try:
        fd, tmp = tempfile.mkstemp(suffix=".py")
        with open(fd, "w") as f:
            f.write(code)
        proc = subprocess.run(
            [sys.executable, tmp],
            capture_output=True,
            text=True,
            timeout=30,
            cwd=Path.home(),
        )
        output = (proc.stdout + proc.stderr).strip()
        if len(output) > 4000:
            output = output[:4000] + "\n[output truncated]"
        return output or "(no output)"
    except subprocess.TimeoutExpired:
        return "Error: script timed out after 30 seconds."
    except Exception as e:
        return f"Error: {e}"
    finally:
        if tmp:
            Path(tmp).unlink(missing_ok=True)


_SYSTEM_PATHS = ("/etc", "/usr", "/bin", "/sbin", "/boot", "/sys", "/proc", "/dev", "/lib")


def _read_file(path: str) -> str:
    p = Path(path).expanduser()
    if not p.exists():
        return f"Error: file not found: {p}"
    if not p.is_file():
        return f"Error: not a file: {p}"
    try:
        content = p.read_text(encoding="utf-8", errors="replace")
        if len(content) > 8000:
            content = content[:8000] + "\n[output truncated]"
        return content
    except Exception as e:
        return f"Error: {e}"


def _write_file(path: str, content: str) -> str:
    p = Path(path).expanduser()
    if any(str(p).startswith(sp) for sp in _SYSTEM_PATHS):
        return "Error: writing to system paths is blocked."
    try:
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(content, encoding="utf-8")
        return f"Written: {p}"
    except Exception as e:
        return f"Error: {e}"


register(Tool(
    name="read_file",
    description="Read the contents of a file. Use this to inspect source code, configs, logs, or any file the user is asking about — rather than asking them to paste it.",
    parameters={
        "type": "object",
        "properties": {
            "path": {"type": "string", "description": "Path to the file (~ supported)"},
        },
        "required": ["path"],
    },
    fn=_read_file,
))

register(Tool(
    name="write_file",
    description="Write content to a file, creating it if it doesn't exist. Use for saving code, configs, or any output the user wants persisted.",
    parameters={
        "type": "object",
        "properties": {
            "path": {"type": "string", "description": "Path to the file (~ supported)"},
            "content": {"type": "string", "description": "Content to write"},
        },
        "required": ["path", "content"],
    },
    fn=_write_file,
    requires_confirmation=True,
))


register(Tool(
    name="run_python",
    description="Execute a Python script and return its output. Use for calculations, data processing, parsing, or anything where running real code gives a more accurate answer than guessing. Print results explicitly — only stdout/stderr is returned.",
    parameters={
        "type": "object",
        "properties": {
            "code": {"type": "string", "description": "Python code to execute"},
        },
        "required": ["code"],
    },
    fn=_run_python,
    requires_confirmation=True,
))


def _web_search(query: str, max_results: int = 5) -> str:
    api_key = os.getenv("BRAVE_API_KEY", "")
    if not api_key:
        return "Error: BRAVE_API_KEY not set in .env"
    resp = requests.get(
        "https://api.search.brave.com/res/v1/web/search",
        headers={"X-Subscription-Token": api_key, "Accept": "application/json"},
        params={"q": query, "count": max_results},
        timeout=10,
    )
    resp.raise_for_status()
    results = resp.json().get("web", {}).get("results", [])
    if not results:
        return "No results found."
    return "\n\n".join(
        f"**{r['title']}**\n{r['url']}\n{r.get('description', '')}"
        for r in results
    )


register(Tool(
    name="web_search",
    description="Search the web for real-time information. Use this for current events, today's news, recent developments, or any question where up-to-date information matters. Do not guess at current information — search instead.",
    parameters={
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "Search query"},
        },
        "required": ["query"],
    },
    fn=_web_search,
))


def _fetch_url(url: str) -> str:
    if not url.startswith(("http://", "https://")):
        return "Error: only http/https URLs are supported."
    try:
        resp = requests.get(
            url,
            headers={"User-Agent": "Mozilla/5.0 (compatible; Nova/1.0)"},
            timeout=15,
        )
        resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        # Drop noise — scripts, styles, nav, ads, etc.
        for tag in soup(["script", "style", "nav", "header", "footer", "aside", "noscript"]):
            tag.decompose()
        text = soup.get_text(separator="\n")
        # Collapse blank lines
        text = re.sub(r"\n{3,}", "\n\n", text).strip()
        if len(text) > 8000:
            text = text[:8000] + "\n[output truncated]"
        return text or "(no readable content)"
    except requests.HTTPError as e:
        return f"Error: HTTP {e.response.status_code}"
    except Exception as e:
        return f"Error: {e}"


register(Tool(
    name="fetch_url",
    description=(
        "Fetch a URL and return its readable text content. Use this after web_search "
        "to get the full content of a page when the search snippet isn't enough — "
        "e.g. to find the latest video on a channel, read an article, or extract "
        "specific data from a page. Note: JS-rendered pages (SPAs) may return limited content."
    ),
    parameters={
        "type": "object",
        "properties": {
            "url": {"type": "string", "description": "The URL to fetch"},
        },
        "required": ["url"],
    },
    fn=_fetch_url,
))