#define _CRT_SECURE_NO_DEPRECATE

#include "types.hpp"
#include "dis.hpp"
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

struct FileHeader
{
  sU32 SizeBefore;      // number of bytes before start of code section
  sU32 SizeAfter;       // number of bytes after code section
  sU32 SizeTransformed; // size of transformed code section
  sU32 SizeOriginal;    // size of untransformed code section
  sU32 Origin;          // virtual address of first byte
};

static sU8 *loadFile(sChar *name,sInt &size)
{
  sU8 *mem;
  FILE *file;

  mem = 0;
  file = fopen(name,"rb");
  if(file)
  {
    fseek(file,0,SEEK_END);
    size = ftell(file);
    fseek(file,0,SEEK_SET);
    mem = new sU8[size];
    if(fread(mem,1,size,file) != size)
    {
      delete[] mem;
      mem = 0;
    }
    fclose(file);
  }

  return mem;
}

int main(int argc,char **argv)
{
  if(argc != 4)
  {
    printf("Syntax: dispack e <input.exe/dll> <output.dat>      encode\n");
    printf("        dispack d <input.dat> <output.exe/dll>      decode\n");
    return 1;
  }

  bool encode;
  if(strcmp(argv[1],"e") == 0)
    encode = true;
  else if(strcmp(argv[1],"d") == 0)
    encode = false;
  else
  {
    printf("Please specify e(ncode) or d(ecode)!\n");
    return 1;
  }

  sInt inSize;
  sU8 *inData = loadFile(argv[2],inSize);
  if(!inData)
  {
    printf("Couldn't load input file!\n");
    return 1;
  }

  if(encode)
  {
    // assume the input file is a PE executable.
    IMAGE_DOS_HEADER *doshdr = (IMAGE_DOS_HEADER *) inData;
    IMAGE_NT_HEADERS *nthdr = (IMAGE_NT_HEADERS *) (inData + doshdr->e_lfanew);
    if(nthdr->FileHeader.Machine != IMAGE_FILE_MACHINE_I386
      || nthdr->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC)
    {
      printf("Only 32-bit PE files for x86 supported.\n");
      return 1;
    }

    sU32 imageBase = nthdr->OptionalHeader.ImageBase;
    sU32 codeStart = nthdr->OptionalHeader.BaseOfCode;
    sU32 codeSize = nthdr->OptionalHeader.SizeOfCode;
    sU32 fileOffs = 0; // find file offset of first section
    printf("Start of code section: %08x, size %08x.\n",codeStart,codeSize);

    // find section containing code
    IMAGE_SECTION_HEADER *sec = IMAGE_FIRST_SECTION(nthdr);
    for(sInt i=0;i<nthdr->FileHeader.NumberOfSections;i++)
    {
      if(codeStart >= sec[i].VirtualAddress && codeStart < sec[i].VirtualAddress + sec[i].SizeOfRawData)
        fileOffs = sec[i].PointerToRawData + (codeStart - sec[i].VirtualAddress);
    }

    if(fileOffs == 0)
    {
      printf("Code section not found!\n");
      return 1;
    }

    // transform code
    sU32 transSize;
    sU8 *transformed = DisFilter(inData + fileOffs,codeSize,imageBase + codeStart,transSize);
    printf("Filtered code: %d -> %d bytes (%d%%)\n",codeSize,transSize,MulDiv(transSize,100,codeSize));

    // write output file
    sU32 before = fileOffs;
    sU32 after = inSize - (fileOffs + codeSize);

    FILE *outFile = fopen(argv[3],"wb");
    if(!outFile)
    {
      printf("Couldn't open output file!\n");
      return 1;
    }

    FileHeader hdr;
    hdr.SizeBefore = fileOffs;
    hdr.SizeAfter = inSize - (fileOffs + codeSize);
    hdr.SizeTransformed = transSize;
    hdr.SizeOriginal = codeSize;
    hdr.Origin = imageBase + codeStart;

    sBool ok = fwrite(&hdr,sizeof(FileHeader),1,outFile) == 1
      && fwrite(transformed,transSize,1,outFile) == 1
      && fwrite(inData,hdr.SizeBefore,1,outFile) == 1
      && fwrite(inData + (fileOffs + codeSize),hdr.SizeAfter,1,outFile) == 1;
    fclose(outFile);

    if(!ok)
    {
      printf("Error writing output file!\n");
      return 1;
    }

    delete[] transformed;
  }
  else
  {
    FileHeader *hdr = (FileHeader *) inData;
    sU8 *transformed = inData + sizeof(FileHeader);
    sU8 *before = transformed + hdr->SizeTransformed;
    sU8 *after = before + hdr->SizeBefore;

    // alloc buffer for unfiltered code
    sU8 *decoded = new sU8[hdr->SizeOriginal];
    if(!decoded || !DisUnFilter(transformed,hdr->SizeTransformed,decoded,hdr->SizeOriginal,hdr->Origin))
    {
      printf("Error decoding transformed code!\n");
      return 1;
    }

    // write output file
    FILE *outFile = fopen(argv[3],"wb");
    if(!outFile)
    {
      printf("Couldn't open output file!\n");
      return 1;
    }

    sBool ok = fwrite(before,hdr->SizeBefore,1,outFile) == 1
      && fwrite(decoded,hdr->SizeOriginal,1,outFile) == 1
      && fwrite(after,hdr->SizeAfter,1,outFile) == 1;
    fclose(outFile);

    if(!ok)
    {
      printf("Error writing output file!\n");
      return 1;
    }

    delete[] decoded;
  }

  delete[] inData;
  return 0;
}