agent_deply/pam_deploy_graph/cli.py

"""Command line interface for the PAM deploy agent."""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict

from .agent import PamDeployAgent
from .checkpoint_store import load_agent_state, redact_mapping
from .interactive import run_interactive_chat
from .llm import build_llm_client
from .params_loader import load_params_file


def add_llm_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--llm-base-url")
    parser.add_argument("--llm-api-key")
    parser.add_argument("--llm-model")


def require_confirm(args: argparse.Namespace) -> None:
    if not getattr(args, "confirm", False):
        raise SystemExit("Refusing to execute actions without --confirm.")


def print_pause_payload(agent: PamDeployAgent, state) -> None:
    if state.pending_confirmation:
        print(json.dumps({"confirmation": agent.build_confirmation_request(state)}, ensure_ascii=False, indent=2))
    if state.checkpoint_path:
        print(json.dumps({"checkpoint": state.checkpoint_path}, ensure_ascii=False, indent=2))


def main() -> None:
    parser = argparse.ArgumentParser(prog="pam-deploy-agent")
    sub = parser.add_subparsers(dest="command", required=True)

    preview = sub.add_parser("preview")
    preview.add_argument("--config", required=True)
    preview.add_argument("--strategy", default="hybrid_node_mcp", choices=["hybrid_node_mcp", "script_only", "fake"])

    analyze = sub.add_parser("analyze")
    analyze.add_argument("--text", required=True)
    analyze.add_argument("--config")
    add_llm_args(analyze)

    chat = sub.add_parser("chat")
    chat.add_argument("--config", required=True)
    chat.add_argument("--strategy", default="fake", choices=["hybrid_node_mcp", "script_only", "fake"])
    chat.add_argument("--target-ip", action="append", default=[])
    chat.add_argument("--checkpoint")
    add_llm_args(chat)

    run = sub.add_parser("run-global")
    run.add_argument("--config", required=True)
    run.add_argument("--strategy", default="fake", choices=["hybrid_node_mcp", "script_only", "fake"])
    run.add_argument("--checkpoint")
    run.add_argument("--confirm", action="store_true")

    deploy = sub.add_parser("run-deploy")
    deploy.add_argument("--config", required=True)
    deploy.add_argument("--strategy", default="fake", choices=["hybrid_node_mcp", "script_only", "fake"])
    deploy.add_argument("--target-ip", action="append", default=[])
    deploy.add_argument("--checkpoint")
    deploy.add_argument("--confirm", action="store_true")

    resume = sub.add_parser("resume")
    resume.add_argument("--checkpoint", required=True)
    resume.add_argument("--confirm", action="store_true")

    confirm = sub.add_parser("confirm")
    confirm.add_argument("--checkpoint", required=True)
    confirm.add_argument("--decision", required=True, choices=["approve", "reject"])
    confirm.add_argument("--note", default="")
    confirm.add_argument("--confirm", action="store_true")

    args = parser.parse_args()
    params = load_params_file(args.config) if getattr(args, "config", None) else {}
    llm_client = None
    if args.command in ("analyze", "chat"):
        llm_client = build_llm_client(
            base_url=args.llm_base_url,
            api_key=args.llm_api_key,
            model=args.llm_model,
        )
    agent = PamDeployAgent(llm_client=llm_client)

    if args.command == "analyze":
        result = agent.analyze_request(args.text, params)
        payload = redact_mapping({key: asdict(value) for key, value in result.items()})
        print(json.dumps(payload, ensure_ascii=False, indent=2))
        return

    if args.command == "chat":
        run_interactive_chat(
            agent=agent,
            params=params,
            strategy=args.strategy,
            checkpoint_path=args.checkpoint,
            target_ips=args.target_ip,
        )
        return

    if args.command == "preview":
        print(agent.preview(params, args.strategy))
        return

    require_confirm(args)
    if args.command == "run-global":
        state = agent.create_state(
            params=params,
            execution_strategy=args.strategy,
            checkpoint_path=args.checkpoint,
        )
        state = agent.run_global_flow(state)
        print(json.dumps({"events": state.events}, ensure_ascii=False, indent=2))
        print_pause_payload(agent, state)
        return

    if args.command == "resume":
        state = load_agent_state(args.checkpoint)
        state.checkpoint_path = state.checkpoint_path or args.checkpoint
        state = agent.run_deploy_flow(state)
        print(agent.render_report(state))
        print_pause_payload(agent, state)
        return

    if args.command == "confirm":
        state = load_agent_state(args.checkpoint)
        state.checkpoint_path = state.checkpoint_path or args.checkpoint
        state = agent.confirm_pending(
            state,
            approved=args.decision == "approve",
            operator_note=args.note,
        )
        print(agent.render_report(state))
        print_pause_payload(agent, state)
        return

    state = agent.create_state(
        params=params,
        execution_strategy=args.strategy,
        checkpoint_path=args.checkpoint,
        target_ips=args.target_ip,
    )
    state = agent.run_deploy_flow(state)
    print(agent.render_report(state))
    print_pause_payload(agent, state)


if __name__ == "__main__":
    main()