Skip to content

Commit 61db1d3

Browse files
njmulsqbkingthorin
njmulsqb
authored and
kingthorin
committed
pass otherinfo to llm
Signed-off-by: Najam Ul Saqib <[email protected]>
1 parent 2134216 commit 61db1d3

File tree

3 files changed

+161
-19
lines changed

3 files changed

+161
-19
lines changed

‎addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmAssistant.java‎

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,22 +34,45 @@ public interface LlmAssistant {
3434
"As a software architect, and based on your previous answer, generate other potential missing endpoints that are not mentioned in the OpenAPI file. For example, if there is GET /product/1, suggest DELETE /product/1 if it's not mentioned")
3535
HttpRequestList complete();
3636

37-
@SystemMessage(
38-
"You are a web application security expert reviewing potential false positives. Answer only in JSON.")
39-
@UserMessage(
40-
"Your task is to review the following finding from ZAP (Zed Attack Proxy).\n"
41-
+ "The confidence level is a pull down field which allows you to specify how confident you are in the validity of the finding : \n"
42-
+ "- 0 if it's False Positive\n"
43-
+ "- 1 if it's Low\n"
44-
+ "- 2 if it's Medium\n"
45-
+ "- 3 if it's High\n"
46-
+ "\n"
47-
+ "The alert is described as follows : {{description}}\n"
48-
+ "\n"
49-
+ "As evidence, the HTTP response contains :\n"
50-
+ "---\n"
51-
+ "{{evidence}}\n"
52-
+ "---\n"
53-
+ "Provide a short consistent explanation of the new score.\n")
37+
static final String ALERT_REVIEW_SYSTEM_MSG =
38+
"You are a web application security expert reviewing potential false positives. Answer only in JSON.";
39+
static final String ALERT_REVIEW_GOAL =
40+
"Provide a short consistent explanation of the new score.\n";
41+
static final String ALERT_REVIEW_PROMPT =
42+
"""
43+
Your task is to review the following finding from ZAP (Zed Attack Proxy).
44+
The confidence level is a pull down field which allows you to specify how confident you are in the validity of the finding:
45+
- 0 if it's False Positive
46+
- 1 if it's Low
47+
- 2 if it's Medium
48+
- 3 if it's High
49+
50+
The alert is described as follows : {{description}}
51+
52+
As evidence, the HTTP response contains:
53+
---
54+
{{evidence}}
55+
---
56+
"""
57+
+ ALERT_REVIEW_GOAL;
58+
59+
static final String ALERT_REVIEW_OTHERINFO_PROMPT =
60+
ALERT_REVIEW_PROMPT
61+
+ """
62+
Also, here's some additional information that may be useful for you to reach your conclusion:
63+
---
64+
{{otherinfo}}
65+
"""
66+
+ ALERT_REVIEW_GOAL;
67+
68+
@SystemMessage(ALERT_REVIEW_SYSTEM_MSG)
69+
@UserMessage(ALERT_REVIEW_PROMPT)
5470
Confidence review(@V("description") String description, @V("evidence") String evidence);
71+
72+
@SystemMessage(ALERT_REVIEW_SYSTEM_MSG)
73+
@UserMessage(ALERT_REVIEW_OTHERINFO_PROMPT)
74+
Confidence review(
75+
@V("description") String description,
76+
@V("evidence") String evidence,
77+
@V("otherinfo") String otherinfo);
5578
}

‎addOns/llm/src/main/java/org/zaproxy/addon/llm/services/LlmCommunicationService.java‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import java.util.Map;
3737
import java.util.stream.Collectors;
3838
import org.apache.commons.httpclient.util.HttpURLConnection;
39+
import org.apache.commons.lang3.StringUtils;
3940
import org.apache.logging.log4j.LogManager;
4041
import org.apache.logging.log4j.Logger;
4142
import org.parosproxy.paros.Constant;
@@ -75,6 +76,9 @@ public LlmCommunicationService(LlmOptions options) {
7576
requestor = new Requestor(HttpSender.MANUAL_REQUEST_INITIATOR, new HistoryPersister());
7677
}
7778

79+
/** For testing purposes only. */
80+
LlmCommunicationService() {}
81+
7882
private ChatLanguageModel buildModel(LlmOptions options) {
7983
return switch (options.getModelProvider()) {
8084
case AZURE_OPENAI ->
@@ -170,7 +174,13 @@ public void reviewAlert(Alert alert) {
170174
LOGGER.debug("Reviewing alert : {}", alert.getName());
171175
LOGGER.debug("Confidence level from ZAP : {}", alert.getConfidence());
172176
Stats.incCounter("stats.llm.alertreview.call");
173-
llmConfidence = llmAssistant.review(alert.getDescription(), alert.getEvidence());
177+
if (StringUtils.isBlank(alert.getOtherInfo())) {
178+
llmConfidence = llmAssistant.review(alert.getDescription(), alert.getEvidence());
179+
} else {
180+
llmConfidence =
181+
llmAssistant.review(
182+
alert.getDescription(), alert.getEvidence(), alert.getOtherInfo());
183+
}
174184

175185
if (llmConfidence.getLevel() == alert.getConfidence()) {
176186
Stats.incCounter("stats.llm.alertreview.result.same");
@@ -206,7 +216,7 @@ public void reviewAlert(Alert alert) {
206216
}
207217
}
208218

209-
private static boolean isPreviouslyReviewed(Alert alert) {
219+
static boolean isPreviouslyReviewed(Alert alert) {
210220
return !alert.getTags().containsKey(AI_REVIEWED_TAG_KEY);
211221
}
212222

@@ -218,4 +228,9 @@ private static String getUpdatedOtherInfo(Alert alert, Confidence llmConfidence)
218228
private static ExtensionAlert getExtAlert() {
219229
return Control.getSingleton().getExtensionLoader().getExtension(ExtensionAlert.class);
220230
}
231+
232+
/* For testing purposes. Not intended for regular use. */
233+
void setAssistant(LlmAssistant assistant) {
234+
this.llmAssistant = assistant;
235+
}
221236
}

‎addOns/llm/src/test/java/org/zaproxy/addon/llm/services/LlmCommunicationServiceUnitTest.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Zed Attack Proxy (ZAP) and its related class files.
3+
*
4+
* ZAP is an HTTP/HTTPS proxy for assessing web application security.
5+
*
6+
* Copyright 2025 The ZAP Development Team
7+
*
8+
* Licensed under the Apache License, Version 2.0 (the "License");
9+
* you may not use this file except in compliance with the License.
10+
* You may obtain a copy of the License at
11+
*
12+
* http://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
20+
package org.zaproxy.addon.llm.services;
21+
22+
import static org.mockito.ArgumentMatchers.anyString;
23+
import static org.mockito.BDDMockito.given;
24+
import static org.mockito.Mockito.mock;
25+
import static org.mockito.Mockito.never;
26+
import static org.mockito.Mockito.verify;
27+
import static org.mockito.Mockito.withSettings;
28+
29+
import java.util.HashMap;
30+
import org.junit.jupiter.api.BeforeAll;
31+
import org.junit.jupiter.api.Disabled;
32+
import org.junit.jupiter.api.Test;
33+
import org.junit.jupiter.params.ParameterizedTest;
34+
import org.junit.jupiter.params.provider.EmptySource;
35+
import org.junit.jupiter.params.provider.NullSource;
36+
import org.mockito.quality.Strictness;
37+
import org.parosproxy.paros.Constant;
38+
import org.parosproxy.paros.core.scanner.Alert;
39+
import org.zaproxy.addon.llm.communication.Confidence;
40+
import org.zaproxy.zap.testutils.TestUtils;
41+
import org.zaproxy.zap.utils.I18N;
42+
43+
/** Unit test for {@link LlmCommunicationServiceUnitTest}. */
44+
class LlmCommunicationServiceUnitTest extends TestUtils {
45+
46+
@BeforeAll
47+
static void beforeAll() {
48+
Constant.messages = mock(I18N.class);
49+
}
50+
51+
@Disabled
52+
@ParameterizedTest
53+
@NullSource
54+
@EmptySource
55+
void shouldUseTwoParamReviewMethodWhenNoOtherInfo(String otherInfo) {
56+
// Given
57+
LlmAssistant assistant =
58+
mock(LlmAssistant.class, withSettings().strictness(Strictness.LENIENT));
59+
LlmCommunicationService service = new LlmCommunicationService();
60+
service.setAssistant(assistant);
61+
Alert alert = new Alert(-1);
62+
alert.setDescription("desc");
63+
alert.setEvidence("evidence");
64+
alert.setConfidence(Alert.CONFIDENCE_MEDIUM);
65+
alert.setOtherInfo(otherInfo);
66+
alert.setTags(new HashMap<String, String>());
67+
given(assistant.review(anyString(), anyString()))
68+
.willReturn(new Confidence(Alert.CONFIDENCE_MEDIUM, "explanation"));
69+
given(assistant.review(anyString(), anyString(), anyString()))
70+
.willReturn(new Confidence(Alert.CONFIDENCE_MEDIUM, "explanation"));
71+
72+
// When
73+
service.reviewAlert(alert);
74+
// Then
75+
verify(assistant).review(anyString(), anyString());
76+
verify(assistant, never()).review(anyString(), anyString(), anyString());
77+
}
78+
79+
@Disabled
80+
@Test
81+
void shouldUseThreeParamReviewMethodWhenNoOtherInfo() {
82+
// Given
83+
LlmAssistant assistant =
84+
mock(LlmAssistant.class, withSettings().strictness(Strictness.LENIENT));
85+
LlmCommunicationService service = new LlmCommunicationService();
86+
service.setAssistant(assistant);
87+
Alert alert = new Alert(-1);
88+
alert.setDescription("desc");
89+
alert.setEvidence("evidence");
90+
alert.setConfidence(Alert.CONFIDENCE_MEDIUM);
91+
alert.setOtherInfo("other info");
92+
alert.setTags(new HashMap<String, String>());
93+
given(assistant.review(anyString(), anyString()))
94+
.willReturn(new Confidence(Alert.CONFIDENCE_MEDIUM, "explanation"));
95+
given(assistant.review(anyString(), anyString(), anyString()))
96+
.willReturn(new Confidence(Alert.CONFIDENCE_MEDIUM, "explanation"));
97+
98+
// When
99+
service.reviewAlert(alert);
100+
// Then
101+
verify(assistant).review(anyString(), anyString(), anyString());
102+
verify(assistant, never()).review(anyString(), anyString());
103+
}
104+
}

0 commit comments

Comments
 (0)