Skip to content

Commit f7758ad

Browse files
ChrisEdwardsChrisEdwards
committed
Addressed PR comments
1 parent cd89117 commit f7758ad

File tree

4 files changed

+46
-50
lines changed

4 files changed

+46
-50
lines changed

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

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -213,18 +213,19 @@ public List<VulnLight> listVulnsByAppIdAndSessionMetadata(
213213
List<VulnLight> vulns = listVulnsByAppId(appID);
214214
List<VulnLight> returnVulns = new ArrayList<>();
215215
for(VulnLight vuln : vulns) {
216-
if(vuln.sessionMetadata()!=null) {
217-
for(SessionMetadata sm : vuln.sessionMetadata()) {
218-
for(MetadataItem metadataItem : sm.getMetadata()) {
219-
if(metadataItem.getDisplayLabel().equalsIgnoreCase(session_Metadata_Name) &&
220-
metadataItem.getValue().equalsIgnoreCase(session_Metadata_Value)) {
221-
returnVulns.add(vuln);
222-
logger.debug("Found matching vulnerability with ID: {}", vuln.vulnID());
223-
break;
224-
}
225-
}
226-
}
216+
if (vuln.sessionMetadata() == null) {
217+
continue;
218+
}
219+
for (SessionMetadata sm : vuln.sessionMetadata()) {
220+
for (MetadataItem metadataItem : sm.getMetadata()) {
221+
if (metadataItem.getDisplayLabel().equalsIgnoreCase(session_Metadata_Name) &&
222+
metadataItem.getValue().equalsIgnoreCase(session_Metadata_Value)) {
223+
returnVulns.add(vuln);
224+
logger.debug("Found matching vulnerability with ID: {}", vuln.vulnID());
225+
break;
226+
}
227227
}
228+
}
228229
}
229230
return returnVulns;
230231
} catch (Exception 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
@@ -64,7 +64,7 @@ public class SCAService {
6464

6565
@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")
6666
public List<LibraryExtended> getApplicationLibrariesByID(String appID) throws IOException {
67-
if (appID == null || appID.isEmpty()) {
67+
if (appID == null || appID.isBlank()) {
6868
throw new IllegalArgumentException("Application ID cannot be null or empty");
6969
}
7070
logger.info("Retrieving libraries for application id: {}", appID);

‎src/test/java/com/contrast/labs/ai/mcp/contrast/ADRServiceTest.java‎

Lines changed: 27 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import com.contrast.labs.ai.mcp.contrast.data.PaginatedResponse;
2020
import com.contrast.labs.ai.mcp.contrast.sdkexstension.SDKExtension;
2121
import com.contrast.labs.ai.mcp.contrast.sdkexstension.SDKHelper;
22+
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData;
23+
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule;
2224
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.adr.Attack;
2325
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.adr.AttacksFilterBody;
2426
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.adr.AttacksResponse;
@@ -545,15 +547,15 @@ void testGetAttacks_MultipleValidationErrors_CombinesErrors() throws Exception {
545547
@Test
546548
void testGetProtectDataByAppID_Success() throws Exception {
547549
// Given
548-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData mockProtectData = createMockProtectData(3);
550+
ProtectData mockProtectData = createMockProtectData(3);
549551

550552
mockedSDKExtension = mockConstruction(SDKExtension.class, (mock, context) -> {
551553
when(mock.getProtectConfig(eq(TEST_ORG_ID), eq(TEST_APP_ID)))
552554
.thenReturn(mockProtectData);
553555
});
554556

555557
// When
556-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData result =
558+
ProtectData result =
557559
adrService.getProtectDataByAppID(TEST_APP_ID);
558560

559561
// Then
@@ -565,15 +567,15 @@ void testGetProtectDataByAppID_Success() throws Exception {
565567
@Test
566568
void testGetProtectDataByAppID_WithRules() throws Exception {
567569
// Given
568-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData mockProtectData = createMockProtectDataWithRules();
570+
ProtectData mockProtectData = createMockProtectDataWithRules();
569571

570572
mockedSDKExtension = mockConstruction(SDKExtension.class, (mock, context) -> {
571573
when(mock.getProtectConfig(eq(TEST_ORG_ID), eq(TEST_APP_ID)))
572574
.thenReturn(mockProtectData);
573575
});
574576

575577
// When
576-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData result =
578+
ProtectData result =
577579
adrService.getProtectDataByAppID(TEST_APP_ID);
578580

579581
// Then
@@ -637,7 +639,7 @@ void testGetProtectDataByAppID_NoProtectDataReturned() throws Exception {
637639
});
638640

639641
// When
640-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData result =
642+
ProtectData result =
641643
adrService.getProtectDataByAppID(TEST_APP_ID);
642644

643645
// Then
@@ -647,8 +649,8 @@ void testGetProtectDataByAppID_NoProtectDataReturned() throws Exception {
647649
@Test
648650
void testGetProtectDataByAppID_EmptyRulesList() throws Exception {
649651
// Given - Protect enabled but no rules configured
650-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData mockProtectData =
651-
new com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData();
652+
ProtectData mockProtectData =
653+
new ProtectData();
652654
mockProtectData.setRules(new ArrayList<>());
653655

654656
mockedSDKExtension = mockConstruction(SDKExtension.class, (mock, context) -> {
@@ -657,7 +659,7 @@ void testGetProtectDataByAppID_EmptyRulesList() throws Exception {
657659
});
658660

659661
// When
660-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData result =
662+
ProtectData result =
661663
adrService.getProtectDataByAppID(TEST_APP_ID);
662664

663665
// Then
@@ -668,19 +670,14 @@ void testGetProtectDataByAppID_EmptyRulesList() throws Exception {
668670

669671
// ========== Helper Methods ==========
670672

671-
/**
672-
* Creates mock AttacksResponse for testing
673-
*/
674673
private AttacksResponse createMockAttacksResponse(int count, Integer totalCount) {
675674
AttacksResponse response = new AttacksResponse();
676675
response.setAttacks(createMockAttacks(count));
677676
response.setCount(totalCount);
678677
return response;
679678
}
680679

681-
/**
682-
* Creates mock Attack objects for testing
683-
*/
680+
684681
private List<Attack> createMockAttacks(int count) {
685682
List<Attack> attacks = new ArrayList<>();
686683
long baseTime = System.currentTimeMillis();
@@ -703,17 +700,15 @@ private List<Attack> createMockAttacks(int count) {
703700
return attacks;
704701
}
705702

706-
/**
707-
* Creates mock ProtectData for testing
708-
*/
709-
private com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData createMockProtectData(int ruleCount) {
710-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData protectData =
711-
new com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData();
712703

713-
List<com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule> rules = new ArrayList<>();
704+
private ProtectData createMockProtectData(int ruleCount) {
705+
ProtectData protectData =
706+
new ProtectData();
707+
708+
List<Rule> rules = new ArrayList<>();
714709
for (int i = 0; i < ruleCount; i++) {
715-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule rule =
716-
new com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule();
710+
Rule rule =
711+
new Rule();
717712
rule.setName("protect-rule-" + i);
718713
rule.setProduction(i % 2 == 0 ? "block" : "monitor");
719714
rules.add(rule);
@@ -723,25 +718,23 @@ private com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData createM
723718
return protectData;
724719
}
725720

726-
/**
727-
* Creates mock ProtectData with realistic rule configuration
728-
*/
729-
private com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData createMockProtectDataWithRules() {
730-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData protectData =
731-
new com.contrast.labs.ai.mcp.contrast.sdkexstension.data.ProtectData();
732721

733-
List<com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule> rules = new ArrayList<>();
722+
private ProtectData createMockProtectDataWithRules() {
723+
ProtectData protectData =
724+
new ProtectData();
725+
726+
List<Rule> rules = new ArrayList<>();
734727

735728
// SQL Injection rule
736-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule sqlRule =
737-
new com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule();
729+
Rule sqlRule =
730+
new Rule();
738731
sqlRule.setName("sql-injection");
739732
sqlRule.setProduction("block");
740733
rules.add(sqlRule);
741734

742735
// XSS rule
743-
com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule xssRule =
744-
new com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Rule();
736+
Rule xssRule =
737+
new Rule();
745738
xssRule.setName("xss-reflected");
746739
xssRule.setProduction("monitor");
747740
rules.add(xssRule);

‎src/test/java/com/contrast/labs/ai/mcp/contrast/SCAServiceTest.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,9 @@
1717

1818
import com.contrast.labs.ai.mcp.contrast.sdkexstension.SDKExtension;
1919
import com.contrast.labs.ai.mcp.contrast.sdkexstension.SDKHelper;
20+
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.App;
2021
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.CveData;
22+
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Library;
2123
import com.contrast.labs.ai.mcp.contrast.sdkexstension.data.LibraryExtended;
2224
import com.contrastsecurity.sdk.ContrastSDK;
2325
import org.junit.jupiter.api.AfterEach;
@@ -360,21 +362,21 @@ private CveData createMockCveData() {
360362
private CveData createMockCveDataWithApps() {
361363
CveData cveData = new CveData();
362364

363-
var app = mock(com.contrast.labs.ai.mcp.contrast.sdkexstension.data.App.class);
365+
var app = mock(App.class);
364366
when(app.getApp_id()).thenReturn(TEST_APP_ID);
365367
when(app.getName()).thenReturn("Test Application");
366368
when(app.getClassCount()).thenReturn(0);
367369

368-
var apps = new ArrayList<com.contrast.labs.ai.mcp.contrast.sdkexstension.data.App>();
370+
var apps = new ArrayList<App>();
369371
apps.add(app);
370372
cveData.setApps(apps);
371373

372-
var lib = mock(com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Library.class);
374+
var lib = mock(Library.class);
373375
when(lib.getHash()).thenReturn("matching-hash-789");
374376
when(lib.getFile_name()).thenReturn("vulnerable-lib.jar");
375377
when(lib.getVersion()).thenReturn("1.0.0");
376378

377-
var libs = new ArrayList<com.contrast.labs.ai.mcp.contrast.sdkexstension.data.Library>();
379+
var libs = new ArrayList<Library>();
378380
libs.add(lib);
379381
cveData.setLibraries(libs);
380382

0 commit comments

Comments
 (0)