basic functional code of va

Author: seran

core/src/main/kotlin/org/evomaster/core/problem/security/ParameterType.kt

@@ -0,0 +1,7 @@
+package org.evomaster.core.problem.security
+
+enum class ParameterType {
+    BODY,
+
+    HEADER
+}

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

@@ -6,6 +6,7 @@ import org.evomaster.core.languagemodel.LanguageModelConnector
 import org.evomaster.core.problem.api.param.Param
 import org.evomaster.core.problem.rest.data.RestCallAction
 import org.evomaster.core.problem.rest.data.RestIndividual
+import org.evomaster.core.problem.security.data.ActionMapping
 import org.evomaster.core.problem.security.data.ParamMapping
 import org.evomaster.core.problem.security.verifiers.SSRFVulnerabilityVerifier
 import org.evomaster.core.search.EvaluatedIndividual
@@ -32,7 +33,7 @@ class VulnerabilityAnalyser {
      */
     private lateinit var ssrfVerifier: SSRFVulnerabilityVerifier
 
-    private var actionVulnerabilityMapping: MutableMap<String, MutableMap<String, ParamMapping>> = mutableMapOf()
+    private var actionVulnerabilityMapping: MutableMap<String, ActionMapping> = mutableMapOf()
 
     /**
      * Individuals in the solution.
@@ -110,9 +111,11 @@ class VulnerabilityAnalyser {
         individualsInSolution.forEach { evaluatedIndividual ->
             val ea = evaluatedIndividual.evaluatedMainActions()[0].action
             if (ea is RestCallAction) {
-                val descriptions : MutableMap<String, ParamMapping> = extractBodyParameterDescriptions(ea.getName(), ea.parameters)
+                val actionMapping = ActionMapping(ea.getName())
+                val paramMapping: MutableMap<String, ParamMapping> =
+                    extractBodyParameters(ea.parameters)
 
-                descriptions.forEach { paramName, description ->
+                paramMapping.forEach { paramName, description ->
                     val answer = if (!description.description.isNullOrBlank()) {
                         languageModelConnector.query(
                             ssrfVerifier.getPromptWithNameAndDescription(
@@ -128,20 +131,24 @@ class VulnerabilityAnalyser {
                         )
                     }
 
-                    if (answer == SSRFVulnerabilityVerifier.SSRF_PROMPT_ANSWER_FOR_POSSIBILITY)
+                    if (answer == SSRFVulnerabilityVerifier.SSRF_PROMPT_ANSWER_FOR_POSSIBILITY) {
                         description.addVulnerabilityClass(VulnerabilityClass.SSRF)
+                        actionMapping.isVulnerable = true
+                    }
                 }
 
-                actionVulnerabilityMapping[ea.getName()] = descriptions
+                // Assign the param mapping
+                actionMapping.paramMapping = paramMapping
+
+                actionVulnerabilityMapping[ea.getName()] = actionMapping
             }
         }
     }
 
     /**
      * Extract descriptions from the Gene of body payloads.
      */
-    private fun extractBodyParameterDescriptions(
-        name: String,
+    private fun extractBodyParameters(
         parameters: List<Param>
     ): MutableMap<String, ParamMapping> {
         val output = mutableMapOf<String, ParamMapping>()
@@ -151,11 +158,15 @@ class VulnerabilityAnalyser {
                 param.seeGenes().forEach { gene ->
                     if (gene.name == "body") {
                         gene.getAllGenesInIndividual().forEach { geneInIndividual ->
-                            output[geneInIndividual.name] =
-                                ParamMapping(geneInIndividual.name, geneInIndividual.description)
+                            val pm = ParamMapping(geneInIndividual.name, geneInIndividual.description)
+                            pm.paramType = ParameterType.BODY
+                            output[geneInIndividual.name] = pm
+
                         }
                     } else {
-                        output[gene.name] = ParamMapping(gene.name, gene.description)
+                        val pm = ParamMapping(gene.name, gene.description)
+                        pm.paramType = ParameterType.HEADER
+                        output[gene.name] = pm
                     }
                 }
             }
@@ -174,8 +185,11 @@ class VulnerabilityAnalyser {
             individualsInSolution.forEach { evaluatedIndividual ->
                 val ea = evaluatedIndividual.evaluatedMainActions()[0].action as RestCallAction
                 // Execute only the action which marked as a possible candidate
-                if (actionVulnerabilityMapping.containsKey(ea.getName())) {
-                    val link = ssrfVerifier.generateLink(ea.getName())
+                if (actionVulnerabilityMapping.containsKey(ea.getName())
+                    && actionVulnerabilityMapping.getValue(ea.getName()).isVulnerable
+                ) {
+                    val honeyPotLink = ssrfVerifier.generateLink(ea.getName())
+                    // TODO: Execute the tests here
                 }
             }
         }
@@ -189,10 +203,11 @@ class VulnerabilityAnalyser {
         if (config.problemType == EMConfig.ProblemType.REST) {
             individualsInSolution.forEach { eI ->
                 val ea = eI.evaluatedMainActions()[0].action as RestCallAction
+                val am = actionVulnerabilityMapping[ea.getName()]
                 val result = ssrfVerifier.verify(ea.getName())
 
                 if (result) {
-                    // TODO: Mark as vulnerable
+                    am?.isExploitable = true
                 }
             }
         }

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

@@ -0,0 +1,21 @@
+package org.evomaster.core.problem.security.data
+
+import org.evomaster.core.problem.security.VulnerabilityClass
+
+class ActionMapping (
+    val name: String,
+
+) {
+
+    var paramMapping: Map<String, ParamMapping> = mutableMapOf()
+
+    var isVulnerable = false
+
+    var isExploitable = false
+
+    var vulnerabilityClasses: MutableList<VulnerabilityClass> = mutableListOf()
+
+    fun addVulnerabilityClass(vulnerabilityClass: VulnerabilityClass) {
+        vulnerabilityClasses.add(vulnerabilityClass)
+    }
+}

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

@@ -1,5 +1,6 @@
 package org.evomaster.core.problem.security.data
 
+import org.evomaster.core.problem.security.ParameterType
 import org.evomaster.core.problem.security.VulnerabilityClass
 
 class ParamMapping (
@@ -9,6 +10,8 @@ class ParamMapping (
 
     var promptId : String? = null
 
+    var paramType: ParameterType? = null
+
     var vulnerabilityClasses: MutableList<VulnerabilityClass> = mutableListOf()
 
     fun addVulnerabilityClass(vulnerabilityClass: VulnerabilityClass) {