refactoring threat modelling agent

.gitignore

@@ -1,7 +1,35 @@
 **/__pycache__/**
-.venv/*
+**/.venv/**
 **/policies/**
+**/memory/**
 .env
 .vscode
 .backup*
-example_sbom*
+example_sbom*
+reports
+.adk
+.venv
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
+
+node_modules
+dist
+dist-ssr
+*.local
+
+# Editor directories and files
+.vscode/*
+!.vscode/extensions.json
+.idea
+.DS_Store
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?

requirements.txt

@@ -30,3 +30,7 @@ pytz
 requests
 Flask
 Flask-Cors
+google-cloud-storage
+networkx
+matplotlib
+scipy

run_agent.py

@@ -11,8 +11,9 @@ async def main(instruction):
         return
 
     # The agent.run method is async
-    response = await secmind.run(instruction)
-    print(response)
+    response = secmind.run_async(instruction)
+    async for chunk in response:
+        print(chunk)
 
 if __name__ == "__main__":
     if len(sys.argv) > 1:

secmind/agent.py

@@ -28,7 +28,7 @@ class AgentConfig:
     """Configuration for the Security Mind agent."""
 
     NAME = "secmind"
-    MODEL = "gemini-2.5-flash"
+    MODEL = "gemini-2.5-pro"
     DESCRIPTION = "Master security agent that delegates tasks."
 
     # Task priorities (lower number = higher priority)
@@ -59,7 +59,7 @@ class AgentConfig:
         "code_reviews": "code_review_agent",
         "cloud_security": "cloud_compliance_agent",
         "cloud_compliance": "cloud_compliance_agent",
-        "threat_modelling": "app_sec_agent",
+        "threat_modelling": "threat_modeling_agent",
         "application_security": "app_sec_agent",
         "policy_governance": "policy_agent",
         "jira_tickets": "jira_agent",
@@ -107,7 +107,7 @@ def build_delegation_rules() -> str:
                 - Vulnerabilities and license checks → vuln_triage_agent
                 - Code reviews → code_review_agent
                 - Cloud security posture/compliance → cloud_compliance_agent
-                - Application security review/Threat Modelling → app_sec_agent
+                - Application security review/Threat Modelling → threat_modeling_agent
                 - Policy governance questions → policy_agent
                 - Jira tickets → jira_agent
 
@@ -184,7 +184,7 @@ def create_secmind_agent(
     Create and configure the Security Mind master agent.
 
     Args:
-        model: The AI model to use (default: gemini-2.5-flash)
+        model: The AI model to use (default: gemini-2.5-pro)
         sub_agents: List of sub-agents to delegate to (optional)
         validate: Whether to validate sub-agents before creating agent
 

secmind/memory_manager.py

@@ -185,6 +185,14 @@ def _setup_sqlite(self):
                 timestamp DATETIME
             )
         """)
+        # Table for storing public GCS buckets
+        cursor.execute("""
+            CREATE TABLE IF NOT EXISTS public_gcs_buckets (
+                project_id TEXT PRIMARY KEY,
+                buckets_json TEXT,
+                timestamp DATETIME
+            )
+        """)
         self.sqlite_conn.commit()
 
     def add_triage_result(self, cve_id: str, severity: str, recommendation: str, details: dict):
@@ -934,6 +942,48 @@ def get_gce_instance_details(self, project_id: str, instance_name: str, zone: st
         logger.info(f"GCE instance details cache miss for key: {cache_key}")
         return None
 
+    def add_public_gcs_buckets(self, project_id: str, buckets: dict):
+        """
+        Adds public GCS buckets to the cache.
+
+        Args:
+            project_id (str): The project ID of the query.
+            buckets (dict): The buckets to cache.
+        """
+        cursor = self.sqlite_conn.cursor()
+        timestamp = datetime.now(timezone.utc)
+
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO public_gcs_buckets (project_id, buckets_json, timestamp)
+            VALUES (?, ?, ?)
+            """,
+            (project_id, json.dumps(buckets), timestamp)
+        )
+        self.sqlite_conn.commit()
+        logger.info(f"Public GCS buckets cached for project: {project_id}")
+
+    def get_public_gcs_buckets(self, project_id: str) -> dict | None:
+        """
+        Retrieves cached public GCS buckets.
+
+        Args:
+            project_id (str): The project ID of the query.
+
+        Returns:
+            The cached buckets, or None if not found.
+        """
+        cursor = self.sqlite_conn.cursor()
+        cursor.execute("SELECT buckets_json FROM public_gcs_buckets WHERE project_id = ?", (project_id,))
+        row = cursor.fetchone()
+
+        if row:
+            logger.info(f"Public GCS buckets cache hit for project: {project_id}")
+            return json.loads(row['buckets_json'])
+
+        logger.info(f"Public GCS buckets cache miss for project: {project_id}")
+        return None
+
     def search_semantic_memory(self, query_text: str, n_results: int = 2) -> list[str]:
         """
         Searches the semantic memory in ChromaDB for contextually relevant information.

secmind/sub_agents/cloud_compliance_agent/agent.py

@@ -275,6 +275,39 @@ def check_access_keys(
     return response.to_dict()
 
 
+def check_public_gcs_buckets(cloud: str, project_id: str) -> dict:
+    """
+    Check for publicly accessible GCS buckets, with caching.
+
+    Args:
+        cloud: The cloud provider to use (e.g., "gcp")
+        project_id: GCP project ID (e.g., "my-project")
+
+    Returns:
+        Dictionary with a list of public buckets and a summary.
+
+    Example:
+        >>> check_public_gcs_buckets("gcp", "my-project")
+    """
+    logger.info(f"Tool called: check_public_gcs_buckets(cloud={cloud}, project_id={project_id})")
+
+    memory = MemoryManager()
+
+    # Check cache first
+    cached_buckets = memory.get_public_gcs_buckets(project_id)
+    if cached_buckets:
+        return cached_buckets
+
+    client = _get_client(cloud)
+    response = client.list_public_gcs_buckets(project_id=project_id)
+
+    # Add to cache
+    if response.status == "success":
+        memory.add_public_gcs_buckets(project_id, response.to_dict())
+
+    return response.to_dict()
+
+
 def generate_compliance_report(cloud: str, parent: str) -> dict:
     """
     Generates a comprehensive compliance report in HTML format.
@@ -309,6 +342,10 @@ def generate_compliance_report(cloud: str, parent: str) -> dict:
         if keys_result.get("status") == "success":
             all_data["access_keys"] = keys_result.get("data", {})
 
+        buckets_result = check_public_gcs_buckets(cloud, project_id)
+        if buckets_result.get("status") == "success":
+            all_data["public_gcs_buckets"] = buckets_result.get("data", {})
+
     if org_id:
         org_policies_result = check_org_policies(cloud, org_id)
         if org_policies_result.get("status") == "success":
@@ -367,14 +404,15 @@ def check_gcp_workload_security(instruction: str) -> dict:
     check_iam_recommendations,
     check_org_policies,
     check_access_keys,
+    check_public_gcs_buckets,
     generate_compliance_report,
     check_gcp_workload_security,
 ]
 
 # Create the agent instance
 cloud_compliance_agent = Agent(
     name=build_agent_name(),
-    model="gemini-2.5-flash",
+    model="gemini-2.5-pro",
     description=build_short_description(),
     instruction=build_agent_instructions(),
     tools=AGENT_TOOLS,

secmind/sub_agents/cloud_compliance_agent/clients/gcp.py

@@ -210,6 +210,9 @@ def _get_client(self, client_type: str) -> Any:
                     self._clients[client_type] = orgpolicy_v2.OrgPolicyClient()
                 elif client_type == "iam_admin":
                     self._clients[client_type] = iam_admin_v1.IAMClient()
+                elif client_type == "storage":
+                    from google.cloud import storage
+                    self._clients[client_type] = storage.Client()
                 else:
                     raise ValueError(f"Unknown client type: {client_type}")
 
@@ -408,6 +411,7 @@ def list_iam_recommendations(
                                 "op": op.get("op"),
                                 "value": op.get("value"),
                                 "originalValue": op.get("originalValue"),
+                                "pathFilters": op.get("pathFilters"),
                             }
                             details["operations"].append(op_details)
                 if not details["operations"]:
@@ -535,3 +539,48 @@ def list_service_account_keys(
             },
             message=f"Analyzed {len(keys)} service account keys"
         )
+
+    # ========================================================================
+    # GCS METHODS
+    # ========================================================================
+
+    @handle_gcp_errors
+    @retry_on_failure()
+    def list_public_gcs_buckets(self, project_id: str) -> APIResponse:
+        """
+        List publicly accessible GCS buckets.
+
+        Args:
+            project_id: GCP project ID
+
+        Returns:
+            APIResponse with list of public buckets and summary
+        """
+        logger.info(f"Listing public GCS buckets for project: {project_id}")
+
+        storage_client = self._get_client("storage")
+
+        public_buckets = []
+
+        for bucket in storage_client.list_buckets(project=project_id):
+            policy = bucket.get_iam_policy(requested_policy_version=3)
+
+            for binding in policy.bindings:
+                if "allUsers" in binding["members"] or "allAuthenticatedUsers" in binding["members"]:
+                    public_buckets.append({
+                        "name": bucket.name,
+                        "url": f"gs://{bucket.name}",
+                        "roles": binding["role"],
+                        "members": list(binding["members"]),
+                    })
+                    break  # Move to the next bucket once a public binding is found
+
+        logger.info(f"Found {len(public_buckets)} public GCS buckets")
+
+        return APIResponse.success(
+            data={
+                "public_buckets": public_buckets,
+                "summary": f"Found {len(public_buckets)} publicly accessible buckets."
+            },
+            message=f"Analyzed GCS buckets in project {project_id}"
+        )

secmind/sub_agents/cloud_compliance_agent/report_generator.py

@@ -4,6 +4,16 @@
 
 import datetime
 from typing import Dict, Any, Optional
+import json
+
+def _format_value(value: Any) -> str:
+    """Format a value for HTML display."""
+    if isinstance(value, dict):
+        return "<br>".join([f"&nbsp;&nbsp;<b>{k}:</b> {_format_value(v)}" for k, v in value.items()])
+    elif isinstance(value, list):
+        return "<br>".join([f"&nbsp;&nbsp;- {_format_value(i)}" for i in value])
+    else:
+        return str(value) if value is not None else "N/A"
 
 def generate_html_report(data: Dict[str, Any], parent: str, cloud: str) -> str:
     """
@@ -28,6 +38,8 @@ def generate_html_report(data: Dict[str, Any], parent: str, cloud: str) -> str:
     org_policies = data.get("org_policies", [])
     access_keys = data.get("access_keys", {})
     non_compliant_keys = access_keys.get("non_compliant", [])
+    public_buckets_data = data.get("public_gcs_buckets", {})
+    public_buckets = public_buckets_data.get("public_buckets", [])
 
     html = f"""
     <!DOCTYPE html>
@@ -135,20 +147,27 @@ def generate_html_report(data: Dict[str, Any], parent: str, cloud: str) -> str:
                     <div>Non-compliant Keys</div>
                     <div class="value">{len(non_compliant_keys)}</div>
                 </div>
+                 <div class="summary-item">
+                    <div>Public GCS Buckets</div>
+                    <div class="value">{len(public_buckets)}</div>
+                </div>
             </div>
 
             <h2>Security Posture Findings</h2>
             {"<table><tr><th>Severity</th><th>Category</th><th>Description</th><th>Resource</th></tr>" + "".join([f"<tr><td>{f['severity']}</td><td>{f['category']}</td><td>{f['description']}</td><td>{f['resource_name']}</td></tr>" for f in findings]) + "</table>" if findings else "<p>No security posture findings.</p>"}
 
             <h2>IAM Recommendations</h2>
-            {"<table><tr><th>Priority</th><th>Description</th><th>Recommender</th><th>Details</th></tr>" + "".join([f"<tr><td>{r['priority']}</td><td>{r['description']}</td><td>{r['recommender_subtype']}</td><td>{'<br>'.join([f'<b>Path:</b> {op.get("path", "N/A")}<br><b>Op:</b> {op.get("op", "N/A")}<br><b>Value:</b> {op.get("value", "N/A")}' for op in r.get('details', {}).get('operations', [])])}</td></tr>" for r in iam_recs]) + "</table>" if iam_recs else "<p>No IAM recommendations found.</p>"}
+            {"<table><tr><th>Priority</th><th>Description</th><th>Recommender</th><th>Details</th></tr>" + "".join([f"<tr><td>{r['priority']}</td><td>{r['description']}</td><td>{r['recommender_subtype']}</td><td>{'<br>'.join([f"<b>Resource:</b> {op.get('resource', 'N/A')}<br><b>Path:</b> {op.get('path', 'N/A')}<br><b>Path Filters:</b> {_format_value(op.get('pathFilters'))}<br><b>Value:</b> {_format_value(op.get('value'))}" for op in r.get('details', {}).get('operations', [])])}</td></tr>" for r in iam_recs]) + "</table>" if iam_recs else "<p>No IAM recommendations found.</p>"}
 
             <h2>Organization Policies</h2>
             {"<table><tr><th>Constraint</th><th>Rules</th></tr>" + "".join([f"<tr><td>{p['constraint']}</td><td>{str(p['rules'])}</td></tr>" for p in org_policies]) + "</table>" if org_policies else "<p>No organization policies found.</p>"}
 
             <h2>Non-Compliant Access Keys (&gt;{access_keys.get('max_age_days', 90)} days)</h2>
             {"<table><tr><th>Service Account</th><th>Key Name</th><th>Age (days)</th></tr>" + "".join([f"<tr><td>{k['service_account']}</td><td>{k['key_name']}</td><td>{k['age_days']}</td></tr>" for k in non_compliant_keys]) + "</table>" if non_compliant_keys else "<p>No non-compliant access keys found.</p>"}
 
+            <h2>Public GCS Buckets</h2>
+            {"<table><tr><th>Bucket Name</th><th>URL</th><th>Exposed Roles</th><th>Exposed Members</th></tr>" + "".join([f"<tr><td>{b['name']}</td><td>{b['url']}</td><td>{b['roles']}</td><td>{', '.join(b['members'])}</td></tr>" for b in public_buckets]) + "</table>" if public_buckets else "<p>No publicly accessible GCS buckets found.</p>"}
+
             <div class="footer">
                 <p>Generated by Security Mind AI</p>
             </div>

secmind/sub_agents/code_review_agent/agent.py

@@ -3,7 +3,7 @@
 import requests
 from typing import List
 from google.adk.agents import Agent
-import google.generativeai as genai
+import google.genai as genai
 
 from pydantic import BaseModel
 
@@ -37,11 +37,9 @@ def review_code(code_snippet: str) -> dict:
     api_key = os.getenv("GOOGLE_API_KEY")
     if not api_key:
         return {"issues": [], "fixes": [], "overall_comments": "Google API key not set."}
-
-    genai.configure(api_key=api_key)
-
+
     model = genai.GenerativeModel(
-        'gemini-2.5-flash', 
+        'gemini-2.5-pro', 
     )
 
     # Step 1: Auto-detect the language
@@ -140,7 +138,7 @@ def get_github_pr_diff(pr_url: str) -> str:
 # - Update the Agent configuration to include the new tool:
 code_review_agent = Agent(
     name="code_review_agent",
-    model="gemini-2.5-flash",
+    model="gemini-2.5-pro",
     description="Reviews code for security,code smells and best practices.And delegates to jira_agent if issues found.",
     instruction="""You are a code review agent. Your role is to review the provided code for security vulnerabilities, code smells, readability, and efficiency.
         The user can provide either a direct code snippet or a GitHub pull request URL. Your goal is to provide a thorough review and suggest necessary fixes.

secmind/sub_agents/jira_agent/agent.py

@@ -27,7 +27,7 @@ def create_jira_issue(project_key: str, summary: str, description: str, issue_ty
 
 jira_agent = Agent(
     name="jira_agent",
-    model="gemini-2.5-flash",
+    model="gemini-2.5-pro",
     description="Creates Jira issues from findings.",
     instruction="""
     Create issues using create_jira_issue with provided context.

secmind/sub_agents/policy_agent/agent.py

@@ -60,7 +60,7 @@ def list_policy_documents() -> dict:
 
 policy_agent = Agent(
     name="policy_agent",
-    model="gemini-2.5-flash",
+    model="gemini-2.5-pro",
     description="Reads policies from local files and answers to user questions from the policy files.",
     instruction="""
     Answer from local policies.