/*
 * ProGuard -- shrinking, optimization, obfuscation, and preverification
 *             of Java bytecode.
 *
 * Copyright (c) 2002-2018 GuardSquare NV
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 */
package proguard.io;

import proguard.util.StringUtil;

import java.io.*;
import java.util.*;
import java.util.zip.*;

/**
 * This class writes zip data to a given output stream. It returns a new
 * output stream for each zip entry that is opened. An entry can be compressed
 * or uncompressed. Uncompressed entries can be aligned to a multiple of a
 * given number of bytes.
 *
 * Multiple entries and output streams can be open at the same time. The entries
 * are added to the central directory in the order in which they are opened, but
 * the corresponding data are only written when their output streams are closed.
 *
 * The code automatically computes the CRC and lengths of the data, for
 * compressed and uncompressed data.
 *
 * @author Eric Lafortune
 */
public class ZipOutput
{
    private static final int MAGIC_LOCAL_FILE_HEADER             = 0x04034b50;
    private static final int MAGIC_CENTRAL_DIRECTORY_FILE_HEADER = 0x02014b50;
    private static final int MAGIC_END_OF_CENTRAL_DIRECTORY      = 0x06054b50;

    private static final int VERSION              = 10;
    private static final int GENERAL_PURPOSE_FLAG =  0;
    private static final int METHOD_UNCOMPRESSED  =  0;
    private static final int METHOD_COMPRESSED    =  8;

    private static final boolean DEBUG = false;


    private       DataOutputStream outputStream;
    private final int              uncompressedAlignment;
    private final String           comment;

    private List zipEntries    = new ArrayList();
    private Set  zipEntryNames = new HashSet();

    private long centralDirectoryOffset;


    /**
     * Creates a new ZipOutput.
     * @param outputStream the output stream to which the zip data will be
     *                     written.
     */
    public ZipOutput(OutputStream outputStream)
    throws IOException
    {
        this(outputStream, null, null, 1);
    }


    /**
     * Creates a new ZipOutput that aligns uncompressed entries.
     * @param outputStream          the output stream to which the zip data will
     *                              be written.
     * @param header                an optional header for the jar file.
     * @param comment               optional comment for the entire zip file.
     * @param uncompressedAlignment the requested alignment of uncompressed data.
     */
    public ZipOutput(OutputStream outputStream,
                     byte[]       header,
                     String       comment,
                     int          uncompressedAlignment)
    throws IOException
    {
        this.outputStream          = new DataOutputStream(outputStream);
        this.comment               = comment;
        this.uncompressedAlignment = uncompressedAlignment;

        if (header != null)
        {
            outputStream.write(header);
        }
    }


    /**
     * Creates a new zip entry, returning an output stream to write its data.
     * It is the caller's responsibility to close the output stream.
     * @param name             the name of the zip entry.
     * @param compress         specifies whether the entry should be compressed.
     * @param modificationTime the modification date and time of the zip entry,
     *                         in DOS format.
     * @return                 an output stream for writing the data of the
     *                         zip entry.
     */
    public OutputStream createOutputStream(String  name,
                                           boolean compress,
                                           int     modificationTime)
    throws IOException
    {
        return createOutputStream(name,
                                  compress,
                                  modificationTime,
                                  null,
                                  null);
    }


    /**
     * Creates a new zip entry, returning an output stream to write its data.
     * It is the caller's responsibility to close the output stream.
     * @param name             the name of the zip entry.
     * @param compress         specifies whether the entry should be compressed.
     * @param modificationTime the modification date and time of the zip entry,
     *                         in DOS format.
     * @param extraField       optional extra field data. These should contain
     *                         chunks, each with a short ID, a short length
     *                         (little endian), and their corresponding data.
     *                         The IDs 0-31 are reserved for Pkware.
     *                         Java's jar tool just specifies an ID 0xcafe on
     *                         its first entry.
     * @param comment          optional comment.
     * @return                 an output stream for writing the data of the
     *                         zip entry.
     */
    public OutputStream createOutputStream(String  name,
                                           boolean compress,
                                           int     modificationTime,
                                           byte[]  extraField,
                                           String  comment)
    throws IOException
    {
        // Check if the name hasn't been used yet.
        if (!zipEntryNames.add(name))
        {
            throw new IOException("Duplicate jar entry ["+name+"]");
        }

        ZipEntry entry = new ZipEntry(name,
                                      compress,
                                      modificationTime,
                                      extraField,
                                      comment);

        // Add the entry to the list that will be put in the central directory.
        zipEntries.add(entry);

        return entry.createOutputStream();
    }


    /**
     * Closes the zip archive, also closing the underlying output stream.
     */
    public void close() throws IOException
    {
        // Write the central directory.
        writeStartOfCentralDirectory();

        for (int index = 0; index < zipEntries.size(); index++)
        {
            ZipEntry entry = (ZipEntry)zipEntries.get(index);

            entry.writeCentralDirectoryFileHeader();
        }

        writeEndOfCentralDirectory();

        // Close the underlying output stream.
        outputStream.close();

        // Make sure the archive can't be used any further.
        outputStream  = null;
        zipEntries    = null;
        zipEntryNames = null;
    }


    /**
     * Starts the central directory.
     */
    private void writeStartOfCentralDirectory()
    {
        // The central directory as such doesn't have a header.
        centralDirectoryOffset = outputStream.size();
    }


    /**
     * Ends the central directory.
     */
    private void writeEndOfCentralDirectory() throws IOException
    {
        if (DEBUG)
        {
            System.out.println("ZipOutput.writeEndOfCentralDirectory ("+zipEntries.size()+" entries)");
        }

       // The size of the central directory, not counting this trailer.
        long centralDirectorySize = outputStream.size() - centralDirectoryOffset;

        writeInt(MAGIC_END_OF_CENTRAL_DIRECTORY);
        writeShort(0);                    // Number of this disk.
        writeShort(0);                    // Number of disk with central directory.
        writeShort(zipEntries.size());    // Number of records on this disk.
        writeShort(zipEntries.size());    // Total number of records.
        writeInt(centralDirectorySize);   // Size of central directory, in bytes.
        writeInt(centralDirectoryOffset); // Offset of central directory.

        if (comment == null)
        {
            // No comment.
            writeShort(0);
        }
        else
        {
            // Comment length and comment.
            byte[] commentBytes = StringUtil.getUtf8Bytes(comment);
            writeShort(commentBytes.length);
            outputStream.write(commentBytes);
        }
    }


    /**
     * This class represents a zip entry in its enclosing zip file. It can
     * provide an output stream and write its headers and its data to the main
     * zip output stream. In fact, it automatically writes its local header and
     * data when the output stream is closed.
     */
    private class ZipEntry
    {
        private boolean compressed;
        private int     modificationTime;
        private int     crc;
        private long    compressedSize;
        private long    uncompressedSize;
        private long    offset;
        private String  name;
        private byte[]  extraField;
        private String  comment;


        /**
         * Creates a new zip entry, returning output stream to write its data.
         * It is the caller's responsibility to close the output stream.
         * @param name             the name of the zip entry.
         * @param compressed       specifies whether the entry should be
         *                         compressed.
         * @param modificationTime the modification date and time of the zip
         *                         entry, in DOS format.
         * @param extraField       optional extra field data. These should
         *                         contain chunks, each with a short ID, a short
         *                         length (little endian), and their
         *                         corresponding data. The IDs 0-31 are reserved
         *                         for Pkware. Java's jar tool just specifies an
         *                         ID 0xcafe on its first entry.
         * @param comment          optional comment.
         * @return                 an output stream for writing the zip data.
         */
        private ZipEntry(String  name,
                         boolean compressed,
                         int     modificationTime,
                         byte[]  extraField,
                         String  comment)
        {
            this.name             = name;
            this.compressed       = compressed;
            this.modificationTime = modificationTime;
            this.extraField       = extraField;
            this.comment          = comment;
        }


        public OutputStream createOutputStream() throws IOException
        {
            return compressed ?
                (OutputStream)new CompressedZipEntryOutputStream() :
                (OutputStream)new UncompressedZipEntryOutputStream();
        }


        /**
         * Writes the local file header, which precedes the data, to the main
         * zip output stream.
         */
        private void writeLocalFileHeader() throws IOException
        {
            if (DEBUG)
            {
                System.out.println("ZipOutput.writeLocalFileHeader ["+name+"] (compressed = "+compressed+", offset = "+offset+", "+compressedSize+"/"+uncompressedSize+" bytes)");
            }

            writeInt(MAGIC_LOCAL_FILE_HEADER);
            writeShort(VERSION);
            writeShort(GENERAL_PURPOSE_FLAG);
            writeShort(compressed ? METHOD_COMPRESSED : METHOD_UNCOMPRESSED);
            writeInt(modificationTime);
            writeInt(crc);
            writeInt(compressedSize);
            writeInt(uncompressedSize);

            byte[] nameBytes     = StringUtil.getUtf8Bytes(name);
            int nameLength       = nameBytes.length;
            int extraFieldLength = extraField == null ? 0 : extraField.length;

            writeShort(nameLength);
            writeShort(extraFieldLength);

            outputStream.write(nameBytes);

            if (extraField != null)
            {
                outputStream.write(extraField);
            }
        }


        /**
         * Writes the file header for the central directory to the main zip
         * output stream.
         */
        public void writeCentralDirectoryFileHeader() throws IOException
        {
            if (DEBUG)
            {
                System.out.println("ZipOutput.writeCentralDirectoryFileHeader ["+name+"] (compressed = "+compressed+", offset = "+offset+", "+compressedSize+"/"+uncompressedSize+" bytes)");
            }

            writeInt(MAGIC_CENTRAL_DIRECTORY_FILE_HEADER);
            writeShort(VERSION); // Creation version.
            writeShort(VERSION); // Extraction Version.
            writeShort(GENERAL_PURPOSE_FLAG);
            writeShort(compressed ? METHOD_COMPRESSED : METHOD_UNCOMPRESSED);
            writeInt(modificationTime);
            writeInt(crc);
            writeInt(compressedSize);
            writeInt(uncompressedSize);

            byte[] nameBytes    = StringUtil.getUtf8Bytes(name);
            byte[] commentBytes = comment == null ? null :
                                  StringUtil.getUtf8Bytes(comment);

            writeShort(nameBytes.length);
            writeShort(extraField   == null ? 0 : extraField.length);
            writeShort(commentBytes == null ? 0 : commentBytes.length);
            writeShort(0); // Disk number of file start.
            writeShort(0); // Internal file attributes.
            writeInt(0);   // External file attributes.
            writeInt(offset);
            outputStream.write(nameBytes);
            if (extraField != null)
            {
                outputStream.write(extraField);
            }

            if (commentBytes != null)
            {
                outputStream.write(commentBytes);
            }
        }


        /**
         * This OutputStream writes its uncompressed zip entry out to its zip
         * output stream when it is closed.
         */
        private class UncompressedZipEntryOutputStream extends ByteArrayOutputStream
        {
            private CRC32 crc32 = new CRC32();


            private UncompressedZipEntryOutputStream()
            {
                super(16 * 1024);
            }


            // Overridden methods for OutputStream.

            public void write(int b)
            {
                super.write(b);

                crc32.update(b);
            }


            //public void write(byte[] b) throws IOException
            //{
            //    // The super implementation delegates to the method below.
            //    super.write(b);
            //}


            public void write(byte[] b, int off, int len)
            {
                super.write(b, off, len);

                crc32.update(b, off, len);
            }


            public void close() throws IOException
            {
                super.close();

                byte[] bytes = super.toByteArray();

                offset           = outputStream.size();
                crc              = (int)crc32.getValue();
                compressedSize   = bytes.length;
                uncompressedSize = bytes.length;

                writeLocalFileHeader();
                outputStream.write(bytes);
            }
        }


        /**
         * This OutputStream writes its compressed zip entry out to its zip
         * output stream when it is closed.
         */
        private class CompressedZipEntryOutputStream extends DeflaterOutputStream
        {
            private CRC32 crc32 = new CRC32();


            private CompressedZipEntryOutputStream()
            {
                super(new ByteArrayOutputStream(16 * 1024),
                      new Deflater(Deflater.BEST_COMPRESSION, true),
                      1024);
            }


            // Overridden methods for OutputStream.

            //public void write(int b) throws IOException
            //{
            //    // The super implementation delegates to the method below.
            //    super.write(b);
            //}
            //
            //
            //public void write(byte[] b) throws IOException
            //{
            //    // The super implementation delegates to the method below.
            //    super.write(b);
            //}


            public void write(byte[] b, int off, int len) throws IOException
            {
                super.write(b, off, len);

                crc32.update(b, off, len);
                uncompressedSize += len;
            }


            public void close() throws IOException
            {
                // Make sure the memory is freed. [JDK-4797189]
                super.finish();
                super.def.end();
                super.close();

                ByteArrayOutputStream byteArrayOutputStream =
                    (ByteArrayOutputStream)super.out;

                byte[] compressedBytes = byteArrayOutputStream.toByteArray();

                offset         = outputStream.size();
                crc            = (int)crc32.getValue();
                compressedSize = compressedBytes.length;

                writeLocalFileHeader();
                outputStream.write(compressedBytes);
            }
        }
    }


    // Small utility methods.

    /**
     * Writes out a little-endian short value to the zip output stream.
     */
    private void writeShort(int value) throws IOException
    {
        outputStream.write(value);
        outputStream.write(value >>> 8);
    }


    /**
     * Writes out a little-endian int value to the zip output stream.
     */
    private void writeInt(int value) throws IOException
    {
        outputStream.write(value);
        outputStream.write(value >>>  8);
        outputStream.write(value >>> 16);
        outputStream.write(value >>> 24);
    }


    /**
     * Writes out a little-endian int value to the zip output stream.
     */
    private void writeInt(long value) throws IOException
    {
        outputStream.write((int)value);
        outputStream.write((int)(value >>>  8));
        outputStream.write((int)(value >>> 16));
        outputStream.write((int)(value >>> 24));
    }


    /**
     * Provides a simple test for this class, creating a zip file with the
     * given name and a few aligned/compressed/uncompressed zip entries.
     */
    public static void main(String[] args)
    {
        try
        {
            ZipOutput output =
                new ZipOutput(new FileOutputStream(args[0]), null, "Main file comment", 4);

            PrintWriter printWriter1 =
                new PrintWriter(output.createOutputStream("file1.txt", false, 0, new byte[] { 0x34, 0x12, 4, 0, 0x48, 0x65, 0x6c, 0x6c, 0x6f }, "Comment"));
            printWriter1.println("This is file 1.");
            printWriter1.println("Hello, world!");
            printWriter1.close();

            PrintWriter printWriter2 =
                new PrintWriter(output.createOutputStream("file2.txt", true, 0, null, "Another comment"));
            printWriter2.println("This is file 2.");
            printWriter2.println("Hello, world!");
            printWriter2.close();

            PrintWriter printWriter3 =
                new PrintWriter(output.createOutputStream("file3.txt", false, 0, null, "Last comment"));
            printWriter3.println("This is file 3.");
            printWriter3.println("Hello, world!");
            printWriter3.close();

            output.close();
        }
        catch (IOException e)
        {
            e.printStackTrace();
        }
    }
}
