leanify.cpp

#include "leanify.h"

#include <algorithm>
#include <iostream>
#include <memory>

#include <zopfli/zlib_container.h>
#include <zopflipng/lodepng/lodepng.h>

#include "formats/data_uri.h"
#include "formats/dwf.h"
#include "formats/format.h"
#include "formats/gft.h"
#include "formats/gz.h"
#include "formats/ico.h"
#include "formats/jpeg.h"
#include "formats/lua.h"
#include "formats/mime.h"
#include "formats/pe.h"
#include "formats/png.h"
#include "formats/rdb.h"
#include "formats/swf.h"
#include "formats/tar.h"
#include "formats/vcf.h"
#include "formats/xml.h"
#include "formats/zip.h"
#include "utils.h"

using std::cerr;
using std::endl;
using std::string;

std::unique_ptr<Format> GetType(void* file_pointer, size_t file_size, const string& filename, int depth) {
  if (depth > max_depth) {
    VerbosePrint("Max depth ", max_depth, " reached, skipping.");
    return std::make_unique<Format>(file_pointer, file_size, depth);
  }

  if (!filename.empty()) {
    size_t dot = filename.find_last_of('.');
    if (dot != string::npos) {
      string ext = filename.substr(dot + 1);
      // toupper
      for (auto& c : ext)
        c &= ~0x20;

      if (ext == "HTML" || ext == "HTM" || ext == "JS" || ext == "CSS") {
        VerbosePrint(ext, " detected.");
        return std::make_unique<DataURI>(file_pointer, file_size, depth);
      }
      if (ext == "VCF" || ext == "VCARD") {
        VerbosePrint(ext, " detected.");
        return std::make_unique<Vcf>(file_pointer, file_size, depth);
      }
      if (ext == "MHT" || ext == "MHTML" || ext == "MIM" || ext == "MIME" || ext == "EML") {
        VerbosePrint(ext, " detected.");
        return std::make_unique<Mime>(file_pointer, file_size, depth);
      }
    }
  }
  if (memcmp(file_pointer, Png::header_magic, sizeof(Png::header_magic)) == 0) {
    VerbosePrint("PNG detected.");
    return std::make_unique<Png>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Jpeg::header_magic, sizeof(Jpeg::header_magic)) == 0) {
    VerbosePrint("JPEG detected.");
    return std::make_unique<Jpeg>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Lua::header_magic, sizeof(Lua::header_magic)) == 0) {
    VerbosePrint("Lua detected.");
    return std::make_unique<Lua>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Zip::header_magic, sizeof(Zip::header_magic)) == 0) {
    VerbosePrint("ZIP detected.");
    return std::make_unique<Zip>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Pe::header_magic, sizeof(Pe::header_magic)) == 0) {
    VerbosePrint("PE detected.");
    return std::make_unique<Pe>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Gz::header_magic, sizeof(Gz::header_magic)) == 0) {
    VerbosePrint("GZ detected.");
    return std::make_unique<Gz>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Ico::header_magic, sizeof(Ico::header_magic)) == 0) {
    VerbosePrint("ICO detected.");
    return std::make_unique<Ico>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Dwf::header_magic, sizeof(Dwf::header_magic)) == 0) {
    VerbosePrint("DWF detected.");
    return std::make_unique<Dwf>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Gft::header_magic, sizeof(Gft::header_magic)) == 0) {
    VerbosePrint("GFT detected.");
    return std::make_unique<Gft>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Rdb::header_magic, sizeof(Rdb::header_magic)) == 0) {
    VerbosePrint("RDB detected.");
    return std::make_unique<Rdb>(file_pointer, file_size, depth);
  } else if (memcmp(file_pointer, Swf::header_magic, sizeof(Swf::header_magic)) == 0 ||
             memcmp(file_pointer, Swf::header_magic_deflate, sizeof(Swf::header_magic_deflate)) == 0 ||
             memcmp(file_pointer, Swf::header_magic_lzma, sizeof(Swf::header_magic_lzma)) == 0) {
    VerbosePrint("SWF detected.");
    return std::make_unique<Swf>(file_pointer, file_size, depth);
  } else {
    // Search for vcard magic which might not be at the very beginning.
    const string vcard_magic = "BEGIN:VCARD";
    const char* fp = static_cast<char*>(file_pointer);
    const char* search_end = fp + std::min(static_cast<size_t>(1024), file_size);
    if (std::search(fp, search_end, vcard_magic.begin(), vcard_magic.end()) < search_end) {
      VerbosePrint("VCF detected.");
      return std::make_unique<Vcf>(file_pointer, file_size, depth);
    }

    // tar file does not have header magic
    // ustar is optional
    {
      auto t = std::make_unique<Tar>(file_pointer, file_size, depth);
      // checking first record checksum
      if (t->IsValid()) {
        VerbosePrint("tar detected.");
        return t;
      }
    }

    // XML file does not have header magic
    // have to parse and see if there are any errors.
    {
      auto x = std::make_unique<Xml>(file_pointer, file_size, depth);
      if (x->IsValid()) {
        VerbosePrint("XML detected.");
        return x;
      }
    }
  }

  VerbosePrint("Format not supported!");
  // for unsupported format, just memmove it.
  return std::make_unique<Format>(file_pointer, file_size, depth);
}

// Leanify the file
// and move the file ahead size_leanified bytes
// the new location of the file will be file_pointer - size_leanified
// it's designed this way to avoid extra memmove or memcpy
// return new size
size_t LeanifyFile(void* file_pointer, size_t file_size, size_t size_leanified /*= 0*/,
                   const string& filename /*= ""*/, int depth /*= 1*/) {
  if (file_pointer == nullptr || file_size == 0) {
    return file_size;
  }
  auto f = GetType(file_pointer, file_size, filename, depth);
  return f->Leanify(size_leanified);
}

size_t ZlibRecompress(uint8_t* src, size_t src_len, size_t size_leanified /*= 0*/) {
  if (!is_fast) {
    size_t uncompressed_size = 0;
    uint8_t* buffer = nullptr;
    if (lodepng_zlib_decompress(&buffer, &uncompressed_size, src, src_len, &lodepng_default_decompress_settings) ||
        !buffer) {
      cerr << "Decompress Zlib data failed." << endl;
    } else {
      ZopfliOptions zopfli_options;
      ZopfliInitOptions(&zopfli_options);
      zopfli_options.numiterations = iterations;

      size_t new_size = 0;
      uint8_t* out_buffer = nullptr;
      ZopfliZlibCompress(&zopfli_options, buffer, uncompressed_size, &out_buffer, &new_size);
      free(buffer);
      if (new_size < src_len) {
        memcpy(src - size_leanified, out_buffer, new_size);
        free(out_buffer);
        return new_size;
      }
      free(out_buffer);
    }
  }

  memmove(src - size_leanified, src, src_len);
  return src_len;
}