diff --git a/src/main/java/org/rostislav/quickdrop/controller/FileViewController.java b/src/main/java/org/rostislav/quickdrop/controller/FileViewController.java index 625a341..895f655 100644 --- a/src/main/java/org/rostislav/quickdrop/controller/FileViewController.java +++ b/src/main/java/org/rostislav/quickdrop/controller/FileViewController.java @@ -81,7 +81,8 @@ public class FileViewController { } } - return fileService.downloadFile(id); + String password = (String) request.getSession().getAttribute("password"); + return fileService.downloadFile(id, password); } @PostMapping("/extend/{id}") diff --git a/src/main/java/org/rostislav/quickdrop/service/FileService.java b/src/main/java/org/rostislav/quickdrop/service/FileService.java index 3984948..8979563 100644 --- a/src/main/java/org/rostislav/quickdrop/service/FileService.java +++ b/src/main/java/org/rostislav/quickdrop/service/FileService.java @@ -15,6 +15,7 @@ import org.springframework.security.crypto.password.PasswordEncoder; import org.springframework.stereotype.Service; import org.springframework.web.multipart.MultipartFile; +import java.io.File; import java.nio.file.Files; import java.nio.file.Path; import java.time.LocalDate; @@ -22,6 +23,9 @@ import java.util.List; import java.util.Optional; import java.util.UUID; +import static org.rostislav.quickdrop.util.FileEncryptionUtils.decryptFile; +import static org.rostislav.quickdrop.util.FileEncryptionUtils.encryptFile; + @Service public class FileService { private static final Logger logger = LoggerFactory.getLogger(FileService.class); @@ -40,13 +44,15 @@ public class FileService { String uuid = UUID.randomUUID().toString(); Path path = Path.of(fileSavePath, uuid); - try { - Files.createFile(path); - Files.write(path, file.getBytes()); - logger.info("File saved: {}", path); - } catch (Exception e) { - logger.error("Error saving file: {}", e.getMessage()); - return null; + + if (fileUploadRequest.password == null) { + if (!saveUnencryptedFile(file, path)) { + return null; + } + } else { + if (!saveEncryptedFile(path, file, fileUploadRequest)) { + return null; + } } FileEntity fileEntity = new FileEntity(); @@ -64,27 +70,78 @@ public class FileService { return fileRepository.save(fileEntity); } + private boolean saveUnencryptedFile(MultipartFile file, Path path) { + try { + Files.createFile(path); + Files.write(path, file.getBytes()); + logger.info("File saved: {}", path); + } catch (Exception e) { + logger.error("Error saving file: {}", e.getMessage()); + return false; + } + return true; + } + + public boolean saveEncryptedFile(Path savePath, MultipartFile file, FileUploadRequest fileUploadRequest) { + Path tempFile; + try { + tempFile = Files.createTempFile("Unencrypted", "tmp"); + Files.write(tempFile, file.getBytes()); + logger.info("Unencrypted temp file saved: {}", tempFile); + } catch (Exception e) { + logger.error("Error saving unencrypted temp file: {}", e.getMessage()); + return false; + } + + try { + Path encryptedFile = Files.createFile(savePath); + encryptFile(tempFile.toFile(), encryptedFile.toFile(), fileUploadRequest.password); + } catch (Exception e) { + logger.error("Error encrypting file: {}", e.getMessage()); + return false; + } + + try { + Files.delete(tempFile); + } catch (Exception e) { + logger.error("Error deleting temp file: {}", e.getMessage()); + return false; + } + + return true; + } + public List getFiles() { return fileRepository.findAll(); } - public void deleteFile(Long id) { - fileRepository.deleteById(id); - } - - public ResponseEntity downloadFile(Long id) { - Optional referenceById = fileRepository.findById(id); - if (referenceById.isEmpty()) { + public ResponseEntity downloadFile(Long id, String password) { + FileEntity fileEntity = fileRepository.findById(id).orElse(null); + if (fileEntity == null) { return ResponseEntity.notFound().build(); } - Path path = Path.of(fileSavePath, referenceById.get().uuid); + Path pathOfFile = Path.of(fileSavePath, fileEntity.uuid); + Path outputFile = null; + if (fileEntity.passwordHash != null) { + try { + outputFile = File.createTempFile("Decrypted", "tmp").toPath(); + decryptFile(pathOfFile.toFile(), outputFile.toFile(), password); + } catch (Exception e) { + logger.error("Error decrypting file: {}", e.getMessage()); + return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).build(); + } + } else { + outputFile = pathOfFile; + } + try { - Resource resource = new UrlResource(path.toUri()); + Resource resource = new UrlResource(outputFile.toUri()); return ResponseEntity.ok() - .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + referenceById.get().name + "\"") + .header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + fileEntity.name + "\"") .body(resource); } catch (Exception e) { + logger.error("Error reading file: {}", e.getMessage()); return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).build(); } } diff --git a/src/main/java/org/rostislav/quickdrop/util/FileEncryptionUtils.java b/src/main/java/org/rostislav/quickdrop/util/FileEncryptionUtils.java new file mode 100644 index 0000000..d8bbd70 --- /dev/null +++ b/src/main/java/org/rostislav/quickdrop/util/FileEncryptionUtils.java @@ -0,0 +1,97 @@ +package org.rostislav.quickdrop.util; + +import javax.crypto.*; +import javax.crypto.spec.IvParameterSpec; +import javax.crypto.spec.PBEKeySpec; +import javax.crypto.spec.SecretKeySpec; +import java.io.File; +import java.io.FileInputStream; +import java.io.FileOutputStream; +import java.io.IOException; +import java.security.InvalidAlgorithmParameterException; +import java.security.InvalidKeyException; +import java.security.NoSuchAlgorithmException; +import java.security.SecureRandom; +import java.security.spec.InvalidKeySpecException; + +public class FileEncryptionUtils { + private static final String ALGORITHM = "AES/CBC/PKCS5Padding"; + private static final int ITERATION_COUNT = 65536; + private static final int KEY_LENGTH = 128; + + public static SecretKey generateKeyFromPassword(String password, byte[] salt) throws NoSuchAlgorithmException, InvalidKeySpecException { + PBEKeySpec spec = new PBEKeySpec(password.toCharArray(), salt, ITERATION_COUNT, KEY_LENGTH); + SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256"); + byte[] keyBytes = keyFactory.generateSecret(spec).getEncoded(); + return new SecretKeySpec(keyBytes, "AES"); + } + + public static byte[] generateRandomBytes() { + byte[] salt = new byte[16]; + SecureRandom random = new SecureRandom(); + random.nextBytes(salt); + return salt; + } + + public static void encryptFile(File inputFile, File outputFile, String password) throws NoSuchAlgorithmException, InvalidKeySpecException, NoSuchPaddingException, InvalidAlgorithmParameterException, InvalidKeyException, IOException, IllegalBlockSizeException, BadPaddingException { + byte[] salt = generateRandomBytes(); + SecretKey secretKey = generateKeyFromPassword(password, salt); + + + byte[] iv = generateRandomBytes(); + IvParameterSpec ivSpec = new IvParameterSpec(iv); + + Cipher cipher = Cipher.getInstance(ALGORITHM); + cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivSpec); + + try (FileOutputStream outputStream = new FileOutputStream(outputFile)) { + outputStream.write(salt); + outputStream.write(iv); + try (FileInputStream inputStream = new FileInputStream(inputFile)) { + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + byte[] encrypted = cipher.update(buffer, 0, bytesRead); + if (encrypted != null) { + outputStream.write(encrypted); + } + } + byte[] encrypted = cipher.doFinal(); + if (encrypted != null) { + outputStream.write(encrypted); + } + } + } + } + + public static void decryptFile(File inputFile, File outputFile, String password) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException, InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException, NoSuchPaddingException { + try (FileInputStream inputStream = new FileInputStream(inputFile)) { + byte[] salt = new byte[16]; + inputStream.read(salt); + + byte[] iv = new byte[16]; + inputStream.read(iv); + IvParameterSpec ivSpec = new IvParameterSpec(iv); + + SecretKey secretKey = generateKeyFromPassword(password, salt); + + Cipher cipher = Cipher.getInstance(ALGORITHM); + cipher.init(Cipher.DECRYPT_MODE, secretKey, ivSpec); + + try (FileOutputStream outputStream = new FileOutputStream(outputFile)) { + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = inputStream.read(buffer)) != -1) { + byte[] decrypted = cipher.update(buffer, 0, bytesRead); + if (decrypted != null) { + outputStream.write(decrypted); + } + } + byte[] decrypted = cipher.doFinal(); + if (decrypted != null) { + outputStream.write(decrypted); + } + } + } + } +}