From c42d1982c31d4d64a27932c2f548a0f73599c04b Mon Sep 17 00:00:00 2001
From: tkomlodi <6026319+tkomlodi@users.noreply.github.com>
Date: Fri, 27 Oct 2023 07:43:37 -0400
Subject: [PATCH] Mocked network calls made in SSRFVulnerabilityTest.

---
 .../vulnerability/ssrf/SSRFVulnerability.java | 20 +++++++++++--------
 .../ssrf/SSRFVulnerabilityTest.java           | 16 ++++++++++++++-
 2 files changed, 27 insertions(+), 9 deletions(-)

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 cd285a66..85313a93 100644
--- a/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java
+++ b/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java
@@ -68,20 +68,24 @@ private ResponseEntity<GenericVulnerabilityResponseBean<String>> invalidUrlRespo
                                 MetaDataServiceMock.getResponse(u), true),
                         HttpStatus.OK);
             } else {
-                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);
-                }
+                return new ResponseEntity<>(
+                    new GenericVulnerabilityResponseBean<>(
+                            getResponseForURLConnection(u), 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 271eebbf..8da34771 100644
--- a/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java
+++ b/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java
@@ -3,13 +3,18 @@
 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;
@@ -36,7 +41,16 @@ class SSRFVulnerabilityTest {
 
     private static String tempFileUrl;
 
-    private final SSRFVulnerability ssrfVulnerability = new SSRFVulnerability(GIST_ID);
+    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;
+    }
 
     @BeforeAll
     static void setUp() throws IOException {