Guided pentest, phase by phase

Assign agents per phase and drive a pentest one phase at a time with full control.

What you’ll build

A guided pentest where you stay in the loop. Unlike automatic mode, guided mode hands control back to you between phases, so you choose which agents run in each phase and review activity before continuing. You will assign agents to phases 1 through 3, stream each phase, and then turn the raw agent output into structured vulnerabilities.

Prerequisites

• A Rank account and an API token (see Authentication).
• Enough pentest agents available for each phase — a phase needs 3 to 4 agents on first assignment.

pip install rank-sdk
export RANK_API_KEY=rk_...

Steps

Run it

Save the following as guided_pentest.py, set RANK_API_KEY, then run python guided_pentest.py.

"""Drive a guided pentest phase by phase with the Rank SDK.

What this script does:
  1. Creates a chat session and a guided web pentest.
  2. Assigns three pentest agents to each of the three phases.
  3. Streams phases 1 -> 3 one at a time, reusing the same chat_id so the
     agents keep context across phases.
  4. Processes the raw operation responses into vulnerabilities with AI.
  5. Marks the pentest as finished and prints the vulnerability breakdown.

Guided mode hands control back to you between phases, which is useful when
you want to review activity or adjust scope before continuing.

The streaming loop applies the golden rule: the assistant's answer is built
ONLY from `content` events. Every `agent_event` (thinking, tool_call, ...) is
internal activity rendered as a timeline and never concatenated to the answer.

Run:
    pip install rank-sdk
    export RANK_API_KEY=rk_...
    python guided_pentest.py

Optional environment variables:
    RANK_TARGET_URL  Target URL to scan (default: https://example.com).
    RANK_MODEL       Model alias for vulnerability processing
                     (default: gemini-2.5-flash).
"""

from __future__ import annotations

import os

import rank

TARGET_URL = os.environ.get("RANK_TARGET_URL", "https://example.com")
MODEL_ALIAS = os.environ.get("RANK_MODEL", "gemini-2.5-flash")
PHASES = (1, 2, 3)
AGENTS_PER_PHASE = 3


def assign_phase_agents(client: rank.Rank, pentest_id: int, phase_id: int) -> None:
    """Pick the first agents available for a phase and assign them.

    Each phase requires a minimum of 3 agents (maximum 4) on first
    assignment.
    """
    available = client.agents.list(type="pentest", phase_id=phase_id)
    chosen = available.items[:AGENTS_PER_PHASE]
    if len(chosen) < AGENTS_PER_PHASE:
        raise SystemExit(
            f"Phase {phase_id} needs at least {AGENTS_PER_PHASE} agents, "
            f"only {len(chosen)} are available."
        )

    assignment = [
        {"agent_id": agent.id, "phase_id": phase_id, "execution_order": order}
        for order, agent in enumerate(chosen, start=1)
    ]
    result = client.pentests.agents.assign(pentest_id, agents=assignment)
    print(f"Phase {phase_id}: assigned {result.assigned_count} agents")


def stream_phase(client: rank.Rank, pentest_id: int, chat_id: int, phase_id: int) -> None:
    """Stream a single guided phase, separating the answer from activity."""
    print(f"\n========== Phase {phase_id} ==========\n")
    with client.ai.chat.stream(
        user_prompt=f"Execute phase {phase_id} against the configured targets",
        pentest_id=pentest_id,
        mode="guided",
        phase_id=phase_id,
        chat_id=chat_id,
    ) as stream:
        for event in stream:
            if event.type == "content":
                # The assistant's answer, token by token.
                print(event.content, end="", flush=True)
            elif event.is_agent_event:
                ev = event.agent_event
                if ev is None:
                    continue
                # A compact activity timeline; never part of the answer.
                if ev.event_type == rank.AgentEvent.TOOL_CALL:
                    print(f"\n  [tool ->] {ev.data.get('tool_name')}")
                elif ev.event_type == rank.AgentEvent.TOOL_RESULT:
                    print(f"\n  [tool <-] {ev.data.get('tool_name')} "
                          f"({ev.data.get('duration_ms', 0)}ms)")
                elif ev.event_type == rank.AgentEvent.AGENT_FINISHED:
                    print(f"\n  [agent #{ev.agent_id} done] "
                          f"{ev.data.get('stop_reason')}")
            elif event.type == "phase_complete":
                print(f"\n--- Phase {event.metadata.get('phase_id')} complete "
                      f"({event.metadata.get('progress', '?')}%) ---")
            elif event.type == "error":
                print(f"\nERROR: {event.error}")


def main() -> None:
    with rank.Rank() as client:
        chat = client.chats.create(name=f"Guided pentest - {TARGET_URL}")

        pentest = client.pentests.create(
            name="Guided pentest",
            url=TARGET_URL,
            type="web",
            mode="guided",
            methodology_id=1,
            assets=[
                {"asset_type": "url", "asset_value": TARGET_URL, "is_primary": True},
            ],
        )
        print(f"Created pentest #{pentest.id} (status: {pentest.status})")

        for phase_id in PHASES:
            assign_phase_agents(client, pentest.id, phase_id)

        for phase_id in PHASES:
            stream_phase(client, pentest.id, chat.id, phase_id)

        # Turn the agents' raw responses into structured vulnerabilities.
        print("\n\nProcessing vulnerabilities with AI...")
        processed = client.pentests.process_vulnerabilities(
            pentest.id, model_alias=MODEL_ALIAS,
        )
        print(f"  {processed.message}")
        print(f"  Created {len(processed.vulnerabilities_created)} vulnerabilities")

        finished = client.pentests.finish(pentest.id)
        print(f"\n{finished.message}")
        if finished.vulnerabilities is not None:
            summary = finished.vulnerabilities
            print(f"  Total      : {summary.total}")
            print(f"  By severity: {summary.by_severity}")
            if summary.risk_score is not None:
                print(f"  Risk score : {summary.risk_score.score} "
                      f"({summary.risk_score.level})")


if __name__ == "__main__":
    try:
        main()
    except rank.AuthenticationError:
        print("ERROR: invalid or missing API key. Set RANK_API_KEY.")
    except rank.APIError as exc:
        print(f"API error ({exc.status_code}): {exc.message}")

Once you have findings, head to Triage vulnerabilities and gate CI.