feat: Add Security Operations Agent with category filtering

itsdevansh · itsdevansh · commit eb59bf40cee5 · 2025-11-10T07:05:10.000-05:00
diff --git a/agents/frameworks/langchain/src/ibmi_agents/agents/__init__.py b/agents/frameworks/langchain/src/ibmi_agents/agents/__init__.py
@@ -7,8 +7,9 @@
 Available agents:
 - Performance Agent: System performance monitoring and analysis
 - SysAdmin Discovery Agent: High-level system discovery and summarization
-- SysAdmin Browse Agent: Detailed system browsing and exploration  
+- SysAdmin Browse Agent: Detailed system browsing and exploration
 - SysAdmin Search Agent: System search and lookup capabilities
+- Security Operations Agent: Security vulnerability assessment and remediation
 """
 
 from .ibmi_agents import (
@@ -17,6 +18,7 @@
     create_sysadmin_discovery_agent,
     create_sysadmin_browse_agent,
     create_sysadmin_search_agent,
+    create_security_ops_agent,
     chat_with_agent,
     list_available_agents,
     set_verbose_logging,
@@ -30,6 +32,7 @@
     "create_sysadmin_discovery_agent",
     "create_sysadmin_browse_agent",
     "create_sysadmin_search_agent",
+    "create_security_ops_agent",
     "chat_with_agent",
     "list_available_agents",
     "set_verbose_logging",
diff --git a/agents/frameworks/langchain/src/ibmi_agents/agents/ibmi_agents.py b/agents/frameworks/langchain/src/ibmi_agents/agents/ibmi_agents.py
@@ -18,7 +18,7 @@
 import os
 import json
 import getpass
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Optional
 from contextlib import _AsyncGeneratorContextManager, asynccontextmanager
 from langchain_ollama import ChatOllama
 from langchain_openai import ChatOpenAI
@@ -370,6 +370,92 @@ async def agent_session():
     
     return agent_session()
 
+async def create_security_ops_agent(
+    model_id: str = "gpt-oss:20b",
+    mcp_url: str = DEFAULT_MCP_URL,
+    transport: str = DEFAULT_TRANSPORT,
+    category: Optional[str] = None,
+    **kwargs
+):
+    """
+    Create IBM i Security Operations Agent.
+    
+    Args:
+        model_id: Model identifier (default: "gpt-oss:20b")
+        mcp_url: MCP server URL
+        transport: Transport type
+        category: Optional category filter for security tools. Options:
+                 - "vulnerability-assessment": Tools for identifying security vulnerabilities
+                 - "audit": Tools for auditing security configurations
+                 - "remediation": Tools for generating and executing security fixes
+                 - "user-management": Tools for managing user capabilities and permissions
+                 - None: Load all security tools (default)
+        **kwargs: Additional agent configuration options
+    
+    Returns an async context manager that yields (agent, session).
+    Usage: async with (await create_security_ops_agent()) as (agent, session): ...
+    """
+    client = get_mcp_client(mcp_url, transport)
+    
+    @asynccontextmanager
+    async def agent_session():
+        async with client.session("ibmi_tools") as session:
+            # Load security tools with optional category filtering
+            if category:
+                # Use annotation filtering to load tools by domain and category
+                tools = await load_filtered_mcp_tools(
+                    session,
+                    annotation_filters={
+                        "domain": "security",
+                        "category": category
+                    },
+                    debug=True
+                )
+                print(f"✅ Loaded {len(tools)} security operations tools (category: {category}) for Security Ops Agent")
+            else:
+                # Load all security tools by domain
+                tools = await load_filtered_mcp_tools(
+                    session,
+                    annotation_filters={"domain": "security"},
+                    debug=True
+                )
+                print(f"✅ Loaded {len(tools)} security operations tools for Security Ops Agent")
+            
+            system_message = """You are a specialized IBM i security operations assistant.
+You help administrators identify security vulnerabilities, audit system configurations, and remediate security issues.
+Your role is to:
+- Identify security vulnerabilities and misconfigurations
+- Assess user privileges and special authorities
+- Audit file and object permissions for *PUBLIC access
+- Detect potential attack vectors (triggers, impersonation, privilege escalation)
+- Generate remediation commands for security lockdown
+- Explain security risks in business terms
+- Provide actionable recommendations for hardening system security
+
+IMPORTANT SECURITY NOTES:
+- Always explain the security implications of findings
+- Distinguish between read-only assessment tools and destructive remediation tools
+- For remediation tools (like execute_impersonation_lockdown), warn users about system changes
+- Recommend testing remediation commands in development before production
+- Prioritize findings by severity (critical vulnerabilities first)
+
+Focus on helping administrators understand their security posture and take appropriate action to protect their IBM i systems."""
+            
+            llm = get_model(model_id)
+            
+            agent = create_agent(
+                model=llm,
+                tools=tools,
+                system_prompt=system_message,
+                checkpointer=get_shared_checkpointer(),
+                store=get_shared_store(),
+                name="IBM i Security Operations",
+                **kwargs
+            )
+            yield agent, session
+    
+    return agent_session()
+
 # -----------------------------------------------------------------------------
 # Agent Registry and Factory Pattern
 # -----------------------------------------------------------------------------
@@ -379,6 +465,7 @@ async def agent_session():
     "discovery": create_sysadmin_discovery_agent,
     "browse": create_sysadmin_browse_agent,
     "search": create_sysadmin_search_agent,
+    "security": create_security_ops_agent,
 }
 
 async def create_ibmi_agent(agent_type: str, **kwargs) -> _AsyncGeneratorContextManager[tuple[Any, Any], None]:
@@ -404,7 +491,8 @@ def list_available_agents() -> Dict[str, str]:
         "performance": "System performance monitoring and analysis",
         "discovery": "High-level system discovery and summarization",
         "browse": "Detailed system browsing and exploration",
-        "search": "System search and lookup capabilities"
+        "search": "System search and lookup capabilities",
+        "security": "Security vulnerability assessment and remediation"
     }
 
 def set_verbose_logging(enabled: bool):
@@ -609,7 +697,7 @@ async def test_agents():
             "gpt-oss:20b",           # Default Ollama model
             "ollama:llama3.1",         # Explicit Ollama model
             "openai:gpt-4o",  # OpenAI model
-            "anthropic:claude-3.7-sonnet"  # Anthropic model
+            "aanthropic:claude-3-7-sonnet-20250219"  # Anthropic model
         ]
         
         print("Available model options:")
diff --git a/agents/frameworks/langchain/src/ibmi_agents/agents/test_agents.py b/agents/frameworks/langchain/src/ibmi_agents/agents/test_agents.py
@@ -12,6 +12,7 @@
 import asyncio
 import sys
 import os
+from typing import Optional
 
 # Disable LangSmith tracing for tests
 os.environ["LANGCHAIN_TRACING_V2"] = "false"
@@ -30,19 +31,30 @@
     "performance": "What is my system status? Give me CPU and memory metrics.",
     "discovery": "Give me an overview of available system services.",
     "browse": "Show me services in the QSYS2 schema.",
-    "search": "Search for services related to system status."
+    "search": "Search for services related to system status.",
+    "security": "Check for user profiles vulnerable to impersonation attacks."
 }
 
-async def test_single_agent(agent_type: str, model_id: str = "gpt-oss:20b"):
+async def test_single_agent(agent_type: str, model_id: str = "gpt-oss:20b", category: Optional[str] = None):
     """Test a single agent with a sample query."""
     print(f"\n{'='*80}")
     print(f"Testing {agent_type.upper()} Agent")
+    if category:
+        print(f"Category Filter: {category}")
     print(f"{'='*80}\n")
     
     try:
         # Create agent context
         print(f"🔧 Creating {agent_type} agent with model {model_id}...")
-        ctx = await create_ibmi_agent(agent_type, model_id=model_id)
+        if category:
+            print(f"   Filtering by category: {category}")
+        
+        # Pass category parameter if provided (only for security agent)
+        kwargs = {"model_id": model_id}
+        if category and agent_type == "security":
+            kwargs["category"] = category
+        
+        ctx = await create_ibmi_agent(agent_type, **kwargs)
         
         async with ctx as (agent, session):
             print(f"✅ Agent created: {agent.name}\n")
@@ -111,16 +123,27 @@ async def test_all_agents(model_id: str = "gpt-oss:20b"):
     
     return all(results.values())
 
-async def interactive_mode(agent_type: str, model_id: str = "gpt-oss:20b"):
+async def interactive_mode(agent_type: str, model_id: str = "gpt-oss:20b", category: Optional[str] = None):
     """Interactive chat mode with a specific agent."""
     print(f"\n{'='*80}")
     print(f"Interactive Mode - {agent_type.upper()} Agent")
+    if category:
+        print(f"Category Filter: {category}")
     print(f"{'='*80}\n")
     
     try:
         # Create agent context
-        print(f"🔧 Initializing {agent_type} agent with {model_id}...\n")
-        ctx = await create_ibmi_agent(agent_type, model_id=model_id)
+        print(f"🔧 Initializing {agent_type} agent with {model_id}...")
+        if category:
+            print(f"   Filtering by category: {category}")
+        print()
+        
+        # Pass category parameter if provided (only for security agent)
+        kwargs = {"model_id": model_id}
+        if category and agent_type == "security":
+            kwargs["category"] = category
+        
+        ctx = await create_ibmi_agent(agent_type, **kwargs)
         
         async with ctx as (agent, session):
             print(f"✅ {agent.name} ready!\n")
@@ -171,19 +194,32 @@ async def interactive_mode(agent_type: str, model_id: str = "gpt-oss:20b"):
         import traceback
         traceback.print_exc()
 
-async def quick_test(model_id: str = "gpt-oss:20b"):
+async def quick_test(model_id: str = "gpt-oss:20b", category: Optional[str] = None, agent_filter: Optional[str] = None):
     """Quick test - just verify all agents can be created."""
     print("\n" + "="*80)
     print("Quick Agent Creation Test")
     print("="*80)
-    print(f"Model: {model_id}\n")
+    print(f"Model: {model_id}")
+    if category:
+        print(f"Category Filter: {category}")
+    print()
     
     results = {}
     
-    for agent_type in AVAILABLE_AGENTS.keys():
+    # Filter to specific agent if requested
+    agents_to_test = [agent_filter] if agent_filter else list(AVAILABLE_AGENTS.keys())
+    
+    for agent_type in agents_to_test:
         try:
             print(f"Creating {agent_type} agent...", end=" ")
-            ctx = await create_ibmi_agent(agent_type, model_id=model_id)
+            
+            # Pass category parameter if provided (only for security agent)
+            kwargs = {"model_id": model_id}
+            if category and agent_type == "security":
+                kwargs["category"] = category
+                print(f"(category: {category})...", end=" ")
+            
+            ctx = await create_ibmi_agent(agent_type, **kwargs)
             async with ctx as (agent, session):
                 print(f"✅ {agent.name}")
                 results[agent_type] = True
@@ -235,6 +271,12 @@ def main():
         help="Test specific agent type"
     )
     
+    parser.add_argument(
+        "--category",
+        choices=["vulnerability-assessment", "audit", "remediation", "user-management"],
+        help="Filter security agent tools by category (only applies to security agent)"
+    )
+    
     parser.add_argument(
         "--interactive",
         action="store_true",
@@ -278,15 +320,21 @@ def main():
     if args.interactive and not args.agent:
         parser.error("--interactive requires --agent to be specified")
     
+    if args.category and not args.agent:
+        parser.error("--category requires --agent to be specified")
+    
+    if args.category and args.agent != "security":
+        parser.error("--category can only be used with --agent security")
+    
     # Run appropriate test mode
     try:
         if args.quick:
-            success = asyncio.run(quick_test(args.model))
+            success = asyncio.run(quick_test(args.model, args.category, args.agent))
         elif args.interactive:
-            asyncio.run(interactive_mode(args.agent, args.model))
+            asyncio.run(interactive_mode(args.agent, args.model, args.category))
             success = True
         elif args.agent:
-            success = asyncio.run(test_single_agent(args.agent, args.model))
+            success = asyncio.run(test_single_agent(args.agent, args.model, args.category))
         else:
             success = asyncio.run(test_all_agents(args.model))
         
diff --git a/tools/security/security-ops.yaml b/tools/security/security-ops.yaml
@@ -1,5 +1,6 @@
 # Security Vulnerability Assessment and Remediation Tools
 # Purpose: Comprehensive security analysis tools for identifying vulnerabilities,
+# Based on SQL script by Scott Forstie
 # assessing user privileges, file permissions, and potential attack vectors on IBM i systems
 
 sources:
