Skip to content

Commit b215a37

Browse files
JoeBeeContrastJoeBeeContrast
committed
updated the queries to find app by appname and get all apps to use a caching mechanism to speed up the response.
1 parent 527c25e commit b215a37

File tree

4 files changed

+50
-32
lines changed

4 files changed

+50
-32
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ public ProtectData getProtectData(String applicationName) throws IOException {
5858

5959
// Get application ID from name
6060
logger.debug("Looking up application ID for name: {}", applicationName);
61-
String appID = SDKHelper.getAppIDFromapp_name(applicationName, orgID, contrastSDK);
61+
String appID = SDKHelper.getAppIDFromName(applicationName, orgID, contrastSDK);
6262
if (appID == null || appID.isEmpty()) {
6363
logger.warn("No application ID found for application: {}", applicationName);
6464
return null;

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

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ public Vulnerability getVulnerabilityById(String vulnID, String appID) throws IO
102102
List<LibraryExtended> libs = SDKHelper.getLibsForID(appID,orgID, new SDKExtension(contrastSDK));
103103
List<LibraryLibraryObservation> lobs = new ArrayList<>();
104104
for(LibraryExtended lib : libs) {
105-
LibraryLibraryObservation llob = new LibraryLibraryObservation(lib,new SDKExtension(contrastSDK).getLibraryObservations(orgID,appID,lib.getHash(),50));
105+
LibraryLibraryObservation llob = new LibraryLibraryObservation(lib, SDKHelper.getLibraryObservationsWithCache(lib.getHash(), appID, orgID, 50,new SDKExtension(contrastSDK)));
106106
lobs.add(llob);
107107
}
108108
List<StackLib> stackLibs = new ArrayList<>();
@@ -157,7 +157,7 @@ public Vulnerability getVulnerability(String vulnID, String app_name) throws IOE
157157
Optional<String> appID = Optional.empty();
158158
logger.debug("Searching for application ID matching name: {}", app_name);
159159

160-
for(Application app : contrastSDK.getApplications(orgID).getApplications()) {
160+
for(Application app : SDKHelper.getApplicationsWithCache(orgID, contrastSDK)) {
161161
if(app.getName().toLowerCase().contains(app_name.toLowerCase())) {
162162
appID = Optional.of(app.getId());
163163
logger.debug("Found matching application - ID: {}, Name: {}", app.getId(), app.getName());
@@ -176,7 +176,6 @@ public Vulnerability getVulnerability(String vulnID, String app_name) throws IOE
176176
public List<VulnLight> listVulnsByAppId(String appID) throws IOException {
177177
logger.info("Listing vulnerabilities for application ID: {}", appID);
178178
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName);
179-
180179
try {
181180
List<Trace> traces = contrastSDK.getTraces(orgID, appID, new TraceFilterBody()).getTraces();
182181
logger.debug("Found {} vulnerability traces for application ID: {}", traces.size(), appID);
@@ -203,7 +202,7 @@ public List<VulnLight> listVulnsInAppByName(String app_name) throws IOException
203202
Optional<String> appID = Optional.empty();
204203
logger.debug("Searching for application ID matching name: {}", app_name);
205204

206-
for(Application app : contrastSDK.getApplications(orgID).getApplications()) {
205+
for(Application app : SDKHelper.getApplicationsWithCache(orgID, contrastSDK)) {
207206
if(app.getName().toLowerCase().contains(app_name.toLowerCase())) {
208207
appID = Optional.of(app.getId());
209208
logger.debug("Found matching application - ID: {}, Name: {}", app.getId(), app.getName());
@@ -225,11 +224,11 @@ public List<VulnLight> listVulnsInAppByName(String app_name) throws IOException
225224

226225

227226
@Tool(name = "list_applications", description = "Takes an application name (app_name) returns a list of active applications matching that name. Please remember to display the name, status and ID.")
228-
public List<ApplicationData> getActiveApplications(String app_name) throws IOException {
227+
public List<ApplicationData> getApplications(String app_name) throws IOException {
229228
logger.info("Listing active applications matching name: {}", app_name);
230229
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName);
231230
try {
232-
List<Application> applications = contrastSDK.getApplications(orgID).getApplications();
231+
List<Application> applications = SDKHelper.getApplicationsWithCache(orgID, contrastSDK);
233232
logger.debug("Retrieved {} total applications from Contrast", applications.size());
234233

235234
List<ApplicationData> filteredApps = new ArrayList<>();
@@ -255,16 +254,16 @@ public List<ApplicationData> getActiveApplications() throws IOException {
255254
logger.info("Listing all applications");
256255
ContrastSDK contrastSDK = SDKHelper.getSDK(hostName, apiKey, serviceKey, userName);
257256
try {
258-
List<Application> applications = contrastSDK.getApplications(orgID).getApplications();
257+
List<Application> applications = SDKHelper.getApplicationsWithCache(orgID, contrastSDK);
259258
logger.debug("Retrieved {} total applications from Contrast", applications.size());
260259

261-
List<ApplicationData> ReturnedApps = new ArrayList<>();
260+
List<ApplicationData> returnedApps = new ArrayList<>();
262261
for(Application app : applications) {
263-
ReturnedApps.add(new ApplicationData(app.getName(), app.getStatus(), app.getId()));
262+
returnedApps.add(new ApplicationData(app.getName(), app.getStatus(), app.getId()));
264263
}
265264

266-
logger.info("Found {} applications matching'", ReturnedApps.size());
267-
return ReturnedApps;
265+
logger.info("Found {} applications", returnedApps.size());
266+
return returnedApps;
268267

269268
} catch (Exception e) {
270269
logger.error("Error listing all applications", e);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ public List<LibraryExtended> getApplicationLibraries(String app_name) throws IOE
7676
Optional<String> appID = Optional.empty();
7777
logger.debug("Searching for application ID matching name: {}", app_name);
7878

79-
for(Application app : contrastSDK.getApplications(orgID).getApplications()) {
79+
for(Application app : SDKHelper.getApplicationsWithCache(orgID, contrastSDK)) {
8080
if(app.getName().toLowerCase().contains(app_name.toLowerCase())) {
8181
appID = Optional.of(app.getId());
8282
logger.info("Found matching application - ID: {}, Name: {}", app.getId(), app.getName());

‎src/main/java/com/contrast/labs/ai/mcp/contrast/sdkexstension/SDKHelper.java

Lines changed: 38 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -47,17 +47,17 @@ public class SDKHelper {
4747
.maximumSize(100)
4848
.expireAfterWrite(5, TimeUnit.MINUTES)
4949
.build();
50-
51-
private static final Cache<String, String> appIDCache = CacheBuilder.newBuilder()
52-
.maximumSize(100)
53-
.expireAfterWrite(10, TimeUnit.MINUTES)
54-
.build();
5550

5651
private static final Cache<String, List<LibraryObservation>> libraryObservationsCache = CacheBuilder.newBuilder()
5752
.maximumSize(5000)
5853
.expireAfterWrite(60, TimeUnit.MINUTES)
5954
.build();
6055

56+
private static final Cache<String, List<Application>> applicationsCache = CacheBuilder.newBuilder()
57+
.maximumSize(5000)
58+
.expireAfterWrite(60, TimeUnit.MINUTES)
59+
.build();
60+
6161
public static List<LibraryExtended> getLibsForID(String appID, String orgID, SDKExtension extendedSDK) throws IOException {
6262
// Check cache for existing result
6363
List<LibraryExtended> cachedLibraries = libraryCache.getIfPresent(appID);
@@ -90,30 +90,21 @@ public static List<LibraryExtended> getLibsForID(String appID, String orgID, SDK
9090
return libs;
9191
}
9292

93-
public static String getAppIDFromapp_name(String app_name, String orgID, ContrastSDK contrastSDK) throws IOException {
94-
// Check cache for existing result
95-
String cachedAppID = appIDCache.getIfPresent(app_name);
96-
if (cachedAppID != null) {
97-
logger.info("Cache hit for application name: {}", app_name);
98-
return cachedAppID;
99-
}
93+
public static String getAppIDFromName(String app_name, String orgID, ContrastSDK contrastSDK) throws IOException {
10094
logger.debug("Cache miss for application name: {}, searching for application ID", app_name);
10195
Optional<String> appID = Optional.empty();
102-
for (Application app : contrastSDK.getApplications(orgID).getApplications()) {
96+
for (Application app : getApplicationsWithCache(orgID, contrastSDK)) {
10397
if (app.getName().toLowerCase().contains(app_name.toLowerCase())) {
10498
appID = Optional.of(app.getId());
10599
logger.info("Found matching application - ID: {}, Name: {}", app.getId(), app.getName());
106100
break;
107101
}
108102
}
109-
110103
if (appID.isPresent()) {
111-
// Store result in cache
112-
appIDCache.put(app_name, appID.get());
113104
return appID.get();
114105
} else {
115-
logger.error("Application not found: {}", app_name);
116-
throw new IOException("Application not found");
106+
logger.warn("No application found with name: {}", app_name);
107+
return null;
117108
}
118109
}
119110

@@ -174,8 +165,36 @@ public static ContrastSDK getSDK(String hostName, String apiKey, String serviceK
174165
.withApiUrl( "https://" + hostName + "/Contrast/api")
175166
.withUserAgentProduct(UserAgentProduct.of(MCP_SERVER_NAME,MCP_VERSION))
176167
.build();
177-
178168
}
179169

170+
/**
171+
* Retrieves all applications from Contrast with caching.
172+
*
173+
* @param orgId The organization ID
174+
* @param contrastSDK The Contrast SDK instance
175+
* @return List of applications
176+
* @throws IOException If an I/O error occurs
177+
*/
178+
public static List<Application> getApplicationsWithCache(String orgId, ContrastSDK contrastSDK) throws IOException {
179+
// Generate cache key based on organization ID
180+
String cacheKey = String.format("applications:%s", orgId);
181+
182+
// Check cache for existing result
183+
List<Application> cachedApplications = applicationsCache.getIfPresent(cacheKey);
184+
if (cachedApplications != null) {
185+
logger.info("Cache hit for applications in org: {}", orgId);
186+
return cachedApplications;
187+
}
180188

189+
logger.info("Cache miss for applications in org: {}, fetching from API", orgId);
190+
List<Application> applications = contrastSDK.getApplications(orgId).getApplications();
191+
192+
logger.info("Successfully retrieved {} applications from organization: {}",
193+
applications.size(), orgId);
194+
195+
// Store result in cache
196+
applicationsCache.put(cacheKey, applications);
197+
198+
return applications;
199+
}
181200
}

0 commit comments

Comments
 (0)