Adding Modules

Adding Modules

Most attack modules follow a four-step pattern. Use aix/modules/inject.py as the reference implementation.

---

Step 1 — Create the module file

# aix/modules/mymodule.py

from aix.core.scanner import BaseScanner
from aix.utils.run import run_scanner
from aix.utils.cli import standard_options
import click


class MyScanner(BaseScanner):
    def __init__(self, target, **kwargs):
        super().__init__(target, **kwargs)
        self.module_name = "mymodule"
        self.console_color = "cyan"
        self.default_payloads = self.load_payloads("mymodule.json")

    async def run(self):
        await self._run_payload_scan(
            payloads=self.default_payloads,
            progress_description="Testing mymodule...",
            finding_title_prefix="MyModule",
        )


def run(target, **kwargs):
    run_scanner(MyScanner, target, **kwargs)

---

Step 2 — Create the payload file

Create aix/payloads/mymodule.json:

[
    {
        "name": "basic_test",
        "payload": "Your attack payload here",
        "indicators": ["keyword1", "keyword2"],
        "severity": "HIGH",
        "level": 1,
        "risk": 1,
        "owasp": ["LLM01"],
        "atlas": ["AML.T0048"]
    }
]

See Payload Schema for all fields.

---

Step 3 — Register in __init__.py

# aix/modules/__init__.py
from aix.modules.mymodule import MyScanner

---

Step 4 — Add CLI command

# aix/cli.py
from aix.modules import mymodule

@main.command()
@standard_options
@click.pass_context
def mymodule_cmd(ctx, target, **kwargs):
    """Short description shown in aix --help."""
    mymodule.run(target, **kwargs)

---

Custom scan logic

If your module can't use _run_payload_scan() (e.g. it needs a custom loop, multiple phases, or non-standard output), inherit BaseScanner and override run() directly. See aix/modules/recon.py for a full example.

---

Key BaseScanner methods

Method	Purpose
load_payloads(filename)	Load from aix/payloads/, apply --level/--risk filters
_run_payload_scan(payloads, ...)	Template scan loop — connector setup, payload iteration, finding, DB write
check_success(response, indicators, payload, technique)	AI eval or keyword match; sets self.last_eval_reason
gather_context(connector)	AI context probe (if ai_engine.enable_context)
_print(status, msg)	Rich-formatted output. Status: info, success, warning, error, blocked, detail
_on_finding(finding)	Hook called after each finding — override for custom post-processing
_on_scan_complete()	Hook called after the scan loop completes

---

Conventions

• Guard all console output with if not self.quiet — modules run silently inside chains
• Use self._print() instead of print() or console.print()
• Add new CLI options via kwargs.get("option_name", default) in __init__ — they're forwarded automatically through run_scanner → Scanner.__init__ → super().__init__
• Severity.INFO.value == "info" — db.add_result() takes strings, Finding() takes enum values