Cleanup, resolve some issues from review.

SasanLabs · Sep 14, 2023 · 10dd8bb · 10dd8bb
1 parent d7ba0a5
commit 10dd8bb
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 32 deletions.
diff --git a/...org/sasanlabs/service/vulnerability/sqlInjection/UnionBasedSQLInjectionVulnerability.java b/...org/sasanlabs/service/vulnerability/sqlInjection/UnionBasedSQLInjectionVulnerability.java
@@ -45,8 +45,7 @@ public ResponseEntity<CarInformation> getCarInformationLevel1(
             @RequestParam final Map<String, String> queryParams) {
         final String id = queryParams.get("id");
         return applicationJdbcTemplate.query(
-                "select * from cars where id=" + id,
-                this::resultSetToResponse);
+                "select * from cars where id=" + id, this::resultSetToResponse);
     }
 
     @AttackVector(
@@ -61,8 +60,7 @@ public ResponseEntity<CarInformation> getCarInformationLevel2(
             @RequestParam final Map<String, String> queryParams) {
         final String id = queryParams.get("id");
         return applicationJdbcTemplate.query(
-                "select * from cars where id='" + id + "'",
-                this::resultSetToResponse);
+                "select * from cars where id='" + id + "'", this::resultSetToResponse);
     }
 
     @AttackVector(
@@ -77,8 +75,7 @@ public ResponseEntity<CarInformation> getCarInformationLevel3(
             @RequestParam final Map<String, String> queryParams) {
         final String id = queryParams.get("id").replaceAll("'", "");
         return applicationJdbcTemplate.query(
-                "select * from cars where id='" + id + "'",
-                this::resultSetToResponse);
+                "select * from cars where id='" + id + "'", this::resultSetToResponse);
     }
 
     @VulnerableAppRequestMapping(
@@ -95,7 +92,8 @@ public ResponseEntity<CarInformation> getCarInformationLevel4(
                 this::resultSetToResponse);
     }
 
-    private ResponseEntity<CarInformation> resultSetToResponse(final ResultSet rs) throws SQLException {
+    private ResponseEntity<CarInformation> resultSetToResponse(final ResultSet rs)
+            throws SQLException {
         final CarInformation carInformation = new CarInformation();
         if (rs.next()) {
             carInformation.setId(rs.getInt(1));

diff --git a/...sasanlabs/service/vulnerability/sqlInjection/UnionBasedSQLInjectionVulnerabilityTest.java b/...sasanlabs/service/vulnerability/sqlInjection/UnionBasedSQLInjectionVulnerabilityTest.java
@@ -1,68 +1,76 @@
 package org.sasanlabs.service.vulnerability.sqlInjection;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.verify;
 
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.jdbc.core.PreparedStatementSetter;
 import org.springframework.jdbc.core.ResultSetExtractor;
 
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-
-import static org.mockito.ArgumentMatchers.*;
-import static org.mockito.Mockito.*;
-
 class UnionBasedSQLInjectionVulnerabilityTest {
 
     private UnionBasedSQLInjectionVulnerability unionBasedSQLInjectionVulnerability;
     private JdbcTemplate template;
 
     @BeforeEach
     void setUp() throws IOException {
-        template = Mockito.spy(new JdbcTemplate());
+        template = Mockito.mock(JdbcTemplate.class);
 
+        // mock database
         doReturn(null)
                 .when(template)
                 .query(anyString(), (ResultSetExtractor<? extends Object>) any());
 
-        unionBasedSQLInjectionVulnerability = Mockito.spy(new UnionBasedSQLInjectionVulnerability(template));
+        unionBasedSQLInjectionVulnerability = new UnionBasedSQLInjectionVulnerability(template);
     }
 
     @Test
     void getCarInformationLevel1_ExpectParamInjected() throws IOException {
         // Act
-        final Map<String, String> params = new HashMap<>();
-        params.put("id", "1 UNION SELECT * FROM cars;");
+        final Map<String, String> params = Collections.singletonMap("id", "1 UNION SELECT * FROM cars;");
         unionBasedSQLInjectionVulnerability.getCarInformationLevel1(params);
 
         // Assert
-        verify(template).query(eq("select * from cars where id=1 UNION SELECT * FROM cars;"), (ResultSetExtractor<? extends Object>) any());
+        verify(template)
+                .query(
+                        eq("select * from cars where id=1 UNION SELECT * FROM cars;"),
+                        (ResultSetExtractor<? extends Object>) any());
     }
 
     @Test
     void getCarInformationLevel2_ExpectParamInjected() throws IOException {
         // Act
-        final Map<String, String> params = new HashMap<>();
-        params.put("id", "1' UNION SELECT * FROM cars; --");
+        final Map<String, String> params = Collections.singletonMap("id", "1' UNION SELECT * FROM cars; --");
         unionBasedSQLInjectionVulnerability.getCarInformationLevel2(params);
 
         // Assert
-        verify(template).query(eq("select * from cars where id='1' UNION SELECT * FROM cars; --'"), (ResultSetExtractor<? extends Object>) any());
+        verify(template)
+                .query(
+                        eq("select * from cars where id='1' UNION SELECT * FROM cars; --'"),
+                        (ResultSetExtractor<? extends Object>) any());
     }
 
     @Test
     void getCarInformationLevel3_ExpectParamEscaped() throws IOException {
         // Act
-        final Map<String, String> params = new HashMap<>();
-        params.put("id", "1' UNION SELECT * FROM cars; --");
+        final Map<String, String> params = Collections.singletonMap("id", "1' UNION SELECT * FROM cars; --");
         unionBasedSQLInjectionVulnerability.getCarInformationLevel3(params);
 
         // Assert
-        verify(template).query(eq("select * from cars where id='1 UNION SELECT * FROM cars; --'"), (ResultSetExtractor<? extends Object>) any());
-
+        verify(template)
+                .query(
+                        eq("select * from cars where id='1 UNION SELECT * FROM cars; --'"),
+                        (ResultSetExtractor<? extends Object>) any());
     }
 
     @Test
@@ -71,17 +79,22 @@ void getCarInformationLevel4_ExpecParamEscaped() throws IOException {
         template = Mockito.spy(new JdbcTemplate());
         doReturn(null)
                 .when(template)
-                .query(anyString(), (PreparedStatementSetter) any(), (ResultSetExtractor<? extends Object>) any());
+                .query(
+                        anyString(),
+                        (PreparedStatementSetter) any(),
+                        (ResultSetExtractor<? extends Object>) any());
 
-        unionBasedSQLInjectionVulnerability = Mockito.spy(new UnionBasedSQLInjectionVulnerability(template));
+        unionBasedSQLInjectionVulnerability = new UnionBasedSQLInjectionVulnerability(template);
 
         // Act
-        final Map<String, String> params = new HashMap<>();
-        params.put("id", "1' UNION SELECT * FROM cars; --");
+        final Map<String, String> params = Collections.singletonMap("id", "1' UNION SELECT * FROM cars; --");
         unionBasedSQLInjectionVulnerability.getCarInformationLevel4(params);
 
         // Assert
-        verify(template).query(eq("select * from cars where id=?"), (PreparedStatementSetter) any(), (ResultSetExtractor<? extends Object>) any());
-
+        verify(template)
+                .query(
+                        eq("select * from cars where id=?"),
+                        (PreparedStatementSetter) any(),
+                        (ResultSetExtractor<? extends Object>) any());
     }
 }