package org.dhatim.fastexcel;

import org.apache.poi.openxml4j.exceptions.InvalidFormatException;
import org.apache.poi.openxml4j.opc.OPCPackage;
import org.apache.poi.poifs.crypt.Decryptor;
import org.apache.poi.poifs.crypt.EncryptionInfo;
import org.apache.poi.poifs.crypt.EncryptionMode;
import org.apache.poi.poifs.crypt.Encryptor;
import org.apache.poi.poifs.filesystem.POIFSFileSystem;
import org.apache.poi.util.IOUtils;
import org.apache.poi.xssf.usermodel.XSSFCell;
import org.apache.poi.xssf.usermodel.XSSFRow;
import org.apache.poi.xssf.usermodel.XSSFSheet;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.junit.jupiter.api.Test;

import java.io.*;
import java.security.GeneralSecurityException;

public class EncryptDecryptTest {
    @Test
    void ProtectTest() throws IOException, InvalidFormatException, GeneralSecurityException {

        final String secretKey = "foobaa";

        final String temp_excel_poi = "D://temp_excel_poi.xlsx";
        final String temp_excel_poi_encrypt = "D://temp_excel_poi_encrypt.xlsx";
        final String temp_excel_poi_decrypt = "D://temp_excel_poi_decrypt.xlsx";

        /* create new excel by poi */
        try (XSSFWorkbook workbook = new XSSFWorkbook(); FileOutputStream foss =new FileOutputStream(temp_excel_poi)) {
            XSSFSheet sheet = workbook.createSheet();
            XSSFRow row = sheet.createRow(0);
            XSSFCell cell = row.createCell(0);
            cell.setCellValue("Hello Apache POI");
            workbook.write(foss);
        }

        /* encrypt excel by poi */
        try ( POIFSFileSystem fs = new POIFSFileSystem()) {
            EncryptionInfo info = new EncryptionInfo(EncryptionMode.agile);
            // EncryptionInfo info = new EncryptionInfo(EncryptionMode.agile, CipherAlgorithm.aes192, HashAlgorithm.sha384, -1, -1, null);
            Encryptor enc = info.getEncryptor();
            enc.confirmPassword(secretKey);
            // Read in an existing OOXML file and write to encrypted output stream
            // don't forget to close the output stream otherwise the padding bytes aren't added
            try (OPCPackage opc = OPCPackage.open(temp_excel_poi); OutputStream os = enc.getDataStream(fs)) {
                opc.save(os);
            }
            // Write out the encrypted version
            try (FileOutputStream fos = new FileOutputStream(temp_excel_poi_encrypt)) {
                fs.writeFilesystem(fos);
            }
        }
        /* decrypt excel by poi */
        try (POIFSFileSystem fileSystem = new POIFSFileSystem(new File(temp_excel_poi_encrypt))) {
            EncryptionInfo info = new EncryptionInfo(fileSystem);
            Decryptor d = Decryptor.getInstance(info);
            if (!d.verifyPassword(secretKey)) {
                throw new RuntimeException("Unable to process: document is encrypted");
            }
            // parse dataStream
            try (InputStream dataStream = d.getDataStream(fileSystem); FileOutputStream fos = new FileOutputStream(temp_excel_poi_decrypt)) {
                IOUtils.copy(dataStream, fos);
            }
        } catch (GeneralSecurityException ex) {
            throw new RuntimeException("Unable to process encrypted document", ex);
        }
    }

}