Merge pull request #26 from Contrast-Security-OSS/AIML-189-remove-duplicate-tools

AIML-189: Remove duplicate app_name/app_id tool variants to optimize AI context usage

Author: ChrisEdwards

src/main/java/com/contrast/labs/ai/mcp/contrast/ADRService.java

@@ -71,39 +71,7 @@ public ADRService(PaginationHandler paginationHandler) {
     private String httpProxyPort;
 
 
-    @Tool(name = "get_ADR_Protect_Rules", description = "takes a application name and returns the protect / adr rules for the application")
-    public ProtectData getProtectData(
-            @ToolParam(description = "Application name") String applicationName) throws IOException {
-        logger.info("Starting retrieval of protection rules for application: {}", applicationName);
-        long startTime = System.currentTimeMillis();
-
-        try {
-            ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
-            logger.debug("ContrastSDK initialized successfully for application: {}", applicationName);
-
-            // Get application ID from name
-            logger.debug("Looking up application ID for name: {}", applicationName);
-            Optional<Application> app = SDKHelper.getApplicationByName(applicationName, orgID, contrastSDK);
-            if (app.isEmpty()) {
-                logger.warn("No application ID found for application: {}", applicationName);
-                return null;
-            }
-            logger.debug("Found application ID: {} for application: {}", app.get().getAppId(), applicationName);
-
-            ProtectData result = getProtectDataByAppID(app.get().getAppId());
-            long duration = System.currentTimeMillis() - startTime;
-            logger.info("Completed retrieval of protection rules for application: {} (took {} ms)", applicationName, duration);
-            return result;
-        } catch (Exception e) {
-            long duration = System.currentTimeMillis() - startTime;
-            logger.error("Error retrieving protection rules for application: {} (after {} ms): {}",
-                    applicationName, duration, e.getMessage(), e);
-            throw e;
-        }
-    }
-
-
-    @Tool(name = "get_ADR_Protect_Rules_by_app_id", description = "takes a application ID and returns the protect / adr rules for the application")
+    @Tool(name = "get_ADR_Protect_Rules", description = "Takes an application ID and returns the Protect/ADR rules for the application. Use list_applications_with_name first to get the application ID from a name")
     public ProtectData getProtectDataByAppID(
             @ToolParam(description = "Application ID") String appID) throws IOException {
         if (appID == null || appID.isEmpty()) {

src/main/java/com/contrast/labs/ai/mcp/contrast/AssessService.java

@@ -102,7 +102,7 @@ public AssessService(VulnerabilityMapper vulnerabilityMapper, PaginationHandler
 
 
 
-    @Tool(name = "get_vulnerability_by_id", description = "takes a vulnerability ID ( vulnID ) and Application ID ( appID ) and returns details about the specific security vulnerability. If based on the stacktrace, the vulnerability looks like it is in code that is not in the codebase, the vulnerability may be in a 3rd party library, review the CVE data attached to that stackframe you believe the vulnerability exists in and if possible upgrade that library to the next non vulnerable version based on the remediation guidance.")
+    @Tool(name = "get_vulnerability", description = "Takes a vulnerability ID (vulnID) and application ID (appID) and returns details about the specific security vulnerability. Use list_applications_with_name first to get the application ID from a name. If based on the stacktrace, the vulnerability looks like it is in code that is not in the codebase, the vulnerability may be in a 3rd party library, review the CVE data attached to that stackframe you believe the vulnerability exists in and if possible upgrade that library to the next non vulnerable version based on the remediation guidance.")
     public Vulnerability getVulnerabilityById(
             @ToolParam(description = "Vulnerability ID (UUID format)") String vulnID,
             @ToolParam(description = "Application ID") String appID) throws IOException {
@@ -189,24 +189,7 @@ private Optional<LibraryLibraryObservation> findMatchingLibraryData(String stack
         return Optional.empty();
     }
 
-    @Tool(name = "get_vulnerability", description = "Takes a vulnerability ID (vulnID) and application name (app_name) and returns details about the specific security vulnerability.  If based on the stacktrace, the vulnerability looks like it is in code that is not in the codebase, the vulnerability may be in a 3rd party library, review the CVE data attached to that stackframe you believe the vulnerability exists in and if possible upgrade that library to the next non vulnerable version based on the remediation guidance.")
-    public Vulnerability getVulnerability(
-            @ToolParam(description = "Vulnerability ID (UUID format)") String vulnID,
-            @ToolParam(description = "Application name") String app_name) throws IOException {
-        logger.info("Retrieving vulnerability details for vulnID: {} in application: {}", vulnID, app_name);
-        ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
-        logger.debug("Searching for application ID matching name: {}", app_name);
-
-        Optional<Application> application = SDKHelper.getApplicationByName(app_name, orgID, contrastSDK);
-        if(application.isPresent()) {
-            return getVulnerabilityById(vulnID, application.get().getAppId());
-        } else {
-            logger.error("Application with name {} not found", app_name);
-            throw new IllegalArgumentException("Application with name " + app_name + " not found");
-        }
-    }
-
-    @Tool(name = "list_vulnerabilities_with_id", description = "Takes a  Application ID ( appID ) and returns a list of vulnerabilities, please remember to include the vulnID in the response.")
+    @Tool(name = "list_vulnerabilities", description = "Takes an application ID (appID) and returns a list of vulnerabilities. Use list_applications_with_name first to get the application ID from a name. Remember to include the vulnID in the response.")
     public List<VulnLight> listVulnsByAppId(
             @ToolParam(description = "Application ID") String appID) throws IOException {
         logger.info("Listing vulnerabilities for application ID: {}", appID);
@@ -337,29 +320,6 @@ public MetadataFilterResponse listSessionMetadataForApplication(
         }
     }
 
-    @Tool(name = "list_vulnerabilities", description = "Takes an application name ( app_name ) and returns a list of vulnerabilities, please remember to include the vulnID in the response.  ")
-    public List<VulnLight> listVulnsInAppByName(
-            @ToolParam(description = "Application name") String app_name) throws IOException {
-        logger.info("Listing vulnerabilities for application: {}", app_name);
-        ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
-        
-        logger.debug("Searching for application ID matching name: {}", app_name);
-
-        Optional<Application> application = SDKHelper.getApplicationByName(app_name, orgID, contrastSDK);
-        if(application.isPresent()) {
-            try {
-              return listVulnsByAppId(application.get().getAppId());
-            } catch (Exception e) {
-                logger.error("Error listing vulnerabilities for application: {}", app_name, e);
-                throw new IOException("Failed to list vulnerabilities: " + e.getMessage(), e);
-            }
-        } else {
-            logger.debug("Application with name {} not found, returning empty list", app_name);
-            return new ArrayList<>();
-        }
-    }
-
-
     @Tool(name = "list_applications_with_name", description = "Takes an application name (app_name) returns a list of active applications that contain that name. Please remember to display the name, status and ID.")
     public List<ApplicationData> getApplications(
             @ToolParam(description = "Application name (supports partial matching, case-insensitive)") String app_name) throws IOException {

src/main/java/com/contrast/labs/ai/mcp/contrast/SCAService.java

@@ -62,8 +62,11 @@ public class SCAService {
     private String httpProxyPort;
 
 
-    @Tool(name = "list_application_libraries_by_app_id", description = "Takes a application ID and returns the libraries used in the application, note if class usage count is 0 the library is unlikely to be used")
+    @Tool(name = "list_application_libraries", description = "Takes an application ID and returns the libraries used in the application. Use list_applications_with_name first to get the application ID from a name. Note: if class usage count is 0 the library is unlikely to be used")
     public List<LibraryExtended> getApplicationLibrariesByID(String appID) throws IOException {
+        if (appID == null || appID.isEmpty()) {
+            throw new IllegalArgumentException("Application ID cannot be null or empty");
+        }
         logger.info("Retrieving libraries for application id: {}", appID);
         ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
         logger.debug("ContrastSDK initialized with host: {}", hostName);
@@ -73,25 +76,6 @@ public List<LibraryExtended> getApplicationLibrariesByID(String appID) throws IO
 
     }
 
-
-    @Tool(name = "list_application_libraries", description = "takes a application name and returns the libraries used in the application, note if class usage count is 0 the library is unlikely to be used")
-    public List<LibraryExtended> getApplicationLibraries(String app_name) throws IOException {
-        logger.info("Retrieving libraries for application: {}", app_name);
-        ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName,httpProxyHost, httpProxyPort);
-        logger.debug("ContrastSDK initialized with host: {}", hostName);
-        
-        SDKExtension extendedSDK = new SDKExtension(contrastSDK);
-        logger.debug("Searching for application ID matching name: {}", app_name);
-
-        Optional<Application> application = SDKHelper.getApplicationByName(app_name, orgID, contrastSDK);
-        if(application.isPresent()) {
-            return SDKHelper.getLibsForID(application.get().getAppId(),orgID, extendedSDK);
-        } else {
-            logger.error("Application not found: {}", app_name);
-            throw new IOException("Application not found");
-        }
-    }
-
     @Tool(name= "list_applications_vulnerable_to_cve", description = "takes a cve id and returns the applications and servers vulnerable to the cve. Please note if the application class usage is 0, its unlikely to be vulnerable")
     public CveData listCVESForApplication(String cveid) throws IOException {
         logger.info("Retrieving applications vulnerable to CVE: {}", cveid);