/pentestgpt_agent/tests/test_plan.py

from pathlib import Path

import pytest

from pentestgpt_agent.execution import ExecutionOutcome, ValidExecution
from pentestgpt_agent.memory import MemoryKernel, RunSnapshot, RunSpec
from pentestgpt_agent.plan import (
    PlanValidationError,
    SupervisorDecision,
    TaskKind,
    TaskProposal,
    TaskStatus,
    compile_plan,
)


def _complete_task(
    memory: MemoryKernel,
    snapshot: RunSnapshot,
    *,
    task_id: str,
    kind: TaskKind,
    target: str,
    evidence: str,
    basis_ids: tuple[str, ...] = (),
    depends_on: tuple[str, ...] = (),
) -> RunSnapshot:
    commit = memory.commit_plan(
        compile_plan(
            SupervisorDecision(
                base_revision=snapshot.revision,
                new_tasks=(
                    TaskProposal(
                        task_id,
                        kind,
                        target,
                        f"Execute {task_id}.",
                        f"Evidence for {task_id} is recorded.",
                        basis_ids=basis_ids,
                        depends_on=depends_on,
                    ),
                ),
                next_task_id=task_id,
                finish=False,
                summary=f"Select {task_id}.",
            ),
            snapshot,
        )
    )
    assert commit.lease is not None
    return memory.commit_execution(
        ValidExecution(
            run_id=snapshot.run_id,
            task_id=task_id,
            attempt_id=commit.lease.attempt_id,
            lease_revision=commit.lease.revision,
            trace_episode_id=f"executor-{task_id}",
            outcome=ExecutionOutcome.DONE,
            summary=f"Completed {task_id}.",
            observation=evidence,
            evidence_sequences=(1,),
        )
    )


def test_valid_supervisor_decision_atomically_adds_and_leases_one_task(
    tmp_path: Path,
) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    snapshot = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    decision = SupervisorDecision(
        base_revision=snapshot.revision,
        new_tasks=(
            TaskProposal(
                id="discover-http",
                kind=TaskKind.DISCOVER,
                target="http://127.0.0.1:8080",
                objective="Inspect the HTTP service and identify its exposed surface.",
                done_when="The response, technology clues, and reachable paths are recorded.",
            ),
        ),
        next_task_id="discover-http",
        finish=False,
        summary="Start with bounded HTTP discovery.",
    )

    commit = memory.commit_plan(compile_plan(decision, snapshot))
    recovered = MemoryKernel(tmp_path / "state.sqlite3").snapshot("run-1")

    assert commit.revision == 1
    assert commit.lease is not None
    assert commit.lease.task_id == "discover-http"
    assert recovered.revision == 1
    assert len(recovered.tasks) == 1
    assert recovered.tasks[0].status is TaskStatus.ACTIVE


def test_discovered_url_path_is_inside_its_authorized_origin(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    snapshot = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized local application.",
            allowed_targets=("http://127.0.0.1:51153",),
        )
    )
    decision = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal(
                "enumerate-apply",
                TaskKind.ENUMERATE,
                "http://127.0.0.1:51153/apply",
                "Enumerate the discovered application form.",
                "The form surface is recorded.",
            ),
        ),
        next_task_id="enumerate-apply",
        finish=False,
        summary="Follow the discovered same-origin path.",
    )

    compiled = compile_plan(decision, snapshot)

    assert compiled.new_tasks[0].target == "http://127.0.0.1:51153/apply"


@pytest.mark.parametrize(
    "target",
    [
        "http://127.0.0.1:51154/apply",
        "http://127.0.0.1:51153@outside.example/apply",
    ],
)
def test_url_scope_rejects_targets_outside_the_authorized_origin(
    tmp_path: Path,
    target: str,
) -> None:
    snapshot = MemoryKernel(tmp_path / "state.sqlite3").create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized local application.",
            allowed_targets=("http://127.0.0.1:51153",),
        )
    )
    decision = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal(
                "outside",
                TaskKind.DISCOVER,
                target,
                "Inspect a target.",
                "The response is recorded.",
            ),
        ),
        next_task_id="outside",
        finish=False,
        summary="Inspect a target.",
    )

    with pytest.raises(PlanValidationError, match="outside scope"):
        compile_plan(decision, snapshot)


@pytest.mark.parametrize(
    "target",
    [
        "http://127.0.0.1:51153/application",
        "http://127.0.0.1:51153/app/../admin",
        "http://127.0.0.1:51153/app/%2e%2e/admin",
        "http://127.0.0.1:51153/app/%252e%252e/admin",
        "http://127.0.0.1:51153/app%2f..%2fadmin",
        "http://127.0.0.1:51153/app#outside",
    ],
)
def test_url_subtree_scope_rejects_prefix_and_ambiguous_path_escapes(
    tmp_path: Path,
    target: str,
) -> None:
    snapshot = MemoryKernel(tmp_path / "state.sqlite3").create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized application subtree.",
            allowed_targets=("http://127.0.0.1:51153/app",),
        )
    )
    decision = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal(
                "candidate",
                TaskKind.DISCOVER,
                target,
                "Inspect a target.",
                "The response is recorded.",
            ),
        ),
        next_task_id="candidate",
        finish=False,
        summary="Inspect a target.",
    )

    with pytest.raises(PlanValidationError, match="outside scope"):
        compile_plan(decision, snapshot)


def test_supervisor_cannot_create_a_speculative_task_backlog(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    snapshot = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    target = "http://127.0.0.1:8080"
    decision = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal("first", TaskKind.DISCOVER, target, "Inspect HTTP.", "HTTP is recorded."),
            TaskProposal("second", TaskKind.TEST, target, "Test HTTP.", "HTTP is tested."),
        ),
        next_task_id="first",
        finish=False,
        summary="Create a broad backlog.",
    )

    with pytest.raises(PlanValidationError, match="at most one new task"):
        compile_plan(decision, snapshot)

    assert memory.snapshot("run-1") == snapshot


def test_new_task_must_be_selected_immediately_instead_of_deferred(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    first = memory.commit_plan(
        compile_plan(
            SupervisorDecision(
                base_revision=0,
                new_tasks=(
                    TaskProposal(
                        "discover",
                        TaskKind.DISCOVER,
                        "http://127.0.0.1:8080",
                        "Inspect the target.",
                        "The surface is recorded.",
                    ),
                ),
                next_task_id="discover",
                finish=False,
                summary="Start discovery.",
            ),
            initial,
        )
    )
    assert first.lease is not None
    progressed = memory.commit_execution(
        ValidExecution(
            run_id="run-1",
            task_id="discover",
            attempt_id=first.lease.attempt_id,
            lease_revision=first.lease.revision,
            trace_episode_id="executor-discover",
            outcome=ExecutionOutcome.PROGRESS,
            summary="One request remains.",
            observation="partial surface",
            evidence_sequences=(1,),
        )
    )

    with pytest.raises(PlanValidationError, match="new task must be selected immediately"):
        compile_plan(
            SupervisorDecision(
                base_revision=progressed.revision,
                new_tasks=(
                    TaskProposal(
                        "test-input",
                        TaskKind.TEST,
                        "http://127.0.0.1:8080/input",
                        "Confirm the named hypothesis.",
                        "The behavior is recorded.",
                        depends_on=("discover",),
                    ),
                ),
                next_task_id="discover",
                finish=False,
                summary="Do not create deferred work.",
            ),
            progressed,
        )


def test_unknown_basis_observation_is_rejected(tmp_path: Path) -> None:
    snapshot = MemoryKernel(tmp_path / "state.sqlite3").create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    decision = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal(
                "test-http",
                TaskKind.TEST,
                "http://127.0.0.1:8080",
                "Test the discovered input.",
                "The behavior is recorded.",
                basis_ids=("fabricated:observation",),
            ),
        ),
        next_task_id="test-http",
        finish=False,
        summary="Test unsupported evidence.",
    )

    with pytest.raises(PlanValidationError, match="unknown basis observations"):
        compile_plan(decision, snapshot)


def test_basis_producer_must_be_a_dependency(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    discovery = memory.commit_plan(
        compile_plan(
            SupervisorDecision(
                base_revision=0,
                new_tasks=(
                    TaskProposal(
                        "discover-http",
                        TaskKind.DISCOVER,
                        "http://127.0.0.1:8080",
                        "Inspect HTTP.",
                        "HTTP is recorded.",
                    ),
                ),
                next_task_id="discover-http",
                finish=False,
                summary="Discover HTTP.",
            ),
            initial,
        )
    )
    assert discovery.lease is not None
    discovered = memory.commit_execution(
        ValidExecution(
            run_id="run-1",
            task_id="discover-http",
            attempt_id=discovery.lease.attempt_id,
            lease_revision=discovery.lease.revision,
            trace_episode_id="executor-discover",
            outcome=ExecutionOutcome.DONE,
            summary="HTTP discovered.",
            observation="200 OK",
            evidence_sequences=(1,),
        )
    )
    observation_id = discovered.observations[0].id
    decision = SupervisorDecision(
        base_revision=discovered.revision,
        new_tasks=(
            TaskProposal(
                "test-http",
                TaskKind.TEST,
                "http://127.0.0.1:8080",
                "Test the discovered input.",
                "The behavior is recorded.",
                basis_ids=(observation_id,),
                depends_on=(),
            ),
        ),
        next_task_id="test-http",
        finish=False,
        summary="Test HTTP.",
    )

    with pytest.raises(PlanValidationError, match="basis-producing tasks missing"):
        compile_plan(decision, discovered)


def test_exploit_requires_the_latest_test_evidence_on_the_same_surface(
    tmp_path: Path,
) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    target = "http://127.0.0.1:8080/submit"
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Capture the flag.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    reflected = _complete_task(
        memory,
        initial,
        task_id="test-reflection",
        kind=TaskKind.TEST,
        target=target,
        evidence="Thank you MARKER",
    )
    reflection_observation = reflected.observations[-1]
    confirmed = _complete_task(
        memory,
        reflected,
        task_id="test-ssti",
        kind=TaskKind.TEST,
        target=target,
        evidence="Thank you 49",
        basis_ids=(reflection_observation.id,),
        depends_on=("test-reflection",),
    )
    ssti_observation = confirmed.observations[-1]
    stale_decision = SupervisorDecision(
        base_revision=confirmed.revision,
        new_tasks=(
            TaskProposal(
                "exploit-ssti",
                TaskKind.EXPLOIT,
                target,
                "Use the confirmed SSTI to capture the flag.",
                "The exact flag is recorded.",
                basis_ids=(reflection_observation.id,),
                depends_on=("test-reflection",),
            ),
        ),
        next_task_id="exploit-ssti",
        finish=False,
        summary="Exploit stale reflection evidence.",
    )

    with pytest.raises(PlanValidationError, match="must cite latest completed test observation"):
        compile_plan(stale_decision, confirmed)

    current_decision = SupervisorDecision(
        base_revision=confirmed.revision,
        new_tasks=(
            TaskProposal(
                "exploit-ssti",
                TaskKind.EXPLOIT,
                target,
                "Use the confirmed SSTI to capture the flag.",
                "The exact flag is recorded.",
                basis_ids=(ssti_observation.id,),
                depends_on=("test-ssti",),
            ),
        ),
        next_task_id="exploit-ssti",
        finish=False,
        summary="Exploit current SSTI evidence.",
    )

    compiled = compile_plan(current_decision, confirmed)

    assert compiled.new_tasks[0].basis_ids == (ssti_observation.id,)


def test_exploit_provenance_requires_the_exact_test_target_including_query(
    tmp_path: Path,
) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Capture the flag.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    tested = _complete_task(
        memory,
        initial,
        task_id="test-mode",
        kind=TaskKind.TEST,
        target="http://127.0.0.1:8080/submit?mode=test",
        evidence="template result: 49",
    )
    observation = tested.observations[-1]

    with pytest.raises(PlanValidationError, match="completed test basis"):
        compile_plan(
            SupervisorDecision(
                base_revision=tested.revision,
                new_tasks=(
                    TaskProposal(
                        "exploit-mode",
                        TaskKind.EXPLOIT,
                        "http://127.0.0.1:8080/submit?mode=admin",
                        "Use the confirmed primitive.",
                        "The goal artifact is recorded.",
                        basis_ids=(observation.id,),
                        depends_on=("test-mode",),
                    ),
                ),
                next_task_id="exploit-mode",
                finish=False,
                summary="Do not cross target queries.",
            ),
            tested,
        )


def test_supervisor_cannot_finish_while_work_is_open(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    target = "http://127.0.0.1:8080"
    plan = compile_plan(
        SupervisorDecision(
            base_revision=0,
            new_tasks=(
                TaskProposal(
                    "discover-http",
                    TaskKind.DISCOVER,
                    target,
                    "Inspect HTTP.",
                    "HTTP is recorded.",
                ),
            ),
            next_task_id="discover-http",
            finish=False,
            summary="Inspect the target.",
        ),
        initial,
    )
    memory.commit_plan(plan)
    active = memory.snapshot("run-1")

    with pytest.raises(PlanValidationError, match="open tasks"):
        compile_plan(
            SupervisorDecision(
                base_revision=active.revision,
                new_tasks=(),
                next_task_id=None,
                finish=True,
                summary="Finish prematurely.",
            ),
            active,
        )


def test_finish_requires_an_explicit_existing_canonical_basis(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    completed = _complete_task(
        memory,
        initial,
        task_id="discover-http",
        kind=TaskKind.DISCOVER,
        target="http://127.0.0.1:8080",
        evidence="200 OK",
    )
    observation_id = completed.observations[0].id

    with pytest.raises(PlanValidationError, match="finish requires canonical basis"):
        compile_plan(
            SupervisorDecision(
                base_revision=completed.revision,
                new_tasks=(),
                next_task_id=None,
                finish=True,
                summary="Finish without naming evidence.",
            ),
            completed,
        )

    with pytest.raises(PlanValidationError, match="unknown finish basis"):
        compile_plan(
            SupervisorDecision(
                base_revision=completed.revision,
                new_tasks=(),
                next_task_id=None,
                finish=True,
                summary="Finish with fabricated evidence.",
                finish_basis_ids=("fabricated:observation",),
            ),
            completed,
        )

    compiled = compile_plan(
        SupervisorDecision(
            base_revision=completed.revision,
            new_tasks=(),
            next_task_id=None,
            finish=True,
            summary="Finish from the recorded response.",
            finish_basis_ids=(observation_id,),
        ),
        completed,
    )

    assert compiled.finish_basis_ids == (observation_id,)


def test_finish_basis_must_come_from_a_completed_task(tmp_path: Path) -> None:
    memory = MemoryKernel(tmp_path / "state.sqlite3")
    initial = memory.create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
            max_attempts_per_task=1,
        )
    )
    commit = memory.commit_plan(
        compile_plan(
            SupervisorDecision(
                base_revision=0,
                new_tasks=(
                    TaskProposal(
                        "test-input",
                        TaskKind.TEST,
                        "http://127.0.0.1:8080",
                        "Confirm whether the input transforms data.",
                        "The behavior is recorded.",
                    ),
                ),
                next_task_id="test-input",
                finish=False,
                summary="Test the input.",
            ),
            initial,
        )
    )
    assert commit.lease is not None
    exhausted = memory.commit_execution(
        ValidExecution(
            run_id="run-1",
            task_id="test-input",
            attempt_id=commit.lease.attempt_id,
            lease_revision=commit.lease.revision,
            trace_episode_id="executor-test",
            outcome=ExecutionOutcome.PROGRESS,
            summary="The premise remains unresolved.",
            observation="partial response",
            evidence_sequences=(1,),
        )
    )
    assert exhausted.tasks[0].status is TaskStatus.FAILED

    with pytest.raises(PlanValidationError, match="completed tasks"):
        compile_plan(
            SupervisorDecision(
                base_revision=exhausted.revision,
                new_tasks=(),
                next_task_id=None,
                finish=True,
                summary="Do not finish from partial progress.",
                finish_basis_ids=(exhausted.observations[0].id,),
            ),
            exhausted,
        )


@pytest.mark.parametrize(
    "objective",
    [
        "Confirm injection, then read the flag.",
        "Test the parameter and capture the flag if it works.",
        "Probe the input and dump credentials from the server.",
        "Try the hypothesis and gain a reverse shell.",
    ],
)
def test_test_task_cannot_embed_exploitation_as_its_objective(
    tmp_path: Path,
    objective: str,
) -> None:
    snapshot = MemoryKernel(tmp_path / "state.sqlite3").create_run(
        RunSpec(
            run_id="run-1",
            goal="Capture the benchmark flag.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )

    with pytest.raises(PlanValidationError, match="TEST task crosses into exploitation"):
        compile_plan(
            SupervisorDecision(
                base_revision=0,
                new_tasks=(
                    TaskProposal(
                        "test-input",
                        TaskKind.TEST,
                        "http://127.0.0.1:8080/check",
                        objective,
                        "The injection behavior is recorded.",
                    ),
                ),
                next_task_id="test-input",
                finish=False,
                summary="Test the input.",
            ),
            snapshot,
        )


def test_plan_text_and_reference_collections_are_bounded(tmp_path: Path) -> None:
    snapshot = MemoryKernel(tmp_path / "state.sqlite3").create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    oversized_objective = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal(
                "discover",
                TaskKind.DISCOVER,
                "http://127.0.0.1:8080",
                "x" * 2_001,
                "The surface is recorded.",
            ),
        ),
        next_task_id="discover",
        finish=False,
        summary="Begin discovery.",
    )
    too_many_references = SupervisorDecision(
        base_revision=0,
        new_tasks=(
            TaskProposal(
                "discover",
                TaskKind.DISCOVER,
                "http://127.0.0.1:8080",
                "Inspect the target.",
                "The surface is recorded.",
                basis_ids=tuple(f"obs-{index}" for index in range(9)),
            ),
        ),
        next_task_id="discover",
        finish=False,
        summary="Begin discovery.",
    )

    with pytest.raises(PlanValidationError, match="objective exceeds"):
        compile_plan(oversized_objective, snapshot)
    with pytest.raises(PlanValidationError, match="at most 8 basis"):
        compile_plan(too_many_references, snapshot)