// See the file "COPYING" in the main distribution directory for copyright.

#include "zeek/analyzer/protocol/zip/ZIP.h"

namespace zeek::analyzer::zip {

ZIP_Analyzer::ZIP_Analyzer(Connection* conn, bool orig, Method arg_method)
    : analyzer::tcp::TCP_SupportAnalyzer("ZIP", conn, orig) {
    zip = nullptr;
    zip_status = Z_OK;
    method = arg_method;

    zip = new z_stream;
    zip->zalloc = 0;
    zip->zfree = 0;
    zip->opaque = 0;
    zip->next_out = 0;
    zip->avail_out = 0;
    zip->next_in = 0;
    zip->avail_in = 0;

    // "32" is a gross overload hack that means "check it
    // for whether it's a gzip file".  Sheesh.
    if ( inflateInit2(zip, MAX_WBITS + 32) != Z_OK ) {
        Weird("inflate_init_failed");
        delete zip;
        zip = nullptr;
    }
}

ZIP_Analyzer::~ZIP_Analyzer() { delete zip; }

void ZIP_Analyzer::Done() {
    Analyzer::Done();

    if ( zip )
        inflateEnd(zip);
}

void ZIP_Analyzer::DeliverStream(int len, const u_char* data, bool orig) {
    analyzer::tcp::TCP_SupportAnalyzer::DeliverStream(len, data, orig);

    if ( ! len || zip_status != Z_OK )
        return;

    static unsigned int unzip_size = 4096;
    auto unzipbuf = std::make_unique<Bytef[]>(unzip_size);

    int allow_restart = 1;

    zip->next_in = (Bytef*)data;
    zip->avail_in = len;

    auto orig_next_in = zip->next_in;
    size_t orig_avail_in = zip->avail_in;

    while ( true ) {
        zip->next_out = unzipbuf.get();
        zip->avail_out = unzip_size;

        zip_status = inflate(zip, Z_SYNC_FLUSH);

        if ( zip_status == Z_STREAM_END || zip_status == Z_OK ) {
            allow_restart = 0;

            int have = unzip_size - zip->avail_out;
            if ( have )
                ForwardStream(have, unzipbuf.get(), IsOrig());

            if ( zip_status == Z_STREAM_END ) {
                inflateEnd(zip);
                return;
            }

            if ( zip->avail_in == 0 )
                return;
        }

        else if ( allow_restart && zip_status == Z_DATA_ERROR ) {
            // Some servers seem to not generate zlib headers,
            // so this is an attempt to fix and continue anyway.
            inflateEnd(zip);

            if ( inflateInit2(zip, -MAX_WBITS) != Z_OK ) {
                Weird("inflate_init_failed");
                return;
            }

            zip->next_in = orig_next_in;
            zip->avail_in = orig_avail_in;
            allow_restart = 0;
            continue;
        }

        else {
            Weird("inflate_failed");
            return;
        }
    }
}

} // namespace zeek::analyzer::zip
