Main.java

import javax.crypto.*;
import javax.crypto.spec.*;
import java.io.*;
import java.nio.file.*;
import java.security.*;
import java.security.spec.*;
import java.util.*;

public class Main {

    // ===== FOLDERS =====
    static final Path INPUT_DIR = Paths.get("input");
    static final Path ENCRYPTED_DIR = Paths.get("encrypted");
    static final Path OUTPUT_DIR = Paths.get("output");

    static final Path USER_KEYS = Paths.get("keys/User");
    static final Path CLIENT_KEYS = Paths.get("keys/Client");

    static final String USER_PRIVATE_KEY = "private.key";
    static final String CLIENT_PUBLIC_KEY = "public.key";

    static final Set<String> IGNORE = Set.of(".jpg", ".png");

    // ===== CRYPTO CONSTANTS =====
    static final int AES_KEY_SIZE = 256;
    static final int GCM_IV_SIZE = 12;
    static final int GCM_TAG_SIZE = 128;
    static final int CHUNK_SIZE = 64 * 1024 * 1024; // 64MB

    public static void main(String[] args) throws Exception {
        Files.createDirectories(ENCRYPTED_DIR);
        Files.createDirectories(OUTPUT_DIR);

        try (Scanner scanner = new Scanner(System.in)) {
            while (true) {
                System.out.println("\n1. Encrypt\n2. Decrypt\n3. Exit");
                System.out.print("> ");

                switch (scanner.nextLine().trim()) {
                    case "1" -> encryptMode();
                    case "2" -> decryptMode();
                    case "3" -> {
                        System.out.println("Exiting...");
                        return;
                    }
                    default -> System.out.println("Invalid option.");
                }
            }
        }
    }

    // ================= ENCRYPTION =================
    static void encryptMode() throws Exception {
        PublicKey clientPub = loadPublicKey(CLIENT_KEYS.resolve(CLIENT_PUBLIC_KEY));

        Files.walk(INPUT_DIR)
                .filter(Files::isRegularFile)
                .forEach(file -> {
                    try {
                        if (shouldIgnore(file))
                            return;
                        encryptFile(file, clientPub);
                        Files.delete(file);
                        System.out.println("✔ Encrypted: " + file);
                    } catch (Exception e) {
                        System.out.println("❌ Failed: " + file + " (" + e.getMessage() + ")");
                    }
                });
    }

    static void encryptFile(Path file, PublicKey rsaPub) throws Exception {

        // Generate AES key
        KeyGenerator gen = KeyGenerator.getInstance("AES");
        gen.init(AES_KEY_SIZE);
        SecretKey aesKey = gen.generateKey();

        Path relative = INPUT_DIR.relativize(file);
        Path encFile = ENCRYPTED_DIR.resolve(relative + ".enc");
        Files.createDirectories(encFile.getParent());

        try (
                InputStream in = new BufferedInputStream(Files.newInputStream(file));
                DataOutputStream out = new DataOutputStream(
                        new BufferedOutputStream(Files.newOutputStream(encFile)))) {
            byte[] buffer = new byte[CHUNK_SIZE];
            int bytesRead;

            while ((bytesRead = in.read(buffer)) != -1) {
                byte[] iv = new byte[GCM_IV_SIZE];
                SecureRandom.getInstanceStrong().nextBytes(iv);

                Cipher aes = Cipher.getInstance("AES/GCM/NoPadding");
                aes.init(Cipher.ENCRYPT_MODE, aesKey,
                        new GCMParameterSpec(GCM_TAG_SIZE, iv));

                byte[] encrypted = aes.doFinal(buffer, 0, bytesRead);

                out.writeInt(encrypted.length); // store chunk length
                out.write(iv); // store IV
                out.write(encrypted); // store encrypted chunk + tag
            }
        }

        // Encrypt AES key with RSA
        Cipher rsa = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        rsa.init(Cipher.ENCRYPT_MODE, rsaPub);
        byte[] encKey = rsa.doFinal(aesKey.getEncoded());

        Files.write(encFile.resolveSibling(encFile.getFileName() + ".key"), encKey);
    }

    // ================= DECRYPTION =================
    static void decryptMode() throws Exception {
        PrivateKey userPriv = loadPrivateKey(USER_KEYS.resolve(USER_PRIVATE_KEY));

        Files.walk(ENCRYPTED_DIR)
                .filter(p -> p.toString().endsWith(".enc"))
                .forEach(encFile -> {
                    try {
                        decryptFile(encFile, userPriv);
                        System.out.println("✔ Decrypted: " + encFile);
                    } catch (Exception e) {
                        System.out.println("❌ Failed: " + encFile + " (" + e.getMessage() + ")");
                    }
                });
    }

    static void decryptFile(Path encFile, PrivateKey rsaPriv) throws Exception {
        Path keyFile = encFile.resolveSibling(encFile.getFileName() + ".key");
        if (!Files.exists(keyFile))
            throw new SecurityException("Missing key file: " + keyFile);

        byte[] encKey = Files.readAllBytes(keyFile);

        Cipher rsa = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding");
        rsa.init(Cipher.DECRYPT_MODE, rsaPriv);
        byte[] aesBytes = rsa.doFinal(encKey);
        SecretKey aesKey = new SecretKeySpec(aesBytes, "AES");

        Path relative = ENCRYPTED_DIR.relativize(encFile);
        Path outFile = OUTPUT_DIR.resolve(relative.toString().replace(".enc", ""));
        Files.createDirectories(outFile.getParent());

        try (
                DataInputStream in = new DataInputStream(
                        new BufferedInputStream(Files.newInputStream(encFile)));
                OutputStream out = Files.newOutputStream(outFile)) {
            while (true) {
                int len;
                try {
                    len = in.readInt();
                } catch (EOFException e) {
                    break; // finished
                }

                byte[] iv = in.readNBytes(GCM_IV_SIZE);
                byte[] encrypted = in.readNBytes(len);

                Cipher aes = Cipher.getInstance("AES/GCM/NoPadding");
                aes.init(Cipher.DECRYPT_MODE, aesKey,
                        new GCMParameterSpec(GCM_TAG_SIZE, iv));

                byte[] plain = aes.doFinal(encrypted);
                out.write(plain);
            }
        }

        Files.delete(encFile);
        Files.delete(keyFile);
    }

    // ================= UTILITIES =================
    static boolean shouldIgnore(Path file) {
        String n = file.getFileName().toString().toLowerCase();
        return IGNORE.stream().anyMatch(n::endsWith);
    }

    static PublicKey loadPublicKey(Path path) throws Exception {
        byte[] b = Base64.getDecoder().decode(Files.readAllBytes(path));
        return KeyFactory.getInstance("RSA")
                .generatePublic(new X509EncodedKeySpec(b));
    }

    static PrivateKey loadPrivateKey(Path path) throws Exception {
        byte[] b = Base64.getDecoder().decode(Files.readAllBytes(path));
        return KeyFactory.getInstance("RSA")
                .generatePrivate(new PKCS8EncodedKeySpec(b));
    }
}