SkipVulnerabilityScanDeciderTest.java

package vulnerabilities;

import dev.aikido.agent_api.background.Endpoint;
import dev.aikido.agent_api.context.ContextObject;
import dev.aikido.agent_api.storage.BypassedContextStore;
import dev.aikido.agent_api.vulnerabilities.SkipVulnerabilityScanDecider;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import utils.EmptyAPIResponses;
import utils.EmptySampleContextObject;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static org.junit.jupiter.api.Assertions.*;

public class SkipVulnerabilityScanDeciderTest {
    @BeforeEach
    public void setUp() {
        BypassedContextStore.clear();
    }

    @AfterEach
    public void tearDown() {
        BypassedContextStore.clear();
    }

    private List<Endpoint> createEndpoints(boolean protectionForcedOff1, boolean protectionForcedOff2) {
        List<Endpoint> endpoints = new ArrayList<>();
        endpoints.add(new Endpoint("POST", "/api/login", 3, 1000, Collections.emptyList(), false, protectionForcedOff1, true));
        endpoints.add(new Endpoint("POST", "/api/*", 1, 1000, Collections.emptyList(), false, protectionForcedOff2, true));
        endpoints.add(new Endpoint("GET", "/", 3, 1000, Collections.emptyList(), false, false, false));
        return endpoints;
    }

    private List<Endpoint> createEndpointsWildcardMethod(boolean protectionForcedOff1, boolean protectionForcedOff2) {
        List<Endpoint> endpoints = new ArrayList<>();
        endpoints.add(new Endpoint("*", "/api/login", 3, 1000, Collections.emptyList(), false, protectionForcedOff1, true));
        endpoints.add(new Endpoint("*", "/api/*", 1, 1000, Collections.emptyList(), false, protectionForcedOff2, true));
        endpoints.add(new Endpoint("GET", "/", 3, 1000, Collections.emptyList(), false, false, false));
        return endpoints;
    }

    @Test
    public void testShouldSkipVulnerabilityScan_NullContext() {
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(null));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_ProtectionForcedOff_1() {
        // Update the ServiceConfiguration with endpoints
        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpoints(true, false));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "GET")
        ));

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login2", "POST")
        ));

        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpoints(false, false));
        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_ProtectionForcedOff_1_WildCard() {
        // Update the ServiceConfiguration with endpoints
        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpointsWildcardMethod(true, false));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login2", "POST")
        ));

        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpointsWildcardMethod(false, false));
        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_ProtectionForcedOff_2() {
        // Update the ServiceConfiguration with endpoints
        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpoints(false, true));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "GET")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login2", "POST")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/", "POST")
        ));

        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpoints(true, true));
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "GET")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login2", "POST")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/", "POST")
        ));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_ProtectionForcedOff_WildcardMethod() {
        // Update the ServiceConfiguration with endpoints
        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpointsWildcardMethod(false, true));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login2", "POST")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/", "POST")
        ));

        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpointsWildcardMethod(true, true));
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login2", "POST")
        ));

        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/", "POST")
        ));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_NoConditionsMet() {
        // Update the ServiceConfiguration with endpoints
        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpoints(false, false));

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(
                new EmptySampleContextObject("", "/api/login", "POST")
        ));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_BypassedIp_NullContext() {
        BypassedContextStore.setBypassed(true);
        // Even with defaultIfNoContext=false (the Stored SSRF path), a bypassed IP must skip.
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(null, false));
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(null, true));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_BypassedIp_WithContext() {
        EmptyAPIResponses.setEmptyConfigWithEndpointList(createEndpoints(false, false));
        ContextObject ctx = new EmptySampleContextObject("", "/api/login", "POST");
        // Without bypass flag this context would return false (no forced protection off).
        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(ctx));

        BypassedContextStore.setBypassed(true);
        ContextObject freshCtx = new EmptySampleContextObject("", "/api/login", "POST");
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(freshCtx));
    }

    @Test
    public void testShouldSkipVulnerabilityScan_NotBypassed_NullContext() {
        // Sanity check: default behavior unchanged when bypass flag is not set.
        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(null, false));
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(null, true));
    }

    @Test
    public void testUsesCache() {
        ContextObject ctx = new EmptySampleContextObject("", "/api/login", "POST");
        assertTrue(ctx.getForcedProtectionOff().isEmpty());

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(ctx));
        assertTrue(ctx.getForcedProtectionOff().isPresent());
        assertFalse(ctx.getForcedProtectionOff().get());

        assertFalse(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(ctx));

        ctx.setForcedProtectionOff(true);
        assertTrue(SkipVulnerabilityScanDecider.shouldSkipVulnerabilityScan(ctx));
    }
}