Skip to content

Commit

Permalink
Merge pull request #1809 from akto-api-security/hotfix/fix_auto_auth_…
Browse files Browse the repository at this point in the history 
…refresh

Adding refresh auth mechanism on auth expiration
  • Loading branch information
@Ark2307
Ark2307 authored Dec 10, 2024
2 parents 6387bd1 + d130d48 commit 5a5bee8
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 36 deletions.
33 changes: 26 additions & 7 deletions apps/testing/src/main/java/com/akto/testing/TestExecutor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,16 +37,13 @@
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;
import com.akto.util.enums.GlobalEnums.Severity;
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.*;

Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -477,13 +482,13 @@ public Void startWithLatch(
TestingUtil testingUtil, ObjectId testRunResultSummaryId, int accountId, CountDownLatch latch, int startTime,
int maxRunTime, Map<String, TestConfig> testConfigMap, TestingRun testingRun,
ConcurrentHashMap<String, String> subCategoryEndpointMap, Map<ApiInfoKey, String> apiInfoKeyToHostMap,
boolean debug, List<TestingRunResult.TestLog> testLogs, SyncLimit syncLimit) {
boolean debug, List<TestingRunResult.TestLog> 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);
}
Expand Down Expand Up @@ -530,6 +535,12 @@ public void trim(List<TestingRunResult> testingRunResults) {
}
}

private synchronized void checkAndUpdateAuthMechanism(int timeNow, AuthMechanism authMechanism){
if(expiryTimeOfAuthToken != -1 && expiryTimeOfAuthToken <= timeNow){
triggerLoginFlow(authMechanism, 3);
}
}

public void insertResultsAndMakeIssues(List<TestingRunResult> testingRunResults, ObjectId testRunResultSummaryId) {
int resultSize = testingRunResults.size();
if (resultSize > 0) {
Expand Down Expand Up @@ -559,7 +570,7 @@ public void startTestNew(ApiInfo.ApiInfoKey apiInfoKey, ObjectId testRunId,
TestingRunConfig testingRunConfig, TestingUtil testingUtil,
ObjectId testRunResultSummaryId, Map<String, TestConfig> testConfigMap,
ConcurrentHashMap<String, String> subCategoryEndpointMap, Map<ApiInfoKey, String> apiInfoKeyToHostMap,
boolean debug, List<TestingRunResult.TestLog> testLogs, int startTime, int timeToKill, SyncLimit syncLimit) {
boolean debug, List<TestingRunResult.TestLog> testLogs, int startTime, int timeToKill, SyncLimit syncLimit, AuthMechanism authMechanism) {

List<String> testSubCategories = testingRunConfig == null ? new ArrayList<>() : testingRunConfig.getTestSubCategoryList();

Expand Down Expand Up @@ -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) {
Expand Down
50 changes: 50 additions & 0 deletions apps/testing/src/main/java/com/akto/testing/workflow_node_executor/Utils.java
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
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.Base64;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
Expand All @@ -12,6 +17,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;
Expand All @@ -24,6 +32,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;
Expand All @@ -34,6 +43,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);
Expand Down Expand Up @@ -298,6 +311,43 @@ 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 {
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();
}
}else{
// check if this cookie with max-age or expiry time
try {
Map<String,String> 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);
Expand Down
62 changes: 33 additions & 29 deletions ...tils/src/main/java/com/akto/test_editor/filter/data_operands_impl/CookieExpireFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,44 +11,21 @@

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<Boolean> querySet = new ArrayList<>();
Boolean queryVal;
String data;
try {

querySet = (List<Boolean>) 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<String,String> cookieMap = parseCookie(Arrays.asList(data));

boolean result = queryVal;
boolean res = false;
public static int getMaxAgeFromCookie(Map<String,String> cookieMap){
if (cookieMap.containsKey("Max-Age") || cookieMap.containsKey("max-age")) {
int maxAge;
if (cookieMap.containsKey("Max-Age")) {
maxAge = Integer.parseInt(cookieMap.get("Max-Age"));
} 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")) {
Expand All @@ -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<Boolean> querySet = new ArrayList<>();
Boolean queryVal;
String data;
try {

querySet = (List<Boolean>) 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<String,String> 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");
Expand Down

0 comments on commit 5a5bee8

Please sign in to comment.