Skip to content

Commit d6fb9ad

Browse files
seranseran
committed
more changes to va
1 parent 260344a commit d6fb9ad

4 files changed

Lines changed: 61 additions & 53 deletions

File tree

‎core/src/main/kotlin/org/evomaster/core/languagemodel/LanguageModelConnector.kt‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import org.evomaster.core.languagemodel.data.OllamaRequest
88
import org.evomaster.core.languagemodel.data.OllamaResponse
99
import org.evomaster.core.languagemodel.data.Prompt
1010
import org.evomaster.core.logging.LoggingUtil
11-
import org.slf4j.LoggerFactory
1211
import java.io.DataOutputStream
1312
import java.net.HttpURLConnection
1413
import java.net.URL
@@ -27,13 +26,10 @@ import kotlin.math.min
2726
*/
2827
class LanguageModelConnector {
2928

30-
companion object {
31-
private val log = LoggerFactory.getLogger(LanguageModelConnector::class.java)
32-
}
3329

3430
@PostConstruct
3531
fun initialise() {
36-
log.debug("Initializing {}", LanguageModelConnector::class.simpleName)
32+
LoggingUtil.getInfoLogger().debug("Initializing {}", LanguageModelConnector::class.simpleName)
3733
}
3834

3935
@Inject
@@ -105,6 +101,8 @@ class LanguageModelConnector {
105101
return null
106102
}
107103

104+
LoggingUtil.getInfoLogger().info("Querying for prompt {}", prompt)
105+
108106
validatePrompt(prompt)
109107

110108
return makeQuery(prompt)
@@ -128,6 +126,8 @@ class LanguageModelConnector {
128126

129127
val response = call(languageModelServerURL, requestBody)
130128

129+
LoggingUtil.getInfoLogger().info("Answer for prompt {} is {}", prompt, response)
130+
131131
if (!id.isNullOrBlank()) {
132132
if (prompts.contains(id)) {
133133
val existingPrompt = prompts[id]
@@ -167,7 +167,7 @@ class LanguageModelConnector {
167167
writer.close()
168168

169169
if (connection.responseCode != HttpURLConnection.HTTP_OK) {
170-
LoggingUtil.Companion.uniqueWarn(log, "Failed to connect to language model server")
170+
LoggingUtil.getInfoLogger().warn("Failed to connect to language model server")
171171
return null
172172
}
173173

@@ -178,7 +178,7 @@ class LanguageModelConnector {
178178

179179
return response.response
180180
} catch (e: Exception) {
181-
LoggingUtil.Companion.uniqueWarn(log, "Failed to connect to language model server: ${e.message}")
181+
LoggingUtil.getInfoLogger().warn("Failed to connect to language model server: ${e.message}")
182182

183183
return null
184184
}

‎core/src/main/kotlin/org/evomaster/core/problem/security/VulnerabilityAnalyser.kt‎

Lines changed: 21 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,14 @@ import org.evomaster.core.languagemodel.LanguageModelConnector
66
import org.evomaster.core.problem.api.param.Param
77
import org.evomaster.core.problem.rest.data.RestCallAction
88
import org.evomaster.core.problem.rest.data.RestIndividual
9-
import org.evomaster.core.problem.rest.service.SecurityRest
109
import org.evomaster.core.problem.security.data.ParamMapping
1110
import org.evomaster.core.problem.security.verifiers.SSRFVulnerabilityVerifier
1211
import org.evomaster.core.search.EvaluatedIndividual
1312
import org.evomaster.core.search.Solution
1413
import org.evomaster.core.search.service.Archive
15-
import org.slf4j.Logger
16-
import org.slf4j.LoggerFactory
1714

1815
class VulnerabilityAnalyser {
1916

20-
companion object {
21-
private val log: Logger = LoggerFactory.getLogger(SecurityRest::class.java)
22-
}
23-
2417
@Inject
2518
private lateinit var config: EMConfig
2619

@@ -34,9 +27,12 @@ class VulnerabilityAnalyser {
3427
private lateinit var languageModelConnector: LanguageModelConnector
3528

3629

30+
/**
31+
* TODO: This can be generalized later to accommodate more vulnerability types
32+
*/
3733
private lateinit var ssrfVerifier: SSRFVulnerabilityVerifier
3834

39-
private var paramMapping: MutableMap<String, ParamMapping> = mutableMapOf()
35+
private var actionVulnerabilityMapping: MutableMap<String, MutableMap<String, ParamMapping>> = mutableMapOf()
4036

4137
/**
4238
* Individuals in the solution.
@@ -114,27 +110,29 @@ class VulnerabilityAnalyser {
114110
individualsInSolution.forEach { evaluatedIndividual ->
115111
val ea = evaluatedIndividual.evaluatedMainActions()[0].action
116112
if (ea is RestCallAction) {
117-
val descriptions = extractBodyParameterDescriptions(ea.getName(), ea.parameters)
118-
paramMapping = descriptions
113+
val descriptions : MutableMap<String, ParamMapping> = extractBodyParameterDescriptions(ea.getName(), ea.parameters)
119114

120-
descriptions.forEach { name, description ->
121-
if (!description.description.isNullOrBlank()) {
122-
val id = languageModelConnector.queryAsync(
123-
getPromptWithNameAndDescription(
115+
descriptions.forEach { paramName, description ->
116+
val answer = if (!description.description.isNullOrBlank()) {
117+
languageModelConnector.query(
118+
ssrfVerifier.getPromptWithNameAndDescription(
124119
description.name,
125120
description.description
126121
)
127122
)
128-
paramMapping[name]?.promptId = id
129123
} else {
130-
val id = languageModelConnector.queryAsync(
131-
getPromptWithNameOnly(
124+
languageModelConnector.query(
125+
ssrfVerifier.getPromptWithNameOnly(
132126
description.name
133127
)
134128
)
135-
paramMapping[name]?.promptId = id
136129
}
130+
131+
if (answer == SSRFVulnerabilityVerifier.SSRF_PROMPT_ANSWER_FOR_POSSIBILITY)
132+
description.addVulnerabilityClass(VulnerabilityClass.SSRF)
137133
}
134+
135+
actionVulnerabilityMapping[ea.getName()] = descriptions
138136
}
139137
}
140138
}
@@ -175,7 +173,10 @@ class VulnerabilityAnalyser {
175173
// TODO: Run the tests again with generated value
176174
individualsInSolution.forEach { evaluatedIndividual ->
177175
val ea = evaluatedIndividual.evaluatedMainActions()[0].action as RestCallAction
178-
val link = ssrfVerifier.generateLink(ea.getName())
176+
// Execute only the action which marked as a possible candidate
177+
if (actionVulnerabilityMapping.containsKey(ea.getName())) {
178+
val link = ssrfVerifier.generateLink(ea.getName())
179+
}
179180
}
180181
}
181182
}
@@ -190,35 +191,12 @@ class VulnerabilityAnalyser {
190191
val ea = eI.evaluatedMainActions()[0].action as RestCallAction
191192
val result = ssrfVerifier.verify(ea.getName())
192193

193-
if (result == true) {
194+
if (result) {
194195
// TODO: Mark as vulnerable
195196
}
196197
}
197198
}
198199
}
199200

200-
/**
201-
* This prompt is for SSRF.
202-
* TODO: In the future, this can be extended to for other classes by using a
203-
* custom large language model.
204-
*/
205-
private fun getPromptWithNameAndDescription(name: String, description: String): String {
206-
return """
207-
Consider the word \"${name}\" and the description as \"${description}\", used as a name identifier
208-
for a parameter inside a OpenAPI/Swagger schema. Would it likely represent a URL value ? give me
209-
answer as boolean with TRUE or FALSE"""
210-
}
211-
212-
/**
213-
* This prompt is for SSRF.
214-
* TODO: In the future, this can be extended to for other classes by using a
215-
* custom large language model.
216-
*/
217-
private fun getPromptWithNameOnly(name: String): String {
218-
return """
219-
Consider the word \"${name}\", used as a name identifier
220-
for a parameter inside a OpenAPI/Swagger schema. Would it likely represent a URL value ? give me
221-
answer as boolean with TRUE or FALSE"""
222-
}
223201

224202
}

‎core/src/main/kotlin/org/evomaster/core/problem/security/data/ParamMapping.kt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,8 @@ class ParamMapping (
1111

1212
var vulnerabilityClasses: MutableList<VulnerabilityClass> = mutableListOf()
1313

14+
fun addVulnerabilityClass(vulnerabilityClass: VulnerabilityClass) {
15+
vulnerabilityClasses.add(vulnerabilityClass)
16+
}
17+
1418
}

‎core/src/main/kotlin/org/evomaster/core/problem/security/verifiers/SSRFVulnerabilityVerifier.kt‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,6 @@ package org.evomaster.core.problem.security.verifiers
22

33
import com.github.tomakehurst.wiremock.WireMockServer
44
import com.github.tomakehurst.wiremock.client.WireMock
5-
import com.github.tomakehurst.wiremock.client.WireMock.containing
6-
import com.github.tomakehurst.wiremock.client.WireMock.findStubsByMetadata
7-
import com.github.tomakehurst.wiremock.client.WireMock.matchingJsonPath
85
import com.github.tomakehurst.wiremock.common.Metadata.metadata
96
import com.github.tomakehurst.wiremock.core.WireMockConfiguration
107
import com.github.tomakehurst.wiremock.extension.responsetemplating.ResponseTemplateTransformer
@@ -14,6 +11,10 @@ import java.util.UUID
1411

1512
class SSRFVulnerabilityVerifier : VulnerabilityVerifier() {
1613

14+
companion object {
15+
const val SSRF_PROMPT_ANSWER_FOR_POSSIBILITY = "TRUE"
16+
}
17+
1718
private var wireMockServer: WireMockServer? = null
1819

1920
private var traceTokens: MutableMap<String, String> = mutableMapOf()
@@ -81,4 +82,29 @@ class SSRFVulnerabilityVerifier : VulnerabilityVerifier() {
8182
wireMockServer?.resetAll()
8283
traceTokens.clear()
8384
}
85+
86+
87+
/**
88+
* This prompt is for SSRF.
89+
* TODO: In the future, this can be extended to for other classes by using a
90+
* custom large language model.
91+
*/
92+
fun getPromptWithNameAndDescription(name: String, description: String): String {
93+
return """
94+
Consider the word "${name}" and the description as "${description}", used as a name identifier
95+
for a parameter inside a OpenAPI/Swagger schema. Would it likely represent a URL value ? give me
96+
answer as boolean with TRUE or FALSE"""
97+
}
98+
99+
/**
100+
* This prompt is for SSRF.
101+
* TODO: In the future, this can be extended to for other classes by using a
102+
* custom large language model.
103+
*/
104+
fun getPromptWithNameOnly(name: String): String {
105+
return """
106+
Consider the word "${name}", used as a name identifier
107+
for a parameter inside a OpenAPI/Swagger schema. Would it likely represent a URL value ? give me
108+
answer as boolean with TRUE or FALSE"""
109+
}
84110
}

0 commit comments

Comments
 (0)