package xyz.docbleach.module.zip;

import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import xyz.docbleach.api.BleachSession;
import xyz.docbleach.api.bleach.Bleach;
import xyz.docbleach.api.exception.BleachException;
import xyz.docbleach.api.exception.RecursionBleachException;
import xyz.docbleach.api.util.CloseShieldInputStream;
import xyz.docbleach.api.util.StreamUtils;

public class ArchiveBleach implements Bleach {

  private static final Logger LOGGER = LoggerFactory.getLogger(ArchiveBleach.class);
  private static final byte[] ZIP_MAGIC = new byte[]{0x50, 0x4B, 0x03, 0x04};

  @Override
  public boolean handlesMagic(InputStream stream) {
    return StreamUtils.hasHeader(stream, ZIP_MAGIC);
  }

  @Override
  public String getName() {
    return "Zip Bleach";
  }

  @Override
  public void sanitize(InputStream inputStream, OutputStream outputStream, BleachSession session)
      throws BleachException {
    ZipInputStream zipIn = new ZipInputStream(inputStream);
    ZipOutputStream zipOut = new ZipOutputStream(outputStream);

    try {
      ZipEntry entry;
      while ((entry = zipIn.getNextEntry()) != null) {
        LOGGER.trace(
            "Entry: {} - Size: (original: {}, compressed: {})",
            entry.getName(),
            entry.getSize(),
            entry.getCompressedSize());

        if (entry.isDirectory()) {
          ZipEntry newEntry = new ZipEntry(entry);
          zipOut.putNextEntry(newEntry);
        } else {
          sanitizeFile(session, zipIn, zipOut, entry);
        }

        zipOut.closeEntry();
      }

      zipOut.finish();
    } catch (IOException e) {
      LOGGER.error("Error in ArchiveBleach", e);
    }
  }

  private void sanitizeFile(
      BleachSession session, ZipInputStream zipIn, ZipOutputStream zipOut, ZipEntry entry)
      throws IOException, BleachException {
    ByteArrayOutputStream streamBuilder = new ByteArrayOutputStream();

    int bytesRead;
    // @TODO: check real file size?
    byte[] tempBuffer = new byte[1024];
    while ((bytesRead = zipIn.read(tempBuffer)) != -1) {
      streamBuilder.write(tempBuffer, 0, bytesRead);
    }
    ByteArrayInputStream bais = new ByteArrayInputStream(streamBuilder.toByteArray());
    CloseShieldInputStream is = new CloseShieldInputStream(new BufferedInputStream(bais));

    ByteArrayOutputStream out = new ByteArrayOutputStream();

    try {
      session.sanitize(is, out);
    } catch (RecursionBleachException e) {
      LOGGER.error("Zip Bomb?", e);
      streamBuilder.writeTo(out);
      // Stream was untouched, o/
    } catch (BleachException e) {
      LOGGER.error("An error occured ", e);
    } finally {
      bais.close();
      is._close();
      streamBuilder.close();
    }

    ZipEntry newEntry = cloneEntry(entry);
    newEntry.setCompressedSize(-1);
    newEntry.setSize(out.size());
    newEntry.setComment(newEntry.getComment() + " - DocBleach");

    zipOut.putNextEntry(newEntry);
    out.writeTo(zipOut);
    out.close();
  }

  // Copies everything except size & CRC-32
  private ZipEntry cloneEntry(ZipEntry entry) {
    ZipEntry newEntry = new ZipEntry(entry.getName());

    newEntry.setTime(entry.getTime());
    if (entry.getCreationTime() != null) {
      newEntry.setCreationTime(entry.getCreationTime());
    }
    if (entry.getLastModifiedTime() != null) {
      newEntry.setLastModifiedTime(entry.getLastModifiedTime());
    }
    if (entry.getLastAccessTime() != null) {
      newEntry.setLastAccessTime(entry.getLastAccessTime());
    }
    newEntry.setComment(entry.getComment());
    newEntry.setExtra(entry.getExtra());

    return newEntry;
  }
}

@@ -1447,219 +1447,92 @@
     
-LOGGER.trace(%0A            %22Entry: %7B%7D - Size: (original: %7B%7D, compressed: %7B%7D)%22,%0A            entry.getName(),%0A            entry.getSize(),%0A            entry.getCompressedSize());%0A%0A        if (entry.isDirectory()) %7B
+if (entry.isDirectory()) %7B%0A          LOGGER.trace(%22Directory: %7B%7D%22, entry.getName());
 %0A   
@@ -1616,24 +1616,24 @@
 (newEntry);%0A
-
         %7D el
@@ -1637,16 +1637,70 @@
  else %7B%0A
+          LOGGER.trace(%22Entry: %7B%7D%22, entry.getName());%0A
         

Remove the erroneous file size