面向企业级复杂环境的 智能体网络攻击能力评估

from __future__ import annotations

import asyncio
import os
import tarfile
import tempfile
from pathlib import Path

from agentcyberrange import AgentOutputPath, ArenaHandle, is_post_challenge, run_agent

CODEX_COMMAND = (
    "codex",
    "exec",
    "--ephemeral",
    "--skip-git-repo-check",
    "--sandbox",
    "workspace-write",
    "--config",
    "sandbox_workspace_write.network_access=true",
    "-",
)


def archive_agent_output(workspace: str, arena_id: str) -> AgentOutputPath | None:
    final_answer = Path(workspace) / "final_answer"
    if not final_answer.is_dir():
        return None

    output_dir = Path.cwd() / ".agentcyberrange" / "agent-output"
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / f"{arena_id}.tar.gz"
    with tarfile.open(output_path, "w:gz") as archive:
        archive.add(final_answer, arcname="final_answer")
    return output_path


async def stop_codex(process: asyncio.subprocess.Process) -> None:
    """Terminate Codex, escalating to kill after five seconds."""
    if process.returncode is not None:
        return
    process.terminate()
    try:
        await asyncio.wait_for(process.wait(), timeout=5)
    except TimeoutError:
        process.kill()
        await process.wait()


async def call_codex(arena: ArenaHandle) -> AgentOutputPath | None:
    """Run Codex non-interactively against one authorized Arena."""
    codex_environment = os.environ.copy()
    codex_environment.pop("AGENTCYBERRANGE_TASK_TOKEN", None)
    codex_environment.pop("AGENTCYBERRANGE_BASE_URL", None)

    with tempfile.TemporaryDirectory(prefix="agentcyberrange-") as workspace:
        process = await asyncio.create_subprocess_exec(
            *CODEX_COMMAND,
            cwd=workspace,
            env=codex_environment,
            stdin=asyncio.subprocess.PIPE,
        )
        try:
            await process.communicate(arena.task_prompt.encode())
        except asyncio.CancelledError:
            await stop_codex(process)
            if not arena.cancel_event.is_set():
                raise
        else:
            if process.returncode != 0:
                raise RuntimeError(f"Codex exited with status {process.returncode}")
        if is_post_challenge(arena.challenge_id):
            return None
        return archive_agent_output(workspace, arena.arena_id)


async def main() -> None:
    result = await run_agent(call_codex)
    print(result.model_dump_json(indent=2))


if __name__ == "__main__":
    asyncio.run(main())