Extending AEGIS

Three of the most common extension paths, with the actual file changes.

Adding a new tool

Say you want to add httprobe. Steps:

1. Add a wrapper

src/aegis/tools/wrappers/httprobe.py:

from aegis.tools.base import BaseTool
from aegis.tools.catalog import TOOL_CATALOG
from aegis.tools.models import ToolResult, Observation

class HttprobeWrapper(BaseTool):
    spec = TOOL_CATALOG["httprobe"]

    async def run(self, target: str, **kw) -> ToolResult:
        stdout, stderr, rc = await self._run(["-c", "50"], cwd=None)
        live_hosts = [l.strip() for l in stdout.splitlines() if l.strip()]
        return ToolResult(
            tool=self.spec.name, target=target, success=rc == 0,
            observations=[Observation(kind="http_alive", value=h) for h in live_hosts],
            raw_stdout=stdout, raw_stderr=stderr, returncode=rc,
        )

2. Register in the catalog

src/aegis/tools/catalog.py:

TOOL_CATALOG["httprobe"] = ToolSpec(
    name="httprobe",
    binary="httprobe",
    arch_install=InstallSpec(manager="go",
        package="github.com/tomnomnom/httprobe@latest"),
    timeout_sec=300,
)

3. Wire into the executor

src/aegis/tools/executor.py:

wrapper_map = {
    # …
    "httprobe": ("aegis.tools.wrappers.httprobe", "HttprobeWrapper"),
}

4. Surface as MCP tool

src/aegis/mcp_server/server.py:

@mcp.tool()
async def httprobe_check(target: str) -> dict:
    """Probe a list of hosts and return which respond on HTTP/HTTPS."""
    result = await ctx.executor.run_tool("httprobe", target)
    return {"alive": [o.value for o in result.observations], "rc": result.returncode}

5. Add a parser (optional, for richer output)

src/aegis/tools/parsers/httprobe.py — only needed if the tool emits structured output you want to lift into typed observations.

That’s it. aegis env install --missing will pick up the catalog entry; the next aegis run will let Claude call your tool.

Adding a verification probe

src/aegis/verify/probes.py:

async def _probe_xss_reflected(target: str, payload: str = '"<svg/onload=1>') -> dict:
    """Check whether a `q` query param is reflected without escaping."""
    async with httpx.AsyncClient() as client:
        r = await client.get(target, params={"q": payload}, timeout=10)
        verified = payload in r.text and "Content-Type" in r.headers
        return {
            "verified": verified,
            "evidence": r.text[:500] if verified else "no reflection",
            "confidence": "high" if verified else "low",
        }

PROBE_TYPES["xss_reflected"] = _probe_xss_reflected

Now aegis_verify(probe_type='xss_reflected', target='https://x/?q=test') routes to your probe.

Adding a custom backend

Your shop has its own LLM gateway? Subclass BaseBackend:

# src/aegis/llm/my_backend.py
from aegis.llm.base import BaseBackend
from aegis.llm.models import CompletionResult, Tier

class MyBackend(BaseBackend):
    supports_tool_use = True
    supports_prompt_caching = True
    name = "my-backend"

    async def complete(self, messages, tools, tier: Tier, **kw) -> CompletionResult:
        # call your gateway, return a CompletionResult
        ...

Register in src/aegis/llm/factory.py:

if preferred == "my-backend":
    from aegis.llm.my_backend import MyBackend
    return MyBackend()

Adding a phase

Phases are configured declaratively in src/aegis/methodology/phases.py. Add an enum entry, a PHASE_CONFIGS entry with objective + budget + prompt template + allowed tools, and the runner’s state machine picks it up.

Consider whether you actually need a new phase or whether the work belongs as an escalation rule in src/aegis/analysis/adaptive.py — most “new phases” are really targeted vulnerability runs that should be triggered by findings.