diff --git a/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java b/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java index aaf92b56..cd285a66 100644 --- a/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java +++ b/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java @@ -68,24 +68,20 @@ private ResponseEntity> invalidUrlRespo MetaDataServiceMock.getResponse(u), true), HttpStatus.OK); } else { - return new ResponseEntity<>( - new GenericVulnerabilityResponseBean<>( - getResponseForURLConnection(u), true), - HttpStatus.OK); + URLConnection urlConnection = u.openConnection(); + try (BufferedReader reader = + new BufferedReader(new InputStreamReader(urlConnection.getInputStream()))) { + return new ResponseEntity<>( + new GenericVulnerabilityResponseBean<>( + reader.lines().collect(Collectors.joining()), true), + HttpStatus.OK); + } } } else { return invalidUrlResponse(); } } - String getResponseForURLConnection(URL u) throws IOException { - URLConnection urlConnection = u.openConnection(); - try (BufferedReader reader = - new BufferedReader(new InputStreamReader(urlConnection.getInputStream()))) { - return reader.lines().collect(Collectors.joining()); - } - } - @AttackVector( vulnerabilityExposed = VulnerabilityType.SIMPLE_SSRF, description = "SSRF_VULNERABILITY_URL_WITHOUT_CHECK", diff --git a/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java b/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java index 4aa6a0ad..271eebbf 100644 --- a/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java +++ b/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java @@ -3,18 +3,13 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; import java.io.File; import java.io.IOException; -import java.net.URL; import java.nio.file.Files; import java.util.Collections; import java.util.stream.Stream; import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -41,18 +36,7 @@ class SSRFVulnerabilityTest { private static String tempFileUrl; - private SSRFVulnerability ssrfVulnerability; - - @BeforeEach - void each() throws IOException { - SSRFVulnerability ssrfSpy = spy(new SSRFVulnerability(GIST_ID)); - // mocks network calls - doReturn(GIST_URL_CONTENT).when(ssrfSpy).getResponseForURLConnection(eq(new URL(GIST_URL))); - doReturn(OTHER_URL_CONTENT) - .when(ssrfSpy) - .getResponseForURLConnection(eq(new URL(OTHER_URL))); - ssrfVulnerability = ssrfSpy; - } + private final SSRFVulnerability ssrfVulnerability = new SSRFVulnerability(GIST_ID); @BeforeAll static void setUp() throws IOException {