#!/usr/bin/env python3
"""
Azure Security Agent - OpenAI Agents SDK + MCP tools
Inspired by https://github.com/ShawhinT/yt-mcp-agent
Uses MCPServerStdio to connect to the Azure Security MCP server and Groq for the model.
"""
import asyncio
import os
import sys
from pathlib import Path
from dotenv import load_dotenv
from openai import AsyncOpenAI
load_dotenv()
# OpenAI Agents SDK
from agents import Agent, Runner
from agents.mcp import MCPServerStdio
# Groq via OpenAI-compatible API (OpenAIChatCompletionsModel accepts custom client)
try:
from agents.models import OpenAIChatCompletionsModel
except ImportError:
from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
GROQ_BASE_URL = "https://api.groq.com/openai/v1"
GROQ_MODEL = "llama-3.3-70b-versatile"
async def main() -> None:
groq_api_key = os.getenv("GROQ_API_KEY")
subscription_id = os.getenv("AZURE_SUBSCRIPTION_ID")
if not groq_api_key:
print("Error: GROQ_API_KEY not found in .env")
return
if not subscription_id:
print("Warning: AZURE_SUBSCRIPTION_ID not found in .env")
subscription_id = input("Enter your Azure Subscription ID: ").strip()
project_dir = Path(__file__).resolve().parent
server_script = project_dir / "server.py"
# MCP server params: same Python (venv), run server.py, project cwd
# PYTHONUNBUFFERED=1 so server stdout is not buffered (client reads JSON-RPC over stdio)
env = {**os.environ, "PYTHONUNBUFFERED": "1"}
params = {
"command": sys.executable,
"args": ["-u", str(server_script)], # -u = unbuffered stdout (required for stdio MCP)
"env": env,
"cwd": str(project_dir),
}
async with MCPServerStdio(
name="Azure Security MCP Server",
params=params,
cache_tools_list=True,
client_session_timeout_seconds=30,
) as server:
# Optional: get system prompt from MCP server (like yt-mcp-agent)
try:
prompt_result = await server.get_prompt("system_prompt")
instructions = prompt_result.messages[0].content.text
except Exception:
instructions = (
"You are an Azure Cloud Security Analyst. Use the MCP tools to discover and analyze "
"Azure resources for security issues. Provide severity ratings and actionable recommendations."
)
# Inject subscription ID into instructions
instructions = instructions.strip() + f"\n\nCurrent Azure subscription ID: {subscription_id}"
# Groq as OpenAI-compatible model
groq_client = AsyncOpenAI(
base_url=GROQ_BASE_URL,
api_key=groq_api_key,
)
model = OpenAIChatCompletionsModel(
model=GROQ_MODEL,
openai_client=groq_client,
)
agent = Agent(
name="Azure Security Agent",
instructions=instructions,
mcp_servers=[server],
model=model,
)
print("=== Azure Security Agent (MCP tools + Groq) ===")
print("Type your security query, or 'exit' to quit.\n")
input_items = []
while True:
user_input = input("You: ").strip()
if user_input.lower() in ("exit", "quit", "bye"):
print("Goodbye!")
break
if not user_input:
continue
input_items.append({"content": user_input, "role": "user"})
# Run agent (tools are provided by MCP server)
result = await Runner.run(agent, input=input_items)
out = getattr(result, "final_output", None) if result else None
if out:
print(out)
input_items.append({"content": out, "role": "assistant"})
else:
print("(No output)")
print()
if __name__ == "__main__":
asyncio.run(main())
