/*
    iwm-forwarding-win32
    Copyright (C) 2006-2007 Kazuki IWAMOTO

    This library is free software; you can redistribute it and/or
    modify it under the terms of the GNU Lesser General Public
    License as published by the Free Software Foundation; either
    version 2.1 of the License, or (at your option) any later version.

    This library 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
    Lesser General Public License for more details.

    You should have received a copy of the GNU Lesser General Public
    License along with this library; if not, write to the Free Software
    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tchar.h>
#include <windows.h>
#include <shlwapi.h>


#define ALIGN 512


int
main (int   argc,
      char *argv[])
{
  char *file, *path, **name;
  int i;
  FILE *fp;
  BYTE bCode[8] = {0xb4, 0x4c, 0xcd, 0x21, 0x00, 0x00, 0x00, 0x00};
  DWORD dwSignature, dwOffset;
  LPDWORD lpdwAddress, lpdwName;
  LPWORD lpwOrdinal;
  IMAGE_DOS_HEADER idh;
  IMAGE_FILE_HEADER ifh;
  IMAGE_OPTIONAL_HEADER ioh;
  IMAGE_EXPORT_DIRECTORY ied;
  PIMAGE_SECTION_HEADER pish;

  if (argc != 3)
    return -1;
  fp = fopen (argv[1], "rb");
  if (!fp || fread (&idh, sizeof (IMAGE_DOS_HEADER), 1, fp) != 1
        || idh.e_magic != IMAGE_DOS_SIGNATURE
        || fseek (fp, idh.e_lfanew, SEEK_SET) != 0
        || fread (&dwSignature, sizeof (DWORD), 1, fp) != 1
        || dwSignature != IMAGE_NT_SIGNATURE
        || fread (&ifh, sizeof (IMAGE_FILE_HEADER), 1, fp) != 1
        || ifh.NumberOfSections == 0
        || fread (&ioh, sizeof (IMAGE_OPTIONAL_HEADER), 1, fp) != 1
        || ioh.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress == 0)
    return -1;
  pish = malloc (ifh.NumberOfSections * sizeof (IMAGE_SECTION_HEADER));
  if (!pish || fread (pish, sizeof (IMAGE_SECTION_HEADER),
                            ifh.NumberOfSections, fp) != ifh.NumberOfSections)
    return -1;
  for (i = 0; i < ifh.NumberOfSections; i++)
    if (pish[i].VirtualAddress
            <= ioh.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
            && ioh.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
                            < pish[i].VirtualAddress + pish[i].SizeOfRawData)
      {
        if (fseek (fp,
            ioh.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress
            - pish[i].VirtualAddress + pish[i].PointerToRawData, SEEK_SET) != 0
                || fread (&ied, sizeof (IMAGE_EXPORT_DIRECTORY), 1, fp) != 1)
          return -1;
        break;
      }
  lpdwName = malloc (ied.NumberOfNames * sizeof (DWORD));
  lpwOrdinal = malloc (ied.NumberOfNames * sizeof (WORD));
  for (i = 0; i < ifh.NumberOfSections; i++)
    if ((pish[i].VirtualAddress <= ied.AddressOfNames
        && ied.AddressOfNames < pish[i].VirtualAddress + pish[i].SizeOfRawData
        && (fseek (fp, ied.AddressOfNames
            - pish[i].VirtualAddress + pish[i].PointerToRawData, SEEK_SET) != 0
                    || fread (lpdwName, sizeof (DWORD), ied.NumberOfNames, fp)
                                                        != ied.NumberOfNames))
      || (pish[i].VirtualAddress <= ied.AddressOfNameOrdinals
        && ied.AddressOfNameOrdinals
                            < pish[i].VirtualAddress + pish[i].SizeOfRawData
        && (fseek (fp, ied.AddressOfNameOrdinals
            - pish[i].VirtualAddress + pish[i].PointerToRawData, SEEK_SET) != 0
                    || fread (lpwOrdinal, sizeof (WORD), ied.NumberOfNames, fp)
                                                        != ied.NumberOfNames)))
      return -1;
  name = malloc (ied.NumberOfNames * sizeof (char *));
  for (i = 0; i < ied.NumberOfNames; i++)
    {
      int j;

      name[i] = NULL;
      for (j = 0; j < ifh.NumberOfSections; j++)
        if (pish[j].VirtualAddress <= lpdwName[i]
            && lpdwName[i] < pish[j].VirtualAddress + pish[j].SizeOfRawData)
          {
            int c, n = 0;

            if (fseek (fp, lpdwName[i] - pish[j].VirtualAddress
                                    + pish[j].PointerToRawData, SEEK_SET) != 0)
              return -1;
            while (c = fgetc (fp), c != EOF && c != '\0')
              {
                name[i] = realloc (name[i], (n + 2) * sizeof (char));
                (name[i])[n++] = c;
              }
            if (name[i])
              (name[i])[n] = '\0';
            break;
          }
    }
  if (fclose (fp) != 0)
    return -1;

  file = malloc ((lstrlenA (argv[1]) + 1) * sizeof (char));
  lstrcpyA (file, argv[1]);
  PathStripPathA (file);
  PathRemoveExtensionA (file);
  path = malloc ((lstrlenA (argv[2]) + 1) * sizeof (char));
  lstrcpyA (path, argv[2]);
  PathStripPathA (path);
  dwOffset = sizeof (IMAGE_DOS_HEADER)
           + sizeof (bCode)
           + sizeof (DWORD)
           + sizeof (IMAGE_FILE_HEADER)
           + sizeof (IMAGE_OPTIONAL_HEADER)
           + sizeof (IMAGE_SECTION_HEADER);
  memset (&idh, 0, sizeof (IMAGE_DOS_HEADER));
  idh.e_magic = IMAGE_DOS_SIGNATURE;
  idh.e_cblp = 0x48;
  idh.e_cblp = 1;
  idh.e_cparhdr = 4;
  idh.e_minalloc = 0x40;
  idh.e_maxalloc = 0xffff;
  idh.e_sp = 0x400;
  idh.e_ip = 0x40;
  idh.e_lfanew = 0x48;
  ifh.Machine = IMAGE_FILE_MACHINE_I386;
  ifh.NumberOfSections = 1;
  ifh.TimeDateStamp = 0;
  ifh.PointerToSymbolTable = 0;
  ifh.NumberOfSymbols = 0;
  ifh.SizeOfOptionalHeader = sizeof (IMAGE_OPTIONAL_HEADER);
  ifh.Characteristics = IMAGE_FILE_DLL | IMAGE_FILE_32BIT_MACHINE
                                                | IMAGE_FILE_EXECUTABLE_IMAGE;
  ioh.Magic = IMAGE_NT_OPTIONAL_HDR32_MAGIC;
  ioh.MajorLinkerVersion = 0;
  ioh.MinorLinkerVersion = 0;
  ioh.SizeOfCode = 0;
  ioh.SizeOfInitializedData = sizeof (IMAGE_EXPORT_DIRECTORY)
                                    + ied.NumberOfFunctions * sizeof (DWORD)
                                    + ied.NumberOfNames * sizeof (DWORD)
                                    + ied.NumberOfNames * sizeof (WORD)
                                    + (lstrlenA (path) + 1) * sizeof (char);
  for (i = 0; i < ied.NumberOfNames; i++)
    if (name[i])
      {
        ioh.SizeOfInitializedData += (lstrlenA (name[i]) + 1) * 2;
      }
    else
      {
        if (0 <= lpwOrdinal[i] && lpwOrdinal[i] <= 9)
          ioh.SizeOfInitializedData += 3;
        else if (10 <= lpwOrdinal[i] && lpwOrdinal[i] <= 99)
          ioh.SizeOfInitializedData += 4;
        else if (100 <= lpwOrdinal[i] && lpwOrdinal[i] <= 999)
          ioh.SizeOfInitializedData += 5;
        else if (1000 <= lpwOrdinal[i] && lpwOrdinal[i] <= 9999)
          ioh.SizeOfInitializedData += 6;
        else
          ioh.SizeOfInitializedData += 7;
      }
  ioh.SizeOfInitializedData += (lstrlenA (file) + 1) * ied.NumberOfNames;
  ioh.SizeOfUninitializedData = 0;
  ioh.AddressOfEntryPoint = 0;
  ioh.BaseOfCode = (dwOffset + ALIGN - 1) / ALIGN * ALIGN;
  ioh.BaseOfData = ioh.BaseOfCode;
  ioh.ImageBase = 0x10000000;
  ioh.SectionAlignment = ALIGN;
  ioh.FileAlignment = ALIGN;
  ioh.MajorOperatingSystemVersion = 4;
  ioh.MinorOperatingSystemVersion = 0;
  ioh.MajorImageVersion = 0;
  ioh.MinorImageVersion = 0;
  ioh.MajorSubsystemVersion = 4;
  ioh.MinorSubsystemVersion = 0;
  ioh.Win32VersionValue = 0;
  ioh.SizeOfImage = (ioh.BaseOfData + ioh.SizeOfInitializedData + ALIGN - 1)
                                                            / ALIGN * ALIGN;
  ioh.SizeOfHeaders = ioh.BaseOfData;
  ioh.CheckSum = 0;
  ioh.Subsystem = IMAGE_SUBSYSTEM_WINDOWS_GUI;
  ioh.DllCharacteristics = IMAGE_DLLCHARACTERISTICS_NO_SEH;
  ioh.SizeOfStackReserve = 0x00100000;
  ioh.SizeOfStackCommit = 0x00001000;
  ioh.SizeOfHeapReserve = 0x00100000;
  ioh.SizeOfHeapCommit = 0x00001000;
  ioh.LoaderFlags = 0;
  ioh.NumberOfRvaAndSizes = IMAGE_NUMBEROF_DIRECTORY_ENTRIES;
  memset (ioh.DataDirectory, 0,
            IMAGE_NUMBEROF_DIRECTORY_ENTRIES * sizeof (IMAGE_DATA_DIRECTORY));
  ioh.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress = 0x200;
  ioh.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].Size
                                                = ioh.SizeOfInitializedData;
  memset (pish->Name, 0, IMAGE_SIZEOF_SHORT_NAME * sizeof (BYTE));
  pish->Misc.VirtualSize = ioh.SizeOfInitializedData;
  pish->VirtualAddress = ioh.BaseOfData;
  pish->SizeOfRawData
                    = (ioh.SizeOfInitializedData + ALIGN - 1) / ALIGN * ALIGN;
  pish->PointerToRawData = ioh.BaseOfData;
  pish->PointerToRelocations = 0;
  pish->PointerToLinenumbers = 0;
  pish->NumberOfRelocations = 0;
  pish->NumberOfLinenumbers = 0;
  pish->Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ;
  ied.Name = ioh.BaseOfData + sizeof (IMAGE_EXPORT_DIRECTORY)
                                    + ied.NumberOfFunctions * sizeof (DWORD)
                                    + ied.NumberOfNames * sizeof (DWORD)
                                    + ied.NumberOfNames * sizeof (WORD);
  ied.AddressOfFunctions = ioh.BaseOfData + sizeof (IMAGE_EXPORT_DIRECTORY);
  ied.AddressOfNames = ied.AddressOfFunctions
                                    + ied.NumberOfFunctions * sizeof (DWORD);
  ied.AddressOfNameOrdinals = ied.AddressOfNames
                                    + ied.NumberOfNames * sizeof (DWORD);

  fp = fopen (argv[2], "wb");
  if (!fp || fwrite (&idh, sizeof (IMAGE_DOS_HEADER), 1, fp) != 1
        || fwrite (bCode, sizeof (BYTE), sizeof (bCode) / sizeof (BYTE), fp)
                                            != sizeof (bCode) / sizeof (BYTE)
        || fwrite (&dwSignature, sizeof (DWORD), 1, fp) != 1
        || fwrite (&ifh, sizeof (IMAGE_FILE_HEADER), 1, fp) != 1
        || fwrite (&ioh, sizeof (IMAGE_OPTIONAL_HEADER), 1, fp) != 1
        || fwrite (pish, sizeof (IMAGE_SECTION_HEADER), 1, fp) != 1)
    return -1;
  free (pish);
  for (i = ioh.BaseOfData - dwOffset; i > 0; i--)
    if (fputc (0,fp) == EOF)
      return -1;
  if (fwrite (&ied, sizeof (IMAGE_EXPORT_DIRECTORY), 1, fp) != 1
                || fseek (fp, ied.NumberOfFunctions * sizeof (DWORD)
                        + ied.NumberOfNames * sizeof (DWORD), SEEK_CUR) != 0
                || fwrite (lpwOrdinal, sizeof (WORD), ied.NumberOfNames, fp)
                                                        != ied.NumberOfNames
                || fwrite (path, sizeof (char), lstrlenA (path) + 1, fp)
                                                        != lstrlenA (path) + 1)
    return -1;
  dwOffset = ioh.BaseOfData + sizeof (IMAGE_EXPORT_DIRECTORY)
                                    + ied.NumberOfFunctions * sizeof (DWORD)
                                    + ied.NumberOfNames * sizeof (DWORD)
                                    + ied.NumberOfNames * sizeof (WORD)
                                    + (lstrlenA (path) + 1) * sizeof (char);
  free (path);
  lpdwAddress = malloc (ied.NumberOfFunctions * sizeof (DWORD));
  memset (lpdwAddress, 0, ied.NumberOfFunctions * sizeof (DWORD));
  for (i = 0; i < ied.NumberOfNames; i++)
    {
      size_t len;

      if (name[i])
        {
          lpdwName[i] = dwOffset;
          len = lstrlenA (name[i]) + 1;
          if (fwrite (name[i], sizeof (char), len, fp) != len)
            return -1;
          dwOffset += len * sizeof (char);
        }
      else
        {
          lpdwName[i] = 0;
        }
      lpdwAddress[lpwOrdinal[i]] = dwOffset;
      len = lstrlenA (file);
      if (fwrite (file, sizeof (char), len, fp) != len
                                                    || fputc ('.', fp) == EOF)
        return -1;
      dwOffset += (len + 1) * sizeof (char);
      if (name[i])
        {
          len = lstrlenA (name[i]) + 1;
          if (fwrite (name[i], sizeof (char), len, fp) != len)
            return -1;
          dwOffset += len * sizeof (char);
          free (name[i]);
        }
      else
        {
          char str[7];

          wsprintfA (str, "#%u", lpwOrdinal[i]);
          len = lstrlenA (str) + 1;
          if (fwrite (str, sizeof (char), len, fp) != len)
            return -1;
          dwOffset += len * sizeof (char);
        }
    }
  free (name);
  free (file);
  free (lpwOrdinal);
  for (i = ioh.SizeOfImage - dwOffset; i > 0; i--)
    if (fputc (0,fp) == EOF)
      return -1;
  if (fseek (fp, ioh.BaseOfData + sizeof (IMAGE_EXPORT_DIRECTORY), SEEK_SET)
                                                                        != 0
            || fwrite (lpdwAddress, sizeof (DWORD), ied.NumberOfFunctions, fp)
                                                    != ied.NumberOfFunctions
            || fwrite (lpdwName, sizeof (DWORD), ied.NumberOfNames, fp)
                                                    != ied.NumberOfNames)
    return -1;
  return fclose (fp);
}
