(AES) Encrypting Folder in Java

AESEncryption.java

import java.io.*;
import java.nio.file.Files;
import java.nio.file.StandardOpenOption;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AESEncryption {

	private static boolean supportsANSICode() {
		return (System.console() != null && System.getenv().get("TERM") != null);
	}

	private static final String AES_ALGORITHM = "AES/CFB/NoPadding";

	private static void print(String msg) {
		System.out.print(msg);
	}

	private static void println(String msg, String type) {
		String colorCode = "", resetCode = "\033[0m";
		
		if(type.equals("error")) {
			colorCode = "\033[0;31m";
		} else if(type.equals("success")) {
			colorCode = "\033[0;32m";
		}

		if(!supportsANSICode()) {
			colorCode = "";
		}
		
		if(colorCode.length() == 0) {
			resetCode = "";
		}

		System.out.println(colorCode + msg + resetCode);
	}
	
	private static void writeToFile(String filePath, byte[] data) throws Exception {
		File fileObj = new File(filePath);
		
		fileObj.createNewFile();
		
		FileOutputStream fout = new FileOutputStream(fileObj);
		fout.write(data);
		fout.close();
	}
	
	private static byte[] readFileContent(String filePath) throws Exception {
		File fileObj = new File(filePath);
		
		byte[] content = Files.readAllBytes(fileObj.toPath());
		
		return content;
	}
	
	private static ArrayList<File> fetchListOfFilesInFolder(ArrayList<File> filesList, File folder) throws Exception {
		File[] files = folder.listFiles();
		
		for(File file : files) {
			if(file.isDirectory()) {
				fetchListOfFilesInFolder(filesList, file);
			} else {
				filesList.add(file);
			}
		}

		return filesList;
	}
	
	private static void copySourceFolderContentToZip(ArrayList<File> filesList, File inputFolder, ZipOutputStream zipOutputStream) throws Exception {
		int inputFolderPathLength = inputFolder.getCanonicalPath().length();

		for(File file : filesList) {
			String filePath = file.getCanonicalPath();
			int filePathLength = filePath.length();
			
			String zipEntryPath = filePath.substring(inputFolderPathLength + 1, filePathLength);
			ZipEntry zipEntry = new ZipEntry(zipEntryPath);
			zipOutputStream.putNextEntry(zipEntry);
			
			byte[] fileContent = Files.readAllBytes(file.toPath());
			zipOutputStream.write(fileContent);
			zipOutputStream.closeEntry();
		}
		
		zipOutputStream.close();
	}
	
	private static void unZipFile(File zipFile, File destinationFolder) throws Exception {
		String pathToUnzip = destinationFolder.getCanonicalPath();
		ZipInputStream zipInputStream = new ZipInputStream(new FileInputStream(zipFile));
		ZipEntry zipEntry = zipInputStream.getNextEntry();
		
		while(zipEntry != null) {
			File newFile = new File(pathToUnzip, zipEntry.getName());

			// Prevent Zip Slip
			if(!newFile.getCanonicalPath().startsWith(pathToUnzip + File.separator)) {
				zipInputStream.closeEntry();
				zipInputStream.close();

				throw new IOException("Entry is outside of the target dir: " + zipEntry.getName());
			}
			
			if(zipEntry.isDirectory()) {
				if(!newFile.isDirectory()) {
					newFile.mkdirs();
				}
			} else {
				File parent = newFile.getParentFile();
				
				if(!parent.isDirectory()) {
					parent.mkdirs();
				}
				
				FileOutputStream fileOutputStream = new FileOutputStream(newFile);
				fileOutputStream.write(zipInputStream.readAllBytes());
				fileOutputStream.close();
			}

			zipEntry = zipInputStream.getNextEntry();
		}

		zipInputStream.closeEntry();
		zipInputStream.close();
	}

	private static void processEncryption(BufferedReader br) {
		try {
			print("Enter Folder Location: ");
			String folderLocation = br.readLine();

			if(folderLocation == null) {
				println("Invalid Folder Location!", "error");
				return;
			}

			File sourceFolder = new File(folderLocation);

			if( !(sourceFolder.exists() || sourceFolder.isDirectory()) ) {
				println("Invalid Folder Location or Folder not Found!", "error");
				return;
			}

			String namePrefix = sourceFolder.getName();
			File parentFolder = sourceFolder.getParentFile();
			File lockedFolder = new File(parentFolder.getCanonicalPath(), namePrefix + "_locked");
			lockedFolder.mkdir();

			File zipFile = new File(lockedFolder.getCanonicalPath(), namePrefix + ".zip");
			ZipOutputStream zipOutputStream = new ZipOutputStream(new FileOutputStream(zipFile));
			
			ArrayList<File> filesList = new ArrayList<File>();
			fetchListOfFilesInFolder(filesList, sourceFolder);

			copySourceFolderContentToZip(filesList, sourceFolder, zipOutputStream);

			// Generate 256 bit Key
			KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
			keyGenerator.init(256);
			SecretKey key = keyGenerator.generateKey();

			// Generate 128 bit IV (Initialization Vector)
			byte[] ivBytes = new byte[16];
			new SecureRandom().nextBytes(ivBytes);
			IvParameterSpec iv = new IvParameterSpec(ivBytes);

			// Encrypt Data
			Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
			cipher.init(Cipher.ENCRYPT_MODE, key, iv);
			byte[] encryptedText = cipher.doFinal(Files.readAllBytes(zipFile.toPath()));

			// Store Encrypted File, Key & IV Files
			writeToFile(new File(lockedFolder.getCanonicalPath(), namePrefix + ".encrypted.bin").getCanonicalPath(), encryptedText);
			writeToFile(new File(lockedFolder.getCanonicalPath(), namePrefix + ".key.bin").getCanonicalPath(), key.getEncoded());
			writeToFile(new File(lockedFolder.getCanonicalPath(), namePrefix + ".iv.bin").getCanonicalPath(), iv.getIV());
			
			zipFile.delete();
			
			println("The Folder has been Locked at:\n" + lockedFolder.getCanonicalPath() + "\nPlease keep the Key and IV files safe as these files are required to Unlock the Folder.", "success");
		} catch(Exception e) {
			e.printStackTrace();
		}
	}
	
	private static void processDecryption(BufferedReader br) {
		String encryptedFileNameRegExp = "^[\\w\\-\\s]+(\\.encrypted\\.bin)$";
		String keyFileNameRegExp = "^[\\w\\-\\s]+(\\.key\\.bin)$";
		String ivFileNameRegExp = "^[\\w\\-\\s]+(\\.iv\\.bin)$";

		try {
			// Get Encrypted File Path
			print("Enter Locked File Location To Unlock: ");
			
			String encryptedFilePath = br.readLine();
			
			if(encryptedFilePath == null) {
				println("Invalid File or File not Found!", "error");
				return;
			}
			
			File encryptedFile = new File(encryptedFilePath);
			
			// Validate Encrypted File Path
			if(!(encryptedFile.exists() && encryptedFile.getName().matches(encryptedFileNameRegExp))) {
				println("Invalid File or File not Found!", "error");
				return;
			}
			
			// Get Key File Path
			print("Enter Key file Location: ");
			
			String keyFilePath = br.readLine();
			
			if(keyFilePath == null) {
				println("Invalid File or File not Found!", "error");
				return;
			}
			
			File keyFile = new File(keyFilePath);
			
			// Validate Key File Path
			if(!(keyFile.exists() && keyFile.getName().matches(keyFileNameRegExp))) {
				println("Invalid File or File not Found!", "error");
				return;
			}
			
			// Get IV File Path
			print("Enter IV file Location: ");
			
			String ivFilePath = br.readLine();
			
			if(ivFilePath == null) {
				println("Invalid File or File not Found!", "error");
				return;
			}
			
			File ivFile = new File(ivFilePath);
			
			// Validate IV File Path
			if(!(ivFile.exists() && ivFile.getName().matches(ivFileNameRegExp))) {
				println("Invalid File or File not Found!", "error");
				return;
			}
			
			byte[] encryptedFileContent = readFileContent(encryptedFilePath);
			byte[] keyFileContent = readFileContent(keyFilePath);
			byte[] ivFileContent = readFileContent(ivFilePath);

			String encryptedFileName = encryptedFile.getName();
			String namePrefix = encryptedFileName.substring(0, encryptedFileName.indexOf(".encrypted.bin"));

			if(namePrefix.length() == 0) {
				println("Unable to Get the Name prefix!", "error");
				return;
			}
			
			File parentFolder = encryptedFile.getParentFile();
			File unlockedFolder = new File(parentFolder.getCanonicalPath(), namePrefix + "_unlocked");
			unlockedFolder.mkdir();

			SecretKey key = new SecretKeySpec(keyFileContent, 0, keyFileContent.length, "AES");
			IvParameterSpec iv = new IvParameterSpec(ivFileContent);
			
			// Decrypt
			Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
			cipher.init(Cipher.DECRYPT_MODE, key, iv);
			byte[] decryptedText = cipher.doFinal(encryptedFileContent);

			File zipFile = new File(unlockedFolder.getCanonicalPath(), namePrefix + ".zip");
			zipFile.createNewFile();
			Files.write(zipFile.toPath(), decryptedText, StandardOpenOption.WRITE);
			
			File destinationFolder = new File(unlockedFolder.getCanonicalPath(), namePrefix);
			destinationFolder.mkdir();

			unZipFile(zipFile, destinationFolder);

			zipFile.delete();

			encryptedFile.delete();
			keyFile.delete();
			ivFile.delete();
			
			println("The Folder has been Unlocked at:\n" + unlockedFolder.getCanonicalPath(), "success");
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static void main(String[] args) {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		try {
			print("Enter any Key to Continue (l/L - Lock, u/U - Unlock): ");
			String action = br.readLine();

			if(action != null) {
				action = action.toLowerCase();
				
				if(action.equals("l")) {
					processEncryption(br);
				} else if(action.equals("u")) {
					processDecryption(br);
				} else {
					println("Invalid Option!", "error");
				}
			} else {
				println("Invalid Option!", "error");
			}
		} catch(Exception e) {
			e.printStackTrace();
		}
	}

}