Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added unit test cases for Preflight Controller #422

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions src/main/java/org/sasanlabs/internal/utility/FileReader.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
package org.sasanlabs.internal.utility;

import org.sasanlabs.service.vulnerability.fileupload.UnrestrictedFileUpload;
import org.springframework.beans.factory.annotation.Autowired;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.InputStream;


public class FileReader {
Copy link
Member

@preetkaran20 preetkaran20 Oct 18, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you have created this class such that you can mock the readAsStream functionality right?
I think you don't need a different class, create a spy of PreflightController in test class and a method in PreflightController which you can mock using spy and everything else will work perfectly.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that was the purpose for the class.
Ok. I will make the changes.


private UnrestrictedFileUpload unrestrictedFileUpload;

public FileReader(UnrestrictedFileUpload unrestrictedFileUpload) {
this.unrestrictedFileUpload = unrestrictedFileUpload;
}


public InputStream readAsStream(String fileName) throws FileNotFoundException {
return (
new FileInputStream(
unrestrictedFileUpload.getContentDispositionRoot().toFile()
+ FrameworkConstants.SLASH
+ fileName));
}

}

Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,13 @@
import static org.sasanlabs.service.vulnerability.fileupload.UnrestrictedFileUpload.CONTENT_DISPOSITION_STATIC_FILE_LOCATION;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import org.apache.commons.io.IOUtils;
import org.sasanlabs.internal.utility.FileReader;
import org.sasanlabs.internal.utility.FrameworkConstants;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
Expand All @@ -24,24 +27,31 @@
*/
@RestController
public class PreflightController {
private UnrestrictedFileUpload unrestrictedFileUpload;
//private UnrestrictedFileUpload unrestrictedFileUpload;

public PreflightController(UnrestrictedFileUpload unrestrictedFileUpload) {
this.unrestrictedFileUpload = unrestrictedFileUpload;
private FileReader fileReader;



public PreflightController(FileReader fileReader) {
this.fileReader=fileReader;
// this.unrestrictedFileUpload = unrestrictedFileUpload;
}

@RequestMapping(
CONTENT_DISPOSITION_STATIC_FILE_LOCATION + FrameworkConstants.SLASH + "{fileName}")
public ResponseEntity<byte[]> fetchFile(@PathVariable("fileName") String fileName)
throws IOException {
InputStream inputStream =
new FileInputStream(
unrestrictedFileUpload.getContentDispositionRoot().toFile()
+ FrameworkConstants.SLASH
+ fileName);
InputStream inputStream = fileReader.readAsStream(fileName);
// new FileInputStream(
// unrestrictedFileUpload.getContentDispositionRoot().toFile()
// + FrameworkConstants.SLASH
// + fileName);
HttpHeaders httpHeaders = new HttpHeaders();
httpHeaders.add(HttpHeaders.CONTENT_DISPOSITION, "attachment");
return new ResponseEntity<byte[]>(
IOUtils.toByteArray(inputStream), httpHeaders, HttpStatus.OK);
}


}
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ public class UnrestrictedFileUpload {
private Path contentDispositionRoot;

private static final String STATIC_FILE_LOCATION = "upload";
static final String CONTENT_DISPOSITION_STATIC_FILE_LOCATION = "contentDispositionUpload";
public static final String CONTENT_DISPOSITION_STATIC_FILE_LOCATION = "contentDispositionUpload";
private static final String BASE_PATH = "static";
private static final String REQUEST_PARAMETER = "file";
private static final Random RANDOM = new Random(new Date().getTime());
Expand All @@ -69,9 +69,9 @@ public UnrestrictedFileUpload() throws IOException, URISyntaxException {
uploadDirectoryURI =
new URI(
Thread.currentThread()
.getContextClassLoader()
.getResource(BASE_PATH)
.toURI()
.getContextClassLoader()
.getResource(BASE_PATH)
.toURI()
+ FrameworkConstants.SLASH
+ STATIC_FILE_LOCATION);
root = Paths.get(uploadDirectoryURI);
Expand Down Expand Up @@ -102,14 +102,14 @@ public UnrestrictedFileUpload() throws IOException, URISyntaxException {
}

private static final ResponseEntity<GenericVulnerabilityResponseBean<String>>
genericFileUploadUtility(
Path root,
String fileName,
Supplier<Boolean> validator,
MultipartFile file,
boolean htmlEncode,
boolean isContentDisposition)
throws IOException {
genericFileUploadUtility(
Path root,
String fileName,
Supplier<Boolean> validator,
MultipartFile file,
boolean htmlEncode,
boolean isContentDisposition)
throws IOException {
if (validator.get()) {
Files.copy(
file.getInputStream(),
Expand All @@ -122,17 +122,17 @@ public UnrestrictedFileUpload() throws IOException, URISyntaxException {
FrameworkConstants.VULNERABLE_APP
+ FrameworkConstants.SLASH
+ (isContentDisposition
? CONTENT_DISPOSITION_STATIC_FILE_LOCATION
: STATIC_FILE_LOCATION)
? CONTENT_DISPOSITION_STATIC_FILE_LOCATION
: STATIC_FILE_LOCATION)
+ FrameworkConstants.SLASH
+ fileName);
} else {
uploadedFileLocation =
FrameworkConstants.VULNERABLE_APP
+ FrameworkConstants.SLASH
+ (isContentDisposition
? CONTENT_DISPOSITION_STATIC_FILE_LOCATION
: STATIC_FILE_LOCATION)
? CONTENT_DISPOSITION_STATIC_FILE_LOCATION
: STATIC_FILE_LOCATION)
+ FrameworkConstants.SLASH
+ fileName;
}
Expand All @@ -145,17 +145,17 @@ public UnrestrictedFileUpload() throws IOException, URISyntaxException {
HttpStatus.OK);
}

Path getContentDispositionRoot() {
public Path getContentDispositionRoot() {
return contentDispositionRoot;
}

// <img src="" onerror=alert(1) /> file name reflected and stored is there.
@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS,
VulnerabilityType.PATH_TRAVERSAL
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS,
VulnerabilityType.PATH_TRAVERSAL
},
description = "UNRESTRICTED_FILE_UPLOAD_NO_VALIDATION_FILE_NAME",
payload = "UNRESTRICTED_FILE_UPLOAD_PAYLOAD_LEVEL_1")
Expand All @@ -173,9 +173,9 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay
// <img src="" onerror=alert(1) /> file name reflected and stored is there.
@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
},
description = "UNRESTRICTED_FILE_UPLOAD_NO_VALIDATION_FILE_NAME",
payload = "UNRESTRICTED_FILE_UPLOAD_PAYLOAD_LEVEL_2")
Expand All @@ -193,9 +193,9 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay
// .htm extension breaks the file upload vulnerability
@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
},
description = "UNRESTRICTED_FILE_UPLOAD_IF_NOT_HTML_FILE_EXTENSION",
payload = "UNRESTRICTED_FILE_UPLOAD_PAYLOAD_LEVEL_3")
Expand All @@ -219,9 +219,9 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay

@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
},
description = "UNRESTRICTED_FILE_UPLOAD_IF_NOT_HTML_NOT_HTM_FILE_EXTENSION",
payload = "UNRESTRICTED_FILE_UPLOAD_PAYLOAD_LEVEL_4")
Expand All @@ -245,9 +245,9 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay

@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
},
description =
"UNRESTRICTED_FILE_UPLOAD_IF_NOT_HTML_NOT_HTM_FILE_EXTENSION_CASE_INSENSITIVE",
Expand Down Expand Up @@ -276,9 +276,9 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay
// WhiteList approach
@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
},
description =
"UNRESTRICTED_FILE_UPLOAD_IF_FILE_NAME_NOT_CONTAINS_.PNG_OR_.JPEG_CASE_INSENSITIVE",
Expand All @@ -305,9 +305,9 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay
// Null Byte Attack
@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS,
VulnerabilityType.REFLECTED_XSS
},
description =
"UNRESTRICTED_FILE_UPLOAD_IF_FILE_NAME_NOT_ENDS_WITH_.PNG_OR_.JPEG_CASE_INSENSITIVE_AND_FILE_NAMES_CONSIDERED_BEFORE_NULL_BYTE",
Expand Down Expand Up @@ -355,8 +355,8 @@ public ResponseEntity<GenericVulnerabilityResponseBean<String>> getVulnerablePay
// ZAP FileUpload Addon.
@AttackVector(
vulnerabilityExposed = {
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS
VulnerabilityType.UNRESTRICTED_FILE_UPLOAD,
VulnerabilityType.PERSISTENT_XSS
},
description =
"UNRESTRICTED_FILE_UPLOAD_IF_FILE_NAME_NOT_ENDS_WITH_.PNG_OR_.JPEG_CASE_INSENSITIVE")
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
package org.sasanlabs.service.vulnerability.commandInjection.fileUpload;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.sasanlabs.internal.utility.FileReader;
import org.sasanlabs.internal.utility.FrameworkConstants;
import org.sasanlabs.service.vulnerability.fileupload.PreflightController;
import org.sasanlabs.service.vulnerability.fileupload.UnrestrictedFileUpload;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.net.URISyntaxException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.sasanlabs.service.vulnerability.fileupload.UnrestrictedFileUpload.CONTENT_DISPOSITION_STATIC_FILE_LOCATION;

public class PreflightControllerTest {

@Mock
UnrestrictedFileUpload unrestrictedFileUpload;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add private.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok

@Mock
FileReader fileReader;
String fileName = "test.txt";
private PreflightController preflightController;
private Path contentDispositionRoot;

private Path textFilePath = Paths.get(fileName);


@BeforeEach
void setUp() throws IOException {
MockitoAnnotations.initMocks(this);
Files.createFile(textFilePath);
}

@AfterEach
void destroy() throws IOException {
Files.delete(textFilePath);
}


@Test
void shouldBeAbleToFetchFile() throws IOException {

//Arrange
preflightController = new PreflightController(fileReader);
contentDispositionRoot =
Paths.get(
FrameworkConstants.SLASH
+ CONTENT_DISPOSITION_STATIC_FILE_LOCATION
+ FrameworkConstants.SLASH);

when(unrestrictedFileUpload.getContentDispositionRoot()).thenReturn(contentDispositionRoot);
when(fileReader.readAsStream(fileName)).thenReturn(new ByteArrayInputStream("test".getBytes()));

//Act
ResponseEntity<byte[]> responseEntity = preflightController.fetchFile(fileName);

//Assert
assertEquals(HttpStatus.OK, responseEntity.getStatusCode());
verify(fileReader).readAsStream(fileName);
}

@Test
void shouldThrowException() throws IOException, URISyntaxException {

//Arrange
FileReader fileReader1 = new FileReader(new UnrestrictedFileUpload());
preflightController = new PreflightController(fileReader1);
String anotherFileName = "test";

//Act and Assert
assertThrows(IOException.class, () -> preflightController.fetchFile(anotherFileName));
}

}