diff --git a/apps/mini-testing/src/main/java/com/akto/test_editor/execution/Executor.java b/apps/mini-testing/src/main/java/com/akto/test_editor/execution/Executor.java index 4cd200b58a..2867aa71d4 100644 --- a/apps/mini-testing/src/main/java/com/akto/test_editor/execution/Executor.java +++ b/apps/mini-testing/src/main/java/com/akto/test_editor/execution/Executor.java @@ -189,7 +189,7 @@ public YamlTestResult execute(ExecutorNode node, RawApi rawApi, Map contentType = origRawApi.getRequest().getHeaders().getOrDefault("content-type", new ArrayList<>()); String contentTypeString = ""; @@ -655,7 +655,8 @@ public ExecutorSingleOperationResp runOperation(String operationType, RawApi raw } case "conversations_list": List conversationsList = (List) value; - return Operations.addHeader(rawApi, "x-agent-conversations", String.join(",", conversationsList)); + rawApi.setConversationsList(conversationsList); + return Operations.addHeader(rawApi, Constants.AKTO_AGENT_CONVERSATIONS , "0"); case "attach_file": return Operations.addHeader(rawApi, Constants.AKTO_ATTACH_FILE , key.toString()); case "add_body_param": diff --git a/apps/mini-testing/src/main/java/com/akto/testing/Main.java b/apps/mini-testing/src/main/java/com/akto/testing/Main.java index e313e1186b..6efc25ad0b 100644 --- a/apps/mini-testing/src/main/java/com/akto/testing/Main.java +++ b/apps/mini-testing/src/main/java/com/akto/testing/Main.java @@ -120,7 +120,6 @@ private static void handleTestEditorPlayground(TestingRunPlayground testingRunPl ApiInfo.ApiInfoKey infoKey = testingRunPlayground.getApiInfoKey(); List sampleData = testingRunPlayground.getSamples(); // get sample data from DB - AgenticUtils.checkAndInitializeAgent(Collections.singleton(infoKey.getApiCollectionId()), true, testConfig.getApiSelectionFilters().getNode()); List testLogs = new ArrayList<>(); Map> sampleDataMap = new HashMap<>(); diff --git a/apps/mini-testing/src/main/java/com/akto/testing/TestExecutor.java b/apps/mini-testing/src/main/java/com/akto/testing/TestExecutor.java index 279bd46654..5429574a94 100644 --- a/apps/mini-testing/src/main/java/com/akto/testing/TestExecutor.java +++ b/apps/mini-testing/src/main/java/com/akto/testing/TestExecutor.java @@ -277,7 +277,6 @@ public void apiWiseInit(TestingRun testingRun, ObjectId summaryId, boolean debug // create count down latch to know when inserting kafka records are completed. CountDownLatch latch = new CountDownLatch(apiInfoKeyList.size()); - AgenticUtils.checkAndInitializeAgent(Main.extractApiCollectionIds(apiInfoKeyList), false, null); int tempRunTime = 10 * 60; if(!Constants.IS_NEW_TESTING_ENABLED){ tempRunTime = testingRun.getTestRunTime() <= 0 ? 30*60 : testingRun.getTestRunTime(); diff --git a/apps/mini-testing/src/main/java/com/akto/testing/workflow_node_executor/YamlNodeExecutor.java b/apps/mini-testing/src/main/java/com/akto/testing/workflow_node_executor/YamlNodeExecutor.java index 0cb15ea06e..1cd7675e60 100644 --- a/apps/mini-testing/src/main/java/com/akto/testing/workflow_node_executor/YamlNodeExecutor.java +++ b/apps/mini-testing/src/main/java/com/akto/testing/workflow_node_executor/YamlNodeExecutor.java @@ -157,7 +157,7 @@ public NodeResult processNode(Node node, Map varMap, Boolean all TestResult res = null; if (AgentClient.isRawApiValidForAgenticTest(testReq)) { // execute agentic test here - res = agentClient.executeAgenticTest(testReq); + res = agentClient.executeAgenticTest(testReq, yamlNodeDetails.getApiCollectionId()); }else{ tsBeforeReq = Context.nowInMillis(); testResponse = ApiExecutor.sendRequest(testReq.getRequest(), followRedirect, testingRunConfig, debug, testLogs, Main.SKIP_SSRF_CHECK); diff --git a/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java b/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java index de8d9a9a00..a7456ef29d 100644 --- a/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java +++ b/apps/testing/src/test/java/com/akto/test_editor/execution/TestVariableResolver.java @@ -54,7 +54,7 @@ public void testResolveWordListVar() { varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/")); String key = "${changed_body_value}${specialCharacters}${randomVar}"; - List result = VariableResolver.resolveWordListVar(key, varMap); + List result = VariableResolver.resolveWordListVar(key, varMap); assertEquals(Arrays.asList("${changed_body_value}.${randomVar}", "${changed_body_value}$${randomVar}", "${changed_body_value}/${randomVar}"), result); @@ -84,7 +84,29 @@ public void testResolveWordListVar() { key = "${changed_body_value}${specialCharacters}${names}"; result = VariableResolver.resolveWordListVar(key, varMap); - assertEquals(Arrays.asList("${changed_body_value}.${names}", "${changed_body_value}$${names}", "${changed_body_value}/${names}"), result); + assertEquals(Arrays.asList("${changed_body_value}..", "${changed_body_value}.$", "${changed_body_value}./", + "${changed_body_value}$.", "${changed_body_value}$$", "${changed_body_value}$/", + "${changed_body_value}/.", "${changed_body_value}/$", "${changed_body_value}//"), result); + + varMap = new HashMap<>(); + varMap.put("changed_body_value", "akto"); + varMap.put("randomVar", "random"); + varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/")); + varMap.put("wordList_names", Arrays.asList(".", "$", "/")); + key = "${changed_body_value}"; + + result = VariableResolver.resolveWordListVar(key, varMap); + assertEquals(Arrays.asList("${changed_body_value}"), result); + + varMap = new HashMap<>(); + varMap.put("changed_body_value", "akto"); + varMap.put("randomVar", "random"); + varMap.put("wordList_specialCharacters", Arrays.asList(".", "$", "/")); + varMap.put("wordList_names", Arrays.asList(".", "$", "/")); + key = "nothing here"; + + result = VariableResolver.resolveWordListVar(key, varMap); + assertEquals(Arrays.asList("nothing here"), result); } } diff --git a/libs/dao/src/main/java/com/akto/dto/ApiCollection.java b/libs/dao/src/main/java/com/akto/dto/ApiCollection.java index 58ec4fc18a..d111ed642f 100644 --- a/libs/dao/src/main/java/com/akto/dto/ApiCollection.java +++ b/libs/dao/src/main/java/com/akto/dto/ApiCollection.java @@ -389,4 +389,11 @@ public boolean isGenAICollection() { return false; } + public boolean isGuardRailCollection() { + if (!CollectionUtils.isEmpty(this.getTagsList())) { + return this.getTagsList().stream().anyMatch(t -> Constants.AKTO_GUARD_RAIL_TAG.equals(t.getKeyName())); + } + return false; + } + } diff --git a/libs/dao/src/main/java/com/akto/dto/OriginalHttpRequest.java b/libs/dao/src/main/java/com/akto/dto/OriginalHttpRequest.java index 374686846e..aa41761849 100644 --- a/libs/dao/src/main/java/com/akto/dto/OriginalHttpRequest.java +++ b/libs/dao/src/main/java/com/akto/dto/OriginalHttpRequest.java @@ -2,19 +2,18 @@ import com.akto.dto.testing.TLSAuthParam; import com.akto.dto.type.RequestTemplate; +import com.akto.util.Constants; import com.akto.util.HttpRequestResponseUtils; import com.alibaba.fastjson2.JSON; import com.google.gson.Gson; import com.mongodb.BasicDBObject; -import java.util.stream.Collectors; import lombok.Getter; import lombok.Setter; import okhttp3.Headers; import okhttp3.HttpUrl; import java.net.URI; import java.util.*; -import org.apache.commons.lang3.StringUtils; public class OriginalHttpRequest { @@ -412,6 +411,32 @@ public Headers toOkHttpHeaders() { return builder.build(); } + public String fetchHeadersJsonString() { + Map> headersMap = this.getHeaders(); + List forbiddenHeaders = Arrays.asList("content-length", "accept-encoding"); + if (headersMap == null) + headersMap = new HashMap<>(); + headersMap.put(Constants.AKTO_IGNORE_FLAG, Collections.singletonList("0")); + Map filteredHeaders = new HashMap<>(); + for (String headerName : headersMap.keySet()) { + if (forbiddenHeaders.contains(headerName)) + continue; + if (headerName.contains(" ")) + continue; + if(headerName.startsWith(":")) continue; + List headerValueList = headersMap.get(headerName); + if (headerValueList == null || headerValueList.isEmpty()) + continue; + for (String headerValue : headerValueList) { + if (headerValue == null) + continue; + filteredHeaders.put(headerName, headerValue); + break; + } + } + return gson.toJson(filteredHeaders); + } + @Override public String toString() { return "OriginalHttpRequest{" + diff --git a/libs/dao/src/main/java/com/akto/dto/RawApi.java b/libs/dao/src/main/java/com/akto/dto/RawApi.java index 30790df2a8..fa283aacbc 100644 --- a/libs/dao/src/main/java/com/akto/dto/RawApi.java +++ b/libs/dao/src/main/java/com/akto/dto/RawApi.java @@ -10,7 +10,8 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.mongodb.BasicDBObject; -import javax.naming.LimitExceededException; +import lombok.Getter; +import lombok.Setter; public class RawApi { @@ -18,6 +19,9 @@ public class RawApi { private OriginalHttpResponse response; private String originalMessage; static ObjectMapper om = new ObjectMapper(); + @Getter + @Setter + private List conversationsList; public RawApi(OriginalHttpRequest request, OriginalHttpResponse response, String originalMessage) { this.request = request; diff --git a/libs/dao/src/main/java/com/akto/util/Constants.java b/libs/dao/src/main/java/com/akto/util/Constants.java index d004c7fcb3..d931ff2dd4 100644 --- a/libs/dao/src/main/java/com/akto/util/Constants.java +++ b/libs/dao/src/main/java/com/akto/util/Constants.java @@ -24,6 +24,7 @@ private Constants() {} public static final String AKTO_TOKEN_KEY = "x-akto-key"; public static final String AKTO_NODE_ID = "x-akto-node"; public static final String AKTO_REMOVE_AUTH= "x-akto-remove-auth"; + public static final String AKTO_AGENT_CONVERSATIONS= "x-agent-conversations"; public static final String LOCAL_KAFKA_BROKER_URL = System.getenv("KAFKA_BROKER_URL") != null ? System.getenv("KAFKA_BROKER_URL") : "localhost:29092"; // run kafka process with name kafka1 in docker public static final String TEST_RESULTS_TOPIC_NAME = "akto.test.messages"; @@ -45,6 +46,7 @@ private Constants() {} public final static String _AKTO = "AKTO"; public static final String AKTO_MCP_SERVER_TAG = "mcp-server"; public static final String AKTO_GEN_AI_TAG = "gen-ai"; + public static final String AKTO_GUARD_RAIL_TAG = "guard-rail"; public static final String AKTO_MCP_TOOLS_TAG = "mcp-tool"; public static final String AKTO_MCP_RESOURCES_TAG = "mcp-resource"; public static final String AKTO_MCP_PROMPTS_TAG = "mcp-prompt"; diff --git a/libs/dao/src/main/java/com/akto/util/enums/GlobalEnums.java b/libs/dao/src/main/java/com/akto/util/enums/GlobalEnums.java index 5812b9468c..cc3c96e941 100644 --- a/libs/dao/src/main/java/com/akto/util/enums/GlobalEnums.java +++ b/libs/dao/src/main/java/com/akto/util/enums/GlobalEnums.java @@ -21,6 +21,9 @@ public String getName() { /* Category of tests perfomred */ public enum TestCategory { + + // whenever adding a new test category, please ensure that if mcp or LLM is used, update the list of categories in TestTemplateUtils.java which will be used to fetch templates in fetchAllSubCategories + BOLA("BOLA", Severity.HIGH, "Broken Object Level Authorization (BOLA)", "BOLA"), NO_AUTH("NO_AUTH", Severity.HIGH, "Broken User Authentication (BUA)", "Broken Authentication"), BFLA("BFLA", Severity.HIGH, "Broken Function Level Authorization (BFLA)", "Broken Function Level Authorization"), @@ -44,7 +47,35 @@ public enum TestCategory { LFI("LFI", Severity.HIGH, "Local File Inclusion (LFI)", "Local File Inclusion"), XSS("XSS", Severity.HIGH, "Cross-site scripting (XSS)", "Cross-site scripting"), IIM("IIM", Severity.HIGH, "Improper Inventory Management (IIM)", "Improper Inventory Management"), - LLM("LLM",Severity.HIGH,"LLM (Large Language Models) Top 10","LLM"); + LLM("LLM",Severity.HIGH,"LLM (Large Language Models) Top 10","LLM"), + INJECT("INJECT", Severity.MEDIUM, "Injection Attacks (INJECT)", "Injection Attacks"), + INPUT("INPUT", Severity.MEDIUM, "Input Validation (INPUT)", "Input Validation"), + PROMPT_INJECTION("PROMPT_INJECTION", Severity.HIGH, "Prompt Injection", "PromptInjection"), + SENSITIVE_INFORMATION_DISCLOSURE("SENSITIVE_INFORMATION_DISCLOSURE", Severity.HIGH, "Sensitive Information Disclosure", "SensitiveInformationDisclosure"), + SUPPLY_CHAIN("SUPPLY_CHAIN", Severity.HIGH, "Supply Chain", "SupplyChain"), + DATA_AND_MODEL_POISONING("DATA_AND_MODEL_POISONING", Severity.HIGH, "DataAndModelPoisoning", "DataAndModelPoisoning"), + IMPROPER_OUTPUT_HANDLING("IMPROPER_OUTPUT_HANDLING", Severity.HIGH, "ImproperOutputHandling", "ImproperOutputHandling"), + EXCESSIVE_AGENCY("EXCESSIVE_AGENCY", Severity.HIGH, "Excessive Agency", "ExcessiveAgency"), + SYSTEM_PROMPT_LEAKAGE("SYSTEM_PROMPT_LEAKAGE", Severity.HIGH, "System Prompt Leakage", "SystemPromptLeakage"), + VECTOR_AND_EMBEDDING_WEAKNESSES("VECTOR_AND_EMBEDDING_WEAKNESSES", Severity.HIGH, "Vector and Embedding Weaknesses", "VectorAndEmbeddingWeaknesses"), + MISINFORMATION("MISINFORMATION", Severity.HIGH, "Misinformation", "Misinformation"), + UNBOUNDED_CONSUMPTION("UNBOUNDED_CONSUMPTION", Severity.HIGH, "Unbounded Consumption", "UnboundedConsumption"), + MCP_PROMPT_INJECTION("MCP_PROMPT_INJECTION", Severity.HIGH, "MCP - Prompt Injection", "MCP_PROMPT_INJECTION"), + MCP_TOOL_POISONING("MCP_TOOL_POISONING", Severity.HIGH, "MCP - Tool Poisoning", "MCP_TOOL_POISONING"), + MCP_PRIVILEGE_ABUSE("MCP_PRIVILEGE_ABUSE", Severity.HIGH, "MCP - Privilege Abuse", "MCP_PRIVILEGE_ABUSE"), + MCP_INDIRECT_PROMPT_INJECTION("MCP_INDIRECT_PROMPT_INJECTION", Severity.HIGH, "MCP - Indirect Prompt Injection", "MCP_INDIRECT_PROMPT_INJECTION"), + MCP_SENSITIVE_DATA_LEAKAGE("MCP_SENSITIVE_DATA_LEAKAGE", Severity.HIGH, "MCP - Data Leak", "MCP_SENSITIVE_DATA_LEAKAGE"), + MCP_DOS("MCP_DOS", Severity.HIGH, "MCP - Denial of Service", "MCP_DOS"), + MCP_AUTH("MCP_AUTH", Severity.HIGH, "MCP - Broken Authentication", "MCP_AUTH"), + MCP_MALICIOUS_CODE_EXECUTION("MCP_MALICIOUS_CODE_EXECUTION", Severity.HIGH, "MCP - Malicious Code Execution", "MCP_MALICIOUS_CODE_EXECUTION"), + MCP("MCP", Severity.HIGH, "Model Context Protocol (MCP) Security", "MCP"), + MCP_INPUT_VALIDATION("MCP_INPUT_VALIDATION", Severity.HIGH, "MCP - Input Validation", "MCP_INPUT_VALIDATION"), + MCP_FUNCTION_MANIPULATION("MCP_FUNCTION_MANIPULATION", Severity.CRITICAL, "Model Context Protocol (MCP) Security - Function Call Manipulation", "MCP_FUNC_MANIP"), + MCP_SECURITY("MCP_SECURITY", Severity.HIGH, "Model Context Protocol (MCP) Security", "MCP_SEC"), + AGENTIC_BUSINESS_ALIGNMENT("AGENTIC_BUSINESS_ALIGNMENT", Severity.HIGH, "Agent Business Alignment", "AGENTIC_BUSINESS_ALIGNMENT"), + AGENTIC_HALLUCINATION_AND_TRUSTWORTHINESS("AGENTIC_HALLUCINATION_AND_TRUSTWORTHINESS", Severity.HIGH, "Agent Hallucination and Trustworthiness", "AGENTIC_HALLUCINATION_AND_TRUSTWORTHINESS"), + AGENTIC_SAFETY("AGENTIC_SAFETY", Severity.HIGH, "Agent Safety", "AGENTIC_SAFETY"), + AGENTIC_SECURITY("AGENTIC_SECURITY", Severity.HIGH, "Agent Security", "AGENTIC_SECURITY"); private final String name; private final Severity severity; @@ -108,6 +139,30 @@ public enum YamlTemplateSource { CUSTOM } + public enum TemplatePlan { + FREE, STANDARD, PRO, ENTERPRISE + } + public enum TemplateFeatureAccess { + PRO_TESTS, ENTERPRISE_TESTS + } + public enum TemplateNature { + INTRUSIVE, NON_INTRUSIVE + } + public enum TemplateDuration { + SLOW, FAST + } + + public enum ENCODING_TYPE{ + BASE_64_ENCODED, JWT + } + + public enum TicketSource { + JIRA, AZURE_BOARDS, SERVICENOW + } + + public enum CONTEXT_SOURCE { + API, MCP, GEN_AI, AGENTIC + } /* ********************************************************************** */ } diff --git a/libs/utils/src/main/java/com/akto/agent/AgentClient.java b/libs/utils/src/main/java/com/akto/agent/AgentClient.java index e759aedd14..a4a5b97670 100644 --- a/libs/utils/src/main/java/com/akto/agent/AgentClient.java +++ b/libs/utils/src/main/java/com/akto/agent/AgentClient.java @@ -42,12 +42,20 @@ public AgentClient(String agentBaseUrl) { this.agentBaseUrl = agentBaseUrl.endsWith("/") ? agentBaseUrl.substring(0, agentBaseUrl.length() - 1) : agentBaseUrl; } - public TestResult executeAgenticTest(RawApi rawApi) throws Exception { + public TestResult executeAgenticTest(RawApi rawApi, int apiCollectionId) throws Exception { String conversationId = UUID.randomUUID().toString(); - String prompts = rawApi.getRequest().getHeaders().get("x-agent-conversations").get(0); - List promptsList = Arrays.asList(prompts.split(",")); + List promptsList = rawApi.getConversationsList(); String testMode = getTestModeFromRole(); + try { + /* + * the rawApi already has been modified by the testRole, + * and should have the updated auth request headers + */ + AgenticUtils.checkAndInitializeAgent(conversationId, rawApi, apiCollectionId); + } catch(Exception e){ + } + try { List conversationResults = processConversations(promptsList, conversationId, testMode); @@ -55,6 +63,7 @@ public TestResult executeAgenticTest(RawApi rawApi) throws Exception { List errors = new ArrayList<>(); TestResult testResult = new TestResult(); + // TODO: Fill in message field testResult.setMessage(null); testResult.setConversationId(conversationId); testResult.setResultTypeAgentic(true); @@ -84,7 +93,7 @@ public TestResult executeAgenticTest(RawApi rawApi) throws Exception { } } - public List processConversations(List prompts, String conversationId, String testMode) throws Exception { + private List processConversations(List prompts, String conversationId, String testMode) throws Exception { List results = new ArrayList<>(); int index = 0; int totalRequests = prompts.size(); @@ -102,7 +111,7 @@ public List processConversations(List prompts, return results; } - public AgentConversationResult sendChatRequest(String prompt, String conversationId, String testMode, boolean isLastRequest) throws Exception { + private AgentConversationResult sendChatRequest(String prompt, String conversationId, String testMode, boolean isLastRequest) throws Exception { Request request = buildOkHttpChatRequest(prompt, conversationId, isLastRequest); try (Response response = agentHttpClient.newCall(request).execute()) { @@ -172,12 +181,13 @@ private void storeConversationResults(List conversation try { dataActor.storeConversationResults(conversationResults); } catch (Exception e) { - loggerMaker.error("Error storing conversation results: " + e.getMessage()); + loggerMaker.errorAndAddToDb("Error storing conversation results: " + e.getMessage()); } } public static boolean isRawApiValidForAgenticTest(RawApi rawApi) { - return rawApi.getRequest().getHeaders().containsKey("x-agent-conversations"); + List temp = rawApi.getConversationsList(); + return (temp != null && !temp.isEmpty()); } public boolean performHealthCheck() { @@ -192,8 +202,24 @@ public boolean performHealthCheck() { } } public void initializeAgent(String sseUrl, String authorizationToken) { + Map requestBody = new HashMap<>(); + requestBody.put("sseUrl", sseUrl); + requestBody.put("authorization", authorizationToken); + initializeAgent(requestBody); + } + + public void initializeAgent(String sessionUrl, String requestHeaders, String apiRequestBody, String conversationId) { + Map requestBody = new HashMap<>(); + requestBody.put("sessionUrl", sessionUrl); + requestBody.put("requestHeaders", requestHeaders); + requestBody.put("requestBody", apiRequestBody); + requestBody.put("conversationId", conversationId); + initializeAgent(requestBody); + } + + public void initializeAgent(Map requestBody) { try { - Request initRequest = buildOkHttpInitializeRequest(sseUrl, authorizationToken); + Request initRequest = buildOkHttpInitializeRequest(requestBody); try (Response response = agentHttpClient.newCall(initRequest).execute()) { if (!response.isSuccessful()) { loggerMaker.errorAndAddToDb("Agent initialization failed with status: " + response.code()); @@ -212,18 +238,15 @@ private Request buildOkHttpHealthCheckRequest() { .build(); } - - private Request buildOkHttpInitializeRequest(String sseUrl, String authorizationToken) { - Map requestBody = new HashMap<>(); - requestBody.put("sseUrl", sseUrl); - requestBody.put("authorization", authorizationToken); - - String body; + private Request buildOkHttpInitializeRequest(Map requestBody) { + + String body = ""; try { body = objectMapper.writeValueAsString(requestBody); } catch (Exception e) { loggerMaker.errorAndAddToDb("Error serializing initialize request body: " + e.getMessage()); - body = "{\"sseUrl\":\"" + sseUrl.replace("\"", "\\\"") + "\"}"; + // TODO: Fix at sse URL. + // body = "{\"sseUrl\":\"" + sseUrl.replace("\"", "\\\"") + "\"}"; } RequestBody requestBodyObj = RequestBody.create(body, MediaType.parse("application/json")); diff --git a/libs/utils/src/main/java/com/akto/agent/AgenticUtils.java b/libs/utils/src/main/java/com/akto/agent/AgenticUtils.java index b96b24ba8f..b16238db87 100644 --- a/libs/utils/src/main/java/com/akto/agent/AgenticUtils.java +++ b/libs/utils/src/main/java/com/akto/agent/AgenticUtils.java @@ -3,20 +3,18 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Set; - import org.apache.commons.lang3.StringUtils; -import com.akto.dao.test_editor.filter.ConfigParser; import com.akto.data_actor.DataActor; import com.akto.data_actor.DataActorFactory; import com.akto.dto.ApiCollection; -import com.akto.dto.test_editor.FilterNode; import com.akto.dto.testing.AuthMechanism; import com.akto.dto.testing.AuthParam; import com.akto.dto.testing.TestRoles; import com.akto.util.Constants; import com.akto.util.enums.LoginFlowEnums; +import com.akto.dto.OriginalHttpRequest; +import com.akto.dto.RawApi; public class AgenticUtils { private static final AgentClient agentClient = new AgentClient(Constants.AGENT_BASE_URL); @@ -25,8 +23,10 @@ public class AgenticUtils { private static Map getMcpAuthPairs() { TestRoles role = dataActor.fetchTestRole("MCP_AUTHENTICATION_ROLE"); + if (role == null) { + return null; + } AuthMechanism authMechanism = role.findMatchingAuthMechanism(null); - List authParams = authMechanism.getAuthParams(); Map authPairs = new HashMap<>(); for(AuthParam authParam : authParams) { @@ -35,48 +35,37 @@ private static Map getMcpAuthPairs() { return authPairs; } - public static void checkAndInitializeAgent(Set apiCollectionIds, boolean isTestEditor, FilterNode filterNode) { - if(apiCollectionIds.isEmpty()){ - return; - } - boolean isMcpCollection = false; - if(apiCollectionIds.size() == 1){ - int apiCollectionId = 0; - for(int i : apiCollectionIds){ - apiCollectionId = i; - } - ApiCollection collection = dataActor.fetchApiCollectionMeta(apiCollectionId); - isMcpCollection = collection.isMcpCollection(); - }else{ - List apiCollections = dataActor.fetchAllApiCollectionsMeta(); - for(ApiCollection apiCollection : apiCollections){ - if(apiCollection.isMcpCollection() && apiCollectionIds.contains(apiCollection.getId())){ - isMcpCollection = true; - break; - } - } - } + public static void checkAndInitializeAgent(String conversationId, RawApi rawApi, int apiCollectionId) { + if (agentClient.performHealthCheck()) { - boolean shouldInitializeAgent = isMcpCollection; - if(isTestEditor){ - shouldInitializeAgent = shouldInitializeAgent && ConfigParser.isFilterNodeValidForAgenticTest(filterNode); - } - - if(shouldInitializeAgent){ - if(agentClient.performHealthCheck()){ - Map authPairs = getMcpAuthPairs(); - String sseUrl = authPairs.get("sseCallBackUrl"); - String authorization = authPairs.get("authorization"); - if(StringUtils.isEmpty(sseUrl) || StringUtils.isEmpty(authorization)){ - return; - } - agentClient.initializeAgent(sseUrl, authorization); + ApiCollection apiCollection = dataActor.fetchApiCollectionMeta(apiCollectionId); + boolean isMcpCollection = apiCollection.isMcpCollection(); + + if (rawApi != null && !isMcpCollection) { + OriginalHttpRequest request = rawApi.getRequest(); + String url = request.getFullUrlWithParams(); + String requestBody = request.getBody(); + String requestHeaders = request.fetchHeadersJsonString(); + agentClient.initializeAgent(url, requestHeaders, requestBody, conversationId); + return; } + Map authPairs = getMcpAuthPairs(); + String sseUrl = authPairs.get("sseCallBackUrl"); + String authorization = authPairs.get("authorization"); + if (StringUtils.isEmpty(sseUrl) || StringUtils.isEmpty(authorization)) { + return; + } + agentClient.initializeAgent(sseUrl, authorization); } } + // TODO: this is not actually being used. + // DELETE IT. public static String getTestModeFromRole() { TestRoles role = dataActor.fetchTestRole("MCP_AUTHENTICATION_ROLE"); + if (role == null) { + return "auto"; + } AuthMechanism authMechanism = role.findMatchingAuthMechanism(null); if(authMechanism == null || authMechanism.getType() == null){ return "auto"; diff --git a/libs/utils/src/main/java/com/akto/test_editor/execution/VariableResolver.java b/libs/utils/src/main/java/com/akto/test_editor/execution/VariableResolver.java index a9e89da0eb..4350e2e719 100644 --- a/libs/utils/src/main/java/com/akto/test_editor/execution/VariableResolver.java +++ b/libs/utils/src/main/java/com/akto/test_editor/execution/VariableResolver.java @@ -67,7 +67,7 @@ public static List resolveExpression(Map varMap, Object return (List) keyContext; } - if (VariableResolver.isWordListVariable(key, varMap)) { + if (key instanceof String && VariableResolver.isWordListVariable(key, varMap)) { varList = (List) VariableResolver.resolveWordListVar(key.toString(), varMap); for (int i = 0; i < varList.size(); i++) { List vals = VariableResolver.resolveExpression(varMap, varList.get(i).toString()); @@ -96,7 +96,7 @@ public static List resolveExpression(Map varMap, Object List keyList = (List) key; int index = 0; for (Object k: keyList) { - List v = VariableResolver.resolveExpression(varMap, k.toString()); + List v = VariableResolver.resolveExpression(varMap, k); if (v != null && v.size() > 0) { keyList.set(index, v.get(0).toString()); } @@ -151,7 +151,8 @@ public static List resolveExpression(Map varMap, String expressionList.remove(index); List valList = (List) val; for (int i = expressionList.size(); i < valList.size(); i++) { - Object v = param.replaceFirst("(\\$\\{[^}]*\\})", valList.get(i).toString()); + String finalVal = (valList.get(i) instanceof String) ? ((String) valList.get(i)) : valList.get(i).toString(); + Object v = param.replaceFirst("(\\$\\{[^}]*\\})", Matcher.quoteReplacement(finalVal)); expressionList.add(v); } } @@ -517,7 +518,7 @@ public static Boolean isWordListVariable(Object key, Map varMap) return false; } - public static List resolveWordListVar(String key, Map varMap) { + public static List resolveWordListVar(String key, Map varMap) { String expression = key.toString(); List wordList = new ArrayList<>(); @@ -525,6 +526,8 @@ public static List resolveWordListVar(String key, Map va Pattern pattern = Pattern.compile("\\$\\{[^}]*\\}"); Matcher matcher = pattern.matcher(expression); + List result = new ArrayList<>(); + result.add(expression); while (matcher.find()) { try { String match = matcher.group(0); @@ -536,23 +539,22 @@ public static List resolveWordListVar(String key, Map va if (isWordListVar) { wordList = (List) varMap.get("wordList_" + match); wordListKey = originalKey; - break; + List tempResult = new ArrayList<>(); + for (String temp : result) { + for (Object word : wordList) { + // TODO: handle case to use numbers as well. + String tempWord = temp.replace(wordListKey, word.toString()); + expression = tempWord; + tempResult.add(tempWord); + } + } + result = tempResult; + matcher = pattern.matcher(expression); } } catch (Exception e) { e.printStackTrace(); } } - - List result = new ArrayList<>(); - for (Object word : wordList) { - String replaced = expression.replace(wordListKey, word.toString()); - Object finalWord = replaced; - try { - finalWord = convertStringToNumber(replaced); - } catch (Exception e) { - } - result.add(finalWord); - } return result; } diff --git a/libs/utils/src/main/java/com/akto/test_editor/filter/FilterAction.java b/libs/utils/src/main/java/com/akto/test_editor/filter/FilterAction.java index cec85e57e7..f58e345420 100644 --- a/libs/utils/src/main/java/com/akto/test_editor/filter/FilterAction.java +++ b/libs/utils/src/main/java/com/akto/test_editor/filter/FilterAction.java @@ -11,10 +11,11 @@ import com.akto.data_actor.DataActorFactory; import com.akto.dto.OriginalHttpResponse; import com.akto.dto.testing.AccessMatrixUrlToRole; - +import com.akto.dao.ApiCollectionsDao; import com.akto.dao.test_editor.TestEditorEnums; import com.akto.dao.test_editor.TestEditorEnums.BodyOperator; import com.akto.dao.test_editor.TestEditorEnums.CollectionOperands; +import com.akto.dto.ApiCollection; import com.akto.dto.ApiInfo; import com.akto.dto.HttpResponseParams; import com.akto.dto.OriginalHttpRequest; @@ -261,13 +262,19 @@ public DataOperandsFilterResponse applyFilterOnTestType(FilterActionRequest filt if (rawApi == null || rawApi.getRequest() == null) { return new DataOperandsFilterResponse(false, null, null, null); } - boolean isMcpRequest = McpRequestResponseUtils.isMcpRequest(rawApi); - if (isMcpRequest) { + int apiCollectionId = filterActionRequest.getApiInfoKey().getApiCollectionId(); + ApiCollection apiCollection = dataActor.fetchApiCollectionMeta(apiCollectionId); + if (apiCollection == null) { + return new DataOperandsFilterResponse(false, null, null, null, "API collection not found"); + } + if (apiCollection.isGenAICollection() || + apiCollection.isMcpCollection() || + apiCollection.isGuardRailCollection()) { return new DataOperandsFilterResponse(true, null, null, null); - } else { - return new DataOperandsFilterResponse(false, null, null, null, "The request is not an MCP request"); } + return new DataOperandsFilterResponse(false, null, null, null, "The request is not an Agentic request"); } + public DataOperandsFilterResponse applyFilterOnRequestPayload(FilterActionRequest filterActionRequest) { RawApi rawApi = filterActionRequest.fetchRawApiBasedOnContext();