open-JSD-10202/src/main/java/com/fr/plugin/xx/zyjn/decode/encode/EncryptOutputStream.java

package com.fr.plugin.xx.zyjn.decode.encode;

import cn.net.drm.edi.client.DrmAgent;
import cn.net.drm.edi.client.DrmAgentInf;
import cn.net.drm.io.IOUtils;
import com.fr.general.CommonIOUtils;
import com.fr.plugin.xx.zyjn.decode.DecodeConfig;
import com.fr.plugin.xx.zyjn.decode.LogUtils;
import com.fr.plugin.xx.zyjn.decode.encode.bi.FileUtil;
import org.apache.tools.zip.ZipOutputStream;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import java.io.*;
import java.util.UUID;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @Author xx
 * @Date 2022/3/3
 * @Description
 **/
public class EncryptOutputStream extends ServletOutputStream {

    private OutputStream original = null;
    private boolean closed = false;
    private ByteArrayOutputStream crypt = new ByteArrayOutputStream();

    public EncryptOutputStream(OutputStream original) {
        this.original = original;
    }

    /**
     * 实际加密的代码就在这里实现了
     *
     * @param bytes
     * @param ignore
     * @return
     */
    public static byte[] encrypt(byte[] bytes, boolean ignore) {
        DecodeConfig config = DecodeConfig.getInstance();
        /*if (!config.isConfiged()) {
            LogUtils.warn("config not finish, not encrypted");
            return bytes;
        }*/
        ByteArrayInputStream bis = new ByteArrayInputStream(bytes);
        byte[] header = new byte[256];
        bis.mark(256);
        byte[] encode = bytes;
        OutputStream fos = null;
        try {
            DrmAgentInf instance = null;
            instance = DrmAgent.getInstance();
            File tempFile = File.createTempFile("tmp" + UUID.randomUUID().toString(), null);
            LogUtils.info("create tmp file path is {}", tempFile.getPath());
            bis.read(header);
            bis.reset();
            boolean isEncryptFile = instance.isEncrypted(tempFile);
            if (!isEncryptFile) {
                if (FileUtil.isArchiveFile(bytes) && !ignore) {
                    encode = encodeZip(bytes, instance, config);
                } else {
                    //fos = instance.encryptBasicFile(tempFile, config.getOwnerId(), config.getSecretLevelId());
                    int len = 0;
                    byte[] buffer = new byte[1024];
                    while ((len = bis.read(buffer)) > 0) {
                        fos.write(buffer, 0, len);
                    }
                    fos.flush();
                    FileInputStream in = new FileInputStream(tempFile);
                    encode = CommonIOUtils.inputStream2Bytes(in);
                    tempFile.deleteOnExit();
                }
            }
        } catch (Exception e) {
            LogUtils.error(e.getMessage(), e);
        } finally {
            IOUtils.closeQuietly(bis);
            IOUtils.closeQuietly(fos);

        }
        //假设逆序就是加密了
        return encode;
    }

    private static byte[] encodeZip(byte[] bytes, DrmAgentInf instance, DecodeConfig config) {
        File zipFile = null, encodeZipFile = null;
        File unzipFileFolder = FileUtil.createTempFolder("tmp_unzip" + UUID.randomUUID().toString());
        File zipFileFolder = FileUtil.createTempFolder("tmp_zip" + UUID.randomUUID().toString());
        try {
            encodeZipFile = FileUtil.createTempFile();
            zipFile = FileUtil.createTempFile();
            //先缓存文件
            com.fr.general.IOUtils.copyBinaryTo(new ByteArrayInputStream(bytes), new FileOutputStream(zipFile));
            LogUtils.debug4plugin("save zip file path:{} size is {}", zipFile.getAbsolutePath(), zipFile.getTotalSpace());
            //再解压出来
            com.fr.general.IOUtils.unzip(zipFile, unzipFileFolder.getAbsolutePath(),"GBK");
            LogUtils.debug4plugin("unzip file path is {} file list is {} unzip size is {}", unzipFileFolder.getAbsolutePath(), Stream.of(unzipFileFolder.listFiles()).map(e->e.getName()).collect(Collectors.joining(",")), unzipFileFolder.listFiles().length);
            File[] files = unzipFileFolder.listFiles();
            ZipOutputStream zipOutputStream = new ZipOutputStream(encodeZipFile);
            zipOutputStream.setEncoding("GBK");
            for (int i = 0; i < files.length; i++) {
                File file = files[i];
                //加密复制到另外一个文件夹
                String encode = zipFileFolder + File.separator + file.getName();
                InputStream bis = new FileInputStream(file);
                //加密
                OutputStream fos =null;// instance.encryptBasicStream(new FileOutputStream(encode), config.getOwnerId(), config.getSecretLevelId());
                int len = 0;
                byte[] buffer = new byte[1024];
                while ((len = bis.read(buffer)) > 0) {
                    fos.write(buffer, 0, len);
                }
                fos.flush();
                IOUtils.closeQuietly(bis);
                IOUtils.closeQuietly(fos);
                LogUtils.debug4plugin("encode file {} to {}", file.getAbsolutePath(), encode);
            }
            //zip
            com.fr.general.IOUtils.zip(zipOutputStream, zipFileFolder);
            zipOutputStream.closeEntry();
            zipOutputStream.close();
            LogUtils.debug4plugin("zip file folder {} to {}", zipFileFolder, encodeZipFile.getAbsolutePath());
            bytes = CommonIOUtils.inputStream2Bytes(new FileInputStream(encodeZipFile));
        } catch (Exception e) {
            LogUtils.error(e.getMessage(), e);
        } finally {
            if (zipFile != null) {
                zipFile.deleteOnExit();
            }
            if (encodeZipFile != null) {
                encodeZipFile.deleteOnExit();
            }
            unzipFileFolder.deleteOnExit();
            zipFileFolder.deleteOnExit();
            zipFile.deleteOnExit();
        }
        return bytes;
    }

    @Override
    public void write(int b) throws IOException {
        if (closed) {
            throw new IOException("Cannot write to a closed output stream");
        }
        crypt.write(b);
    }

    @Override
    public void close() throws IOException {
        if (!closed) {
            byte[] bytes = crypt.toByteArray();
            bytes = encrypt(bytes, false);
            original.write(bytes);
            original.flush();
            original.close();
            closed = true;
        }
    }

    @Override
    public void flush() throws IOException {
        crypt.flush();
    }

    @Override
    public void write(byte[] data) throws IOException {
        write(data, 0, data.length);
    }

    @Override
    public void write(byte[] data, int off, int len) throws IOException {
        if (closed) {
            throw new IOException("Cannot write to a closed output stream");
        }
        crypt.write(data, off, len);
    }

    @Override
    public boolean isReady() {
        return true;
    }

    @Override
    public void setWriteListener(WriteListener listener) {

    }

}