agent/tools.py

"""The three real tools the agent can call, plus their function declarations."""

import ast
import html
import ipaddress
import re
import socket
from urllib.parse import urljoin, urlparse

import httpx
import wikipediaapi
from asteval import Interpreter
from google.genai import types

from agent import config


# asteval's default symtable exposes open() and other file/introspection builtins. A calculator
# only needs arithmetic (evaluated via the AST, not the symtable) plus a few math names, so we
# default-deny: keep this allowlist and drop everything else from the symtable.
# Symtable allowlist: only these math names survive pruning; everything else (open, eval, getattr,
# pow, factorial, ...) is stripped. pow/factorial are excluded because they bypass asteval's
# safe_pow exponent guard and can build huge integers (CPU/memory DoS).
_CALC_ALLOWED_NAMES = frozenset({
    "abs", "round", "min", "max", "sum", "sqrt",
    "sin", "cos", "tan", "asin", "acos", "atan", "atan2",
    "log", "log10", "log2", "exp", "pi", "e",
    "floor", "ceil", "trunc", "int", "float", "fmod",
    "degrees", "radians", "hypot", "copysign",
})

# AST allowlist: numbers, arithmetic operators, and bare-name function calls only. Rejects
# string/list/tuple/dict literals (closing file-path access and the sequence-multiplication memory
# DoS, e.g. [0]*10**9 or (0,)*10**9) and attribute access — structurally, not by character.
_CALC_ALLOWED_NODES = (
    ast.Expression, ast.BinOp, ast.UnaryOp, ast.Call, ast.Name, ast.Load, ast.Constant,
    ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, ast.Mod, ast.Pow, ast.USub, ast.UAdd,
)


def _is_numeric_expression(expression: str) -> bool:
    """True only for pure arithmetic / math-function-call expressions."""
    try:
        tree = ast.parse(expression, mode="eval")
    except (SyntaxError, ValueError):
        return False
    for node in ast.walk(tree):
        if not isinstance(node, _CALC_ALLOWED_NODES):
            return False
        if isinstance(node, ast.Constant) and not isinstance(node.value, (int, float)):
            return False
        if isinstance(node, ast.Call) and not isinstance(node.func, ast.Name):
            return False
    return True


def calculator(expression: str) -> str:
    """Evaluate a safe arithmetic expression. Returns the result or 'ERROR: ...'."""
    if len(expression) > config.CALC_MAX_CHARS:
        return "ERROR: expression too long"
    if not _is_numeric_expression(expression):
        return "ERROR: only numeric arithmetic expressions are allowed"
    aeval = Interpreter()
    for name in list(aeval.symtable):
        if name not in _CALC_ALLOWED_NAMES:
            del aeval.symtable[name]
    try:
        result = aeval(expression)
        if aeval.error:
            first = aeval.error[0]
            msg = first.get_error()[1] if hasattr(first, "get_error") else str(first)
            return f"ERROR: {msg}"
        return str(result)
    except Exception as exc:
        return f"ERROR: {type(exc).__name__}: {exc}"


def strip_html(raw: str) -> str:
    """Remove scripts/styles/tags, unescape entities, collapse whitespace."""
    text = re.sub(r"(?is)<(script|style)\b.*?</\1>", " ", raw)
    text = re.sub(r"(?s)<[^>]+>", " ", text)
    text = html.unescape(text)
    return re.sub(r"\s+", " ", text).strip()


def _host_is_public(host, *, resolver=socket.getaddrinfo) -> bool:
    """True only if the host resolves and every resolved IP is a public address."""
    host = (host or "").rstrip(".").lower()
    if not host:
        return False
    try:
        infos = resolver(host, None)
    except Exception:
        return False
    if not infos:
        return False
    for info in infos:
        ip = ipaddress.ip_address(info[4][0])
        if (ip.is_private or ip.is_loopback or ip.is_link_local
                or ip.is_reserved or ip.is_multicast or ip.is_unspecified):
            return False
    # Residual risk: this validates the resolved IPs at check time; a DNS-rebinding
    # attacker with a low TTL could swap to a private IP before the connection. Fully
    # closing that needs IP-pinned connections — out of scope for this demo.
    return True


def web_get(url: str, *, client_factory=httpx.Client, host_check=_host_is_public,
            max_redirects=config.WEB_MAX_REDIRECTS) -> str:
    """Fetch http/https text from a public host, cleaned and length-capped, or 'ERROR: ...'.

    SSRF-hardened: rejects non-public/loopback/link-local hosts and does not auto-follow
    redirects (each hop's host is re-validated manually).
    """
    current = url
    for _ in range(max_redirects + 1):
        parsed = urlparse(current)
        if parsed.scheme not in ("http", "https"):
            return "ERROR: only http and https URLs are allowed"
        if not host_check(parsed.hostname):
            return "ERROR: refusing to fetch a non-public or unresolvable host"
        try:
            with client_factory(
                timeout=config.WEB_TIMEOUT,
                follow_redirects=False,
                headers={"User-Agent": config.WIKI_USER_AGENT},
            ) as client:
                resp = client.get(current)
        except Exception as exc:
            return f"ERROR: failed to fetch ({type(exc).__name__})"
        location = resp.headers.get("location") if getattr(resp, "headers", None) else None
        if getattr(resp, "is_redirect", False) and location:
            current = urljoin(current, location)
            continue
        try:
            resp.raise_for_status()
        except Exception as exc:
            return f"ERROR: failed to fetch ({type(exc).__name__})"
        return strip_html(resp.text)[: config.WEB_MAX_CHARS]
    return "ERROR: too many redirects"


def _build_wiki():
    return wikipediaapi.Wikipedia(user_agent=config.WIKI_USER_AGENT, language="en")


def wikipedia_search(query: str, *, wiki=None) -> str:
    """Return the first part of a Wikipedia page summary, or a NO_RESULTS/ERROR marker."""
    client = wiki if wiki is not None else _build_wiki()
    try:
        page = client.page(query)
        if not page.exists():
            return f"NO_RESULTS: no Wikipedia page found for '{query}'"
        summary = (page.summary or "").strip()
    except Exception as exc:
        return f"ERROR: wikipedia lookup failed ({type(exc).__name__})"
    if not summary:
        return f"NO_RESULTS: empty summary for '{query}'"
    return summary[: config.WIKI_SUMMARY_CHARS]


_CALCULATOR_DECL = types.FunctionDeclaration(
    name="calculator",
    description="Evaluate an arithmetic expression (+, -, *, /, %, **, parentheses). Use for all math.",
    parameters_json_schema={
        "type": "object",
        "properties": {"expression": {"type": "string", "description": "Arithmetic expression to evaluate."}},
        "required": ["expression"],
    },
)

_WIKIPEDIA_DECL = types.FunctionDeclaration(
    name="wikipedia_search",
    description="Get a short summary of the best-matching English Wikipedia page for a query.",
    parameters_json_schema={
        "type": "object",
        "properties": {"query": {"type": "string", "description": "Search term, e.g. a person, place, or event."}},
        "required": ["query"],
    },
)

_WEB_GET_DECL = types.FunctionDeclaration(
    name="web_get",
    description="Fetch readable text from an http/https URL.",
    parameters_json_schema={
        "type": "object",
        "properties": {"url": {"type": "string", "description": "An http or https URL."}},
        "required": ["url"],
    },
)

# name -> (callable, FunctionDeclaration)
TOOLS = {
    "calculator": (calculator, _CALCULATOR_DECL),
    "wikipedia_search": (wikipedia_search, _WIKIPEDIA_DECL),
    "web_get": (web_get, _WEB_GET_DECL),
}


def tool_callables():
    """Return {name: callable} for execution."""
    return {name: fn for name, (fn, _decl) in TOOLS.items()}


def declarations():
    """Return the genai tools list (one Tool holding all function declarations)."""
    return [types.Tool(function_declarations=[decl for _name, (_fn, decl) in TOOLS.items()])]