From 635bd4c84958c71cb085db4eb01e9132169ed766 Mon Sep 17 00:00:00 2001 From: Ark2307 Date: Tue, 10 Dec 2024 16:52:49 +0530 Subject: [PATCH 1/2] Adding refresh auth mechanism on auth expiration --- .../java/com/akto/testing/TestExecutor.java | 33 +++++++--- .../testing/workflow_node_executor/Utils.java | 45 ++++++++++++++ .../CookieExpireFilter.java | 62 ++++++++++--------- 3 files changed, 104 insertions(+), 36 deletions(-) diff --git a/apps/testing/src/main/java/com/akto/testing/TestExecutor.java b/apps/testing/src/main/java/com/akto/testing/TestExecutor.java index 5ea6f719f5..c40d6ed941 100644 --- a/apps/testing/src/main/java/com/akto/testing/TestExecutor.java +++ b/apps/testing/src/main/java/com/akto/testing/TestExecutor.java @@ -37,7 +37,6 @@ import com.akto.test_editor.filter.data_operands_impl.ValidationResult; import com.akto.testing.yaml_tests.YamlTestTemplate; import com.akto.testing_issues.TestingIssuesHandler; -import com.akto.testing_utils.TestingUtils; import com.akto.usage.UsageMetricCalculator; import com.akto.util.Constants; import com.akto.util.JSONUtils; @@ -45,8 +44,6 @@ import com.akto.util.enums.LoginFlowEnums; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.gson.Gson; -import com.mongodb.MongoInterruptedException; -import org.apache.commons.lang3.StringUtils; import com.mongodb.BasicDBObject; import com.mongodb.client.model.*; @@ -73,6 +70,13 @@ public class TestExecutor { public static final String REQUEST_HOUR = "requestHour"; public static final String COUNT = "count"; public static final int ALLOWED_REQUEST_PER_HOUR = 100; + + private static int expiryTimeOfAuthToken = -1; + + public static synchronized void setExpiryTimeOfAuthToken(int newExpiryTime) { + expiryTimeOfAuthToken = newExpiryTime; + } + public void init(TestingRun testingRun, ObjectId summaryId, SyncLimit syncLimit) { if (testingRun.getTestIdConfig() != 1) { apiWiseInit(testingRun, summaryId, false, new ArrayList<>(), syncLimit); @@ -253,7 +257,8 @@ public void apiWiseInit(TestingRun testingRun, ObjectId summaryId, boolean debug () -> startWithLatch(apiInfoKey, testingRun.getTestIdConfig(), testingRun.getId(), testingRun.getTestingRunConfig(), testingUtil, summaryId, - accountId, latch, now, maxRunTime, testConfigMap, testingRun, subCategoryEndpointMap, apiInfoKeyToHostMap, debug, testLogs, syncLimit)); + accountId, latch, now, maxRunTime, testConfigMap, testingRun, subCategoryEndpointMap, + apiInfoKeyToHostMap, debug, testLogs, syncLimit, authMechanism)); futureTestingRunResults.add(future); } catch (Exception e) { loggerMaker.errorAndAddToDb("Error in API " + apiInfoKey + " : " + e.getMessage(), LogDb.TESTING); @@ -477,13 +482,13 @@ public Void startWithLatch( TestingUtil testingUtil, ObjectId testRunResultSummaryId, int accountId, CountDownLatch latch, int startTime, int maxRunTime, Map testConfigMap, TestingRun testingRun, ConcurrentHashMap subCategoryEndpointMap, Map apiInfoKeyToHostMap, - boolean debug, List testLogs, SyncLimit syncLimit) { + boolean debug, List testLogs, SyncLimit syncLimit, AuthMechanism authMechanism) { Context.accountId.set(accountId); loggerMaker.infoAndAddToDb("Starting test for " + apiInfoKey, LogDb.TESTING); try { - startTestNew(apiInfoKey, testRunId, testingRunConfig, testingUtil, testRunResultSummaryId, testConfigMap, subCategoryEndpointMap, apiInfoKeyToHostMap, debug, testLogs, startTime, maxRunTime, syncLimit); + startTestNew(apiInfoKey, testRunId, testingRunConfig, testingUtil, testRunResultSummaryId, testConfigMap, subCategoryEndpointMap, apiInfoKeyToHostMap, debug, testLogs, startTime, maxRunTime, syncLimit, authMechanism); } catch (Exception e) { loggerMaker.errorAndAddToDb(e, "error while running tests: " + e); } @@ -530,6 +535,12 @@ public void trim(List testingRunResults) { } } + private synchronized void checkAndUpdateAuthMechanism(int timeNow, AuthMechanism authMechanism){ + if(expiryTimeOfAuthToken != -1 && expiryTimeOfAuthToken <= timeNow){ + triggerLoginFlow(authMechanism, 3); + } + } + public void insertResultsAndMakeIssues(List testingRunResults, ObjectId testRunResultSummaryId) { int resultSize = testingRunResults.size(); if (resultSize > 0) { @@ -559,7 +570,7 @@ public void startTestNew(ApiInfo.ApiInfoKey apiInfoKey, ObjectId testRunId, TestingRunConfig testingRunConfig, TestingUtil testingUtil, ObjectId testRunResultSummaryId, Map testConfigMap, ConcurrentHashMap subCategoryEndpointMap, Map apiInfoKeyToHostMap, - boolean debug, List testLogs, int startTime, int timeToKill, SyncLimit syncLimit) { + boolean debug, List testLogs, int startTime, int timeToKill, SyncLimit syncLimit, AuthMechanism authMechanism) { List testSubCategories = testingRunConfig == null ? new ArrayList<>() : testingRunConfig.getTestSubCategoryList(); @@ -606,6 +617,14 @@ public void startTestNew(ApiInfo.ApiInfoKey apiInfoKey, ObjectId testRunId, try { if(testingRunResult==null){ + + // check and update automated auth token here + int diffTimeInMinutes = (Context.now() - startTime)/60; + if(diffTimeInMinutes != 0 && (diffTimeInMinutes % 10) == 0){ + // check for expiry in every 10 minutes + checkAndUpdateAuthMechanism(Context.now(), authMechanism); + } + testingRunResult = runTestNew(apiInfoKey,testRunId,testingUtil,testRunResultSummaryId, testConfig, testingRunConfig, debug, testLogs); } } catch (Exception e) { diff --git a/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java b/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java index 84d9b9e9df..b4222ed773 100644 --- a/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java +++ b/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java @@ -1,7 +1,11 @@ package com.akto.testing.workflow_node_executor; +import static com.akto.runtime.utils.Utils.parseCookie; +import static org.mockito.Answers.values; + import java.io.File; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -12,6 +16,9 @@ import com.akto.dto.testing.*; import com.akto.test_editor.execution.Memory; +import com.akto.test_editor.filter.data_operands_impl.CookieExpireFilter; +import com.akto.testing.TestExecutor; + import org.apache.commons.lang3.StringUtils; import org.bson.conversions.Bson; import org.json.JSONObject; @@ -24,6 +31,7 @@ import com.akto.dto.RecordedLoginFlowInput; import com.akto.dto.api_workflow.Graph; import com.akto.dto.api_workflow.Node; +import com.akto.dto.type.KeyTypes; import com.akto.dto.type.RequestTemplate; import com.akto.log.LoggerMaker; import com.akto.log.LoggerMaker.LogDb; @@ -34,6 +42,10 @@ import com.mongodb.client.model.Filters; import com.mongodb.client.model.Updates; +import io.jsonwebtoken.Claims; +import io.jsonwebtoken.Jws; +import io.jsonwebtoken.Jwts; + public class Utils { private static final LoggerMaker loggerMaker = new LoggerMaker(Utils.class, LogDb.TESTING); @@ -298,6 +310,39 @@ public static LoginFlowResponse runLoginFlow(WorkflowTest workflowTest, AuthMech return new LoginFlowResponse(responses, "auth param not found at specified path " + param.getValue(), false); } + + // checking on the value of if this is valid jwt token or valid cookie which has expiry time + String tempVal = new String(value); + if(tempVal.contains("Bearer")){ + tempVal = value.split("Bearer ")[1]; + } + if(KeyTypes.isJWT(tempVal)){ + try { + Jws jws = Jwts.parserBuilder() + .build() + .parseClaimsJws(tempVal); + + if(jws.getBody() != null && jws.getBody().getExpiration() != null && jws.getBody().getExpiration().getTime() != 0){ + int newExpiryTime = (int) jws.getBody().getExpiration().getTime(); + TestExecutor.setExpiryTimeOfAuthToken(newExpiryTime); + } + } catch (Exception e) { + e.printStackTrace(); + } + }else{ + // check if this cookie with max-age or expiry time + try { + Map cookieMap = parseCookie(Arrays.asList(value)); + int expiryTsEpoch = CookieExpireFilter.getMaxAgeFromCookie(cookieMap); + if(expiryTsEpoch > 0){ + int newExpiryTime = Context.now() + expiryTsEpoch; + TestExecutor.setExpiryTimeOfAuthToken(newExpiryTime); + } + } catch (Exception e) { + e.printStackTrace(); + } + } + param.setValue(value); } catch(Exception e) { return new LoginFlowResponse(responses, "error resolving auth param " + param.getValue(), false); diff --git a/libs/utils/src/main/java/com/akto/test_editor/filter/data_operands_impl/CookieExpireFilter.java b/libs/utils/src/main/java/com/akto/test_editor/filter/data_operands_impl/CookieExpireFilter.java index d9367cf4b2..145d49bcf3 100644 --- a/libs/utils/src/main/java/com/akto/test_editor/filter/data_operands_impl/CookieExpireFilter.java +++ b/libs/utils/src/main/java/com/akto/test_editor/filter/data_operands_impl/CookieExpireFilter.java @@ -11,34 +11,13 @@ import com.akto.dao.test_editor.TestEditorEnums; import com.akto.dto.test_editor.DataOperandFilterRequest; +import com.akto.util.Constants; import static com.akto.runtime.utils.Utils.parseCookie; public class CookieExpireFilter extends DataOperandsImpl { - - @Override - public ValidationResult isValid(DataOperandFilterRequest dataOperandFilterRequest) { - - List querySet = new ArrayList<>(); - Boolean queryVal; - String data; - try { - - querySet = (List) dataOperandFilterRequest.getQueryset(); - queryVal = (Boolean) querySet.get(0); - data = (String) dataOperandFilterRequest.getData(); - } catch(Exception e) { - return new ValidationResult(false, ValidationResult.GET_QUERYSET_CATCH_ERROR); - } - if (data == null || queryVal == null) { - return new ValidationResult(false, queryVal == null ? TestEditorEnums.DataOperands.COOKIE_EXPIRE_FILTER.name().toLowerCase() + " is not set true": "no data to be matched for validation"); - } - - Map cookieMap = parseCookie(Arrays.asList(data)); - - boolean result = queryVal; - boolean res = false; + public static int getMaxAgeFromCookie(Map cookieMap){ if (cookieMap.containsKey("Max-Age") || cookieMap.containsKey("max-age")) { int maxAge; if (cookieMap.containsKey("Max-Age")) { @@ -46,9 +25,7 @@ public ValidationResult isValid(DataOperandFilterRequest dataOperandFilterReques } else { maxAge = Integer.parseInt(cookieMap.get("max-age")); } - if (maxAge/(60*60*24) > 30) { - res = true; - } + return maxAge; } else if (cookieMap.containsKey("Expires") || cookieMap.containsKey("expires")) { String expiresTs; if (cookieMap.containsKey("Expires")) { @@ -68,10 +45,37 @@ public ValidationResult isValid(DataOperandFilterRequest dataOperandFilterReques LocalDateTime now = LocalDateTime.now(); Duration duration = Duration.between(now, dateTime); long seconds = duration.getSeconds(); - if (seconds/(60*60*24) > 30) { - res = true; - } + return (int) seconds; } + return -1; + } + + @Override + public ValidationResult isValid(DataOperandFilterRequest dataOperandFilterRequest) { + + List querySet = new ArrayList<>(); + Boolean queryVal; + String data; + try { + + querySet = (List) dataOperandFilterRequest.getQueryset(); + queryVal = (Boolean) querySet.get(0); + data = (String) dataOperandFilterRequest.getData(); + } catch(Exception e) { + return new ValidationResult(false, ValidationResult.GET_QUERYSET_CATCH_ERROR); + } + + if (data == null || queryVal == null) { + return new ValidationResult(false, queryVal == null ? TestEditorEnums.DataOperands.COOKIE_EXPIRE_FILTER.name().toLowerCase() + " is not set true": "no data to be matched for validation"); + } + + Map cookieMap = parseCookie(Arrays.asList(data)); + + boolean result = queryVal; + boolean res = false; + + int maxAgeOfCookieTs = getMaxAgeFromCookie(cookieMap); + res = maxAgeOfCookieTs/(Constants.ONE_MONTH_TIMESTAMP) > 1; if (result == res) { return new ValidationResult(true, result? TestEditorEnums.DataOperands.COOKIE_EXPIRE_FILTER.name().toLowerCase() + ": true passed because cookie:"+ data+" expired": TestEditorEnums.DataOperands.COOKIE_EXPIRE_FILTER.name().toLowerCase() + ": false passed because cookie:"+ data+" not expired"); From d130d4838024a2fd84daf5ff4ce634a9232b9f4e Mon Sep 17 00:00:00 2001 From: Ark2307 Date: Tue, 10 Dec 2024 17:27:19 +0530 Subject: [PATCH 2/2] fixing expiry time calculation from jwt --- .../testing/workflow_node_executor/Utils.java | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java b/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java index b4222ed773..052c9863e2 100644 --- a/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java +++ b/apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java @@ -6,6 +6,7 @@ import java.io.File; import java.util.ArrayList; import java.util.Arrays; +import java.util.Base64; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -318,13 +319,17 @@ public static LoginFlowResponse runLoginFlow(WorkflowTest workflowTest, AuthMech } if(KeyTypes.isJWT(tempVal)){ try { - Jws jws = Jwts.parserBuilder() - .build() - .parseClaimsJws(tempVal); - - if(jws.getBody() != null && jws.getBody().getExpiration() != null && jws.getBody().getExpiration().getTime() != 0){ - int newExpiryTime = (int) jws.getBody().getExpiration().getTime(); + String[] parts = tempVal.split("\\."); + if (parts.length != 3) { + throw new IllegalArgumentException("Invalid JWT token format"); + } + String payload = new String(Base64.getUrlDecoder().decode(parts[1])); + JSONObject payloadJson = new JSONObject(payload); + if (payloadJson.has("exp")) { + int newExpiryTime = payloadJson.getInt("exp"); TestExecutor.setExpiryTimeOfAuthToken(newExpiryTime); + } else { + throw new IllegalArgumentException("JWT does not have an 'exp' claim"); } } catch (Exception e) { e.printStackTrace();