Refactor: Extract helper methods and improve code quality

Copilot · sgerlach · Copilot · commit fcef6cbc4f45 · 2025-10-27T21:23:58.000Z
Co-authored-by: sgerlach &lt;4514084+sgerlach@users.noreply.github.com&gt;
diff --git a/stackhawk_mcp/server.py b/stackhawk_mcp/server.py
@@ -2192,24 +2192,12 @@ async def _analyze_sensitive_data_trends(self, org_id: str, analysis_period: str
     async def _check_repository_attack_surface(self, repo_name: str = None, org_id: str = None, include_vulnerabilities: bool = True, include_apps: bool = True, **kwargs) -> Dict[str, Any]:
         """Check if a repository name exists in StackHawk attack surface and get security information"""
         try:
-            # Get repo name from current directory if not provided
-            if not repo_name:
-                repo_name = os.path.basename(os.getcwd())
-            
-            # Get org_id if not provided
-            if not org_id:
-                user_info = await self.client.get_user_info()
-                org_id = user_info["user"]["external"]["organizations"][0]["organization"]["id"]
+            # Get repo name and org_id using helper methods
+            repo_name = self._get_current_repository_name(repo_name)
+            org_id = await self._get_organization_id(org_id)
 
-            # List all repositories in the organization
-            repos_response = await self.client.list_repositories(org_id, pageSize=1000)
-            repositories = repos_response.get("repositories", [])
-            
-            # Find matching repositories (case-insensitive)
-            matching_repos = [
-                repo for repo in repositories
-                if repo.get("name", "").lower() == repo_name.lower()
-            ]
+            # Find matching repositories
+            matching_repos = await self._find_matching_repositories(repo_name, org_id)
             
             result = {
                 "repository_name": repo_name,
@@ -2276,24 +2264,12 @@ async def _check_repository_attack_surface(self, repo_name: str = None, org_id:
     async def _check_repository_sensitive_data(self, repo_name: str = None, org_id: str = None, data_type_filter: str = "All", include_remediation: bool = True, **kwargs) -> Dict[str, Any]:
         """Check if a repository has sensitive data findings in StackHawk"""
         try:
-            # Get repo name from current directory if not provided
-            if not repo_name:
-                repo_name = os.path.basename(os.getcwd())
-            
-            # Get org_id if not provided
-            if not org_id:
-                user_info = await self.client.get_user_info()
-                org_id = user_info["user"]["external"]["organizations"][0]["organization"]["id"]
+            # Get repo name and org_id using helper methods
+            repo_name = self._get_current_repository_name(repo_name)
+            org_id = await self._get_organization_id(org_id)
 
-            # List all repositories in the organization
-            repos_response = await self.client.list_repositories(org_id, pageSize=1000)
-            repositories = repos_response.get("repositories", [])
-            
-            # Find matching repositories (case-insensitive)
-            matching_repos = [
-                repo for repo in repositories
-                if repo.get("name", "").lower() == repo_name.lower()
-            ]
+            # Find matching repositories
+            matching_repos = await self._find_matching_repositories(repo_name, org_id)
             
             result = {
                 "repository_name": repo_name,
@@ -2354,10 +2330,8 @@ async def _check_repository_sensitive_data(self, repo_name: str = None, org_id:
     async def _list_application_repository_connections(self, org_id: str = None, include_repo_details: bool = True, include_app_details: bool = True, filter_connected_only: bool = False, **kwargs) -> Dict[str, Any]:
         """List connections between StackHawk applications and code repositories"""
         try:
-            # Get org_id if not provided
-            if not org_id:
-                user_info = await self.client.get_user_info()
-                org_id = user_info["user"]["external"]["organizations"][0]["organization"]["id"]
+            # Get org_id using helper method
+            org_id = await self._get_organization_id(org_id)
 
             # Get all applications and repositories
             apps_response = await self.client.list_applications(org_id, pageSize=1000)
@@ -2463,6 +2437,36 @@ async def _list_application_repository_connections(self, org_id: str = None, inc
                 "organization_id": org_id
             }
 
+    def _get_current_repository_name(self, repo_name: str = None) -> str:
+        """Get repository name from parameter or auto-detect from current directory"""
+        if repo_name:
+            return repo_name
+        return os.path.basename(os.getcwd())
+    
+    async def _get_organization_id(self, org_id: str = None) -> str:
+        """Get organization ID from parameter or auto-detect from user info"""
+        if org_id:
+            return org_id
+        
+        user_info = await self.client.get_user_info()
+        organizations = user_info.get("user", {}).get("external", {}).get("organizations", [])
+        
+        if not organizations:
+            raise ValueError("No organizations found for user")
+        
+        return organizations[0]["organization"]["id"]
+    
+    async def _find_matching_repositories(self, repo_name: str, org_id: str) -> List[Dict[str, Any]]:
+        """Find repositories matching the given name (case-insensitive)"""
+        repos_response = await self.client.list_repositories(org_id, pageSize=1000)
+        repositories = repos_response.get("repositories", [])
+        
+        # Find matching repositories (case-insensitive)
+        return [
+            repo for repo in repositories
+            if repo.get("name", "").lower() == repo_name.lower()
+        ]
+
     def _calculate_name_similarity(self, name1: str, name2: str) -> float:
         """Calculate similarity between two names using simple string matching"""
         if not name1 or not name2:
@@ -2483,10 +2487,8 @@ def _calculate_name_similarity(self, name1: str, name2: str) -> float:
     async def _get_comprehensive_sensitive_data_summary(self, org_id: str = None, time_period: str = "30d", include_trends: bool = True, include_critical_only: bool = False, include_recommendations: bool = True, group_by: str = "data_type", **kwargs) -> Dict[str, Any]:
         """Get a comprehensive sensitive data summary combining multiple analysis approaches"""
         try:
-            # Get org_id if not provided
-            if not org_id:
-                user_info = await self.client.get_user_info()
-                org_id = user_info["user"]["external"]["organizations"][0]["organization"]["id"]
+            # Get org_id using helper method
+            org_id = await self._get_organization_id(org_id)
 
             # Get all sensitive data findings
             if time_period == "all":
diff --git a/test_new_tools.py b/test_new_tools.py
@@ -47,8 +47,8 @@ async def test_new_tools():
         # Test 1: Check Repository Attack Surface
         print("1. Testing check_repository_attack_surface...")
         try:
-            # Test with the current repository name
-            current_repo = "stackhawk-mcp"  # This repo
+            # Test with the current repository name (auto-detected from directory)
+            current_repo = os.path.basename(os.getcwd())  # Dynamic detection
             result = await server._check_repository_attack_surface(
                 repo_name=current_repo,
                 include_vulnerabilities=True,
@@ -159,7 +159,16 @@ async def test_new_tools():
                 "get_sensitive_data_summary"
             ]
             
-            found_tools = [tool.name for tool in tools if tool.name in new_tool_names]
+            # Handle different possible return types from _list_tools_handler
+            tool_names = []
+            if tools:
+                for tool in tools:
+                    if hasattr(tool, 'name'):
+                        tool_names.append(tool.name)
+                    elif isinstance(tool, dict) and 'name' in tool:
+                        tool_names.append(tool['name'])
+            
+            found_tools = [name for name in tool_names if name in new_tool_names]
             print(f"&#9989; Found {len(found_tools)}/{len(new_tool_names)} new tools in MCP interface")
             for tool_name in found_tools:
                 print(f"   &#10003; {tool_name}")
