home *** CD-ROM | disk | FTP | other *** search
/ Chip 2004 March / CMCD0304.ISO / Software / Freeware / Programare / nullsoft / nsis20.exe / Source / ResourceEditor.cpp < prev    next >
Encoding:
C/C++ Source or Header  |  2004-02-05  |  29.9 KB  |  828 lines
  1. /*
  2.   Copyright (C) 2002 Amir Szekely <kichik@netvision.net.il>
  3.  
  4.   This software is provided 'as-is', without any express or implied
  5.   warranty.  In no event will the authors be held liable for any damages
  6.   arising from the use of this software.
  7.  
  8.   Permission is granted to anyone to use this software for any purpose,
  9.   including commercial applications, and to alter it and redistribute it
  10.   freely, subject to the following restrictions:
  11.  
  12.   1. The origin of this software must not be misrepresented; you must not
  13.   claim that you wrote the original software. If you use this software
  14.   in a product, an acknowledgment in the product documentation would be
  15.   appreciated but is not required.
  16.  
  17.   2. Altered source versions must be plainly marked as such, and must not be
  18.   misrepresented as being the original software.
  19.  
  20.   3. This notice may not be removed or altered from any source distribution.
  21. */
  22.  
  23. #define RESOURCE_EDITOR_NOT_API
  24. #include "ResourceEditor.h"
  25.  
  26. //////////////////////////////////////////////////////////////////////
  27. // Utilities
  28. //////////////////////////////////////////////////////////////////////
  29.  
  30. #define ALIGN(dwToAlign, dwAlignOn) dwToAlign = (dwToAlign%dwAlignOn == 0) ? dwToAlign : dwToAlign - (dwToAlign%dwAlignOn) + dwAlignOn
  31. #define RALIGN(dwToAlign, dwAlignOn) ((dwToAlign%dwAlignOn == 0) ? dwToAlign : dwToAlign - (dwToAlign%dwAlignOn) + dwAlignOn)
  32.  
  33. //////////////////////////////////////////////////////////////////////
  34. //////////////////////////////////////////////////////////////////////
  35. // CResourceEditor
  36. //////////////////////////////////////////////////////////////////////
  37. //////////////////////////////////////////////////////////////////////
  38.  
  39. //////////////////////////////////////////////////////////////////////
  40. // Construction/Destruction
  41. //////////////////////////////////////////////////////////////////////
  42.  
  43. CResourceEditor::CResourceEditor(BYTE* pbPE, int iSize) {
  44.   // Copy the data pointer
  45.   m_pbPE = pbPE;
  46.   m_iSize = iSize;
  47.  
  48.   // Get dos header
  49.   m_dosHeader = (PIMAGE_DOS_HEADER)m_pbPE;
  50.   if (m_dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
  51.     throw runtime_error("PE file contains invalid DOS header");
  52.  
  53.   // Get NT headers
  54.   m_ntHeaders = (PIMAGE_NT_HEADERS)(m_pbPE + m_dosHeader->e_lfanew);
  55.   if (m_ntHeaders->Signature != IMAGE_NT_SIGNATURE)
  56.     throw runtime_error("PE file missing NT signature");
  57.  
  58.   // No check sum support yet...
  59.   if (m_ntHeaders->OptionalHeader.CheckSum)
  60.     throw runtime_error("CResourceEditor doesn't yet support check sum");
  61.  
  62.   // Get resource section virtual address
  63.   m_dwResourceSectionVA = m_ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress;
  64.   // Pointer to the sections headers array
  65.   PIMAGE_SECTION_HEADER sectionHeadersArray = IMAGE_FIRST_SECTION(m_ntHeaders);
  66.  
  67.   m_dwResourceSectionIndex = -1;
  68.  
  69.   // Find resource section index in the array
  70.   for (int i = 0; i < m_ntHeaders->FileHeader.NumberOfSections; i++) {
  71.     if (m_dwResourceSectionVA == sectionHeadersArray[i].VirtualAddress) {
  72.       // Remember resource section index
  73.       m_dwResourceSectionIndex = i;
  74.       // Check for invalid resource section pointer
  75.       if (!sectionHeadersArray[i].PointerToRawData)
  76.         throw runtime_error("Invalid resource section pointer");
  77.     }
  78.  
  79.     // Invalid section pointer (goes beyond the PE image)
  80.     if (sectionHeadersArray[i].PointerToRawData > (unsigned int)m_iSize)
  81.       throw runtime_error("Invalid section pointer");
  82.   }
  83.  
  84.   // No resource section...
  85.   if (m_dwResourceSectionIndex == m_ntHeaders->FileHeader.NumberOfSections)
  86.     throw runtime_error("PE file doesn't contain any resource section");
  87.  
  88.   // Pointer to section data, the first resource directory
  89.   PRESOURCE_DIRECTORY rdRoot = PRESOURCE_DIRECTORY(m_pbPE + sectionHeadersArray[m_dwResourceSectionIndex].PointerToRawData);
  90.  
  91.   // Scan the resource directory
  92.   m_cResDir = ScanDirectory(rdRoot, rdRoot);
  93. }
  94.  
  95. CResourceEditor::~CResourceEditor() {
  96.   if (m_cResDir) {
  97.     m_cResDir->Destroy();
  98.     delete m_cResDir;
  99.   }
  100. }
  101.  
  102. //////////////////////////////////////////////////////////////////////
  103. // Methods
  104. //////////////////////////////////////////////////////////////////////
  105.  
  106. // Adds/Replaces/Removes a resource.
  107. // If lpData is 0 UpdateResource removes the resource.
  108. bool CResourceEditor::UpdateResource(char* szType, char* szName, LANGID wLanguage, BYTE* lpData, DWORD dwSize) {
  109.   CResourceDirectory* nameDir = 0;
  110.   CResourceDirectory* langDir = 0;
  111.   CResourceDataEntry* data = 0;
  112.   IMAGE_RESOURCE_DIRECTORY rd = {0, /*time(0),*/};
  113.   int iTypeIdx, iNameIdx, iLangIdx;
  114.  
  115.   iTypeIdx = m_cResDir->Find(szType);
  116.   if (iTypeIdx > -1) {
  117.     nameDir = m_cResDir->GetEntry(iTypeIdx)->GetSubDirectory();
  118.     iNameIdx = nameDir->Find(szName);
  119.     if (iNameIdx > -1) {
  120.       langDir = nameDir->GetEntry(iNameIdx)->GetSubDirectory();
  121.       iLangIdx = langDir->Find(wLanguage);
  122.       if (iLangIdx > -1) {
  123.         data = langDir->GetEntry(iLangIdx)->GetDataEntry();
  124.       }
  125.     }
  126.   }
  127.   
  128.   if (lpData) {
  129.     // Replace/Add the resource
  130.     if (data) {
  131.       data->SetData(lpData, dwSize);
  132.       return true;
  133.     }
  134.  
  135.     if (!nameDir) {
  136.       // Type doesn't yet exist
  137.       nameDir = new CResourceDirectory(&rd);
  138.       m_cResDir->AddEntry(new CResourceDirectoryEntry(szType, nameDir));
  139.     }
  140.     if (!langDir) {
  141.       // Name doesn't yet exist
  142.       langDir = new CResourceDirectory(&rd);
  143.       nameDir->AddEntry(new CResourceDirectoryEntry(szName, langDir));
  144.     }
  145.     if (!data) {
  146.       // Language doesn't yet exist, hence data nither
  147.       data = new CResourceDataEntry(lpData, dwSize);
  148.       langDir->AddEntry(new CResourceDirectoryEntry(MAKEINTRESOURCE(wLanguage), data));
  149.     }
  150.   }
  151.   else if (data) {
  152.     // Delete the resource
  153.     delete data;
  154.     langDir->RemoveEntry(iLangIdx);
  155.     // Delete directories holding the resource if empty
  156.     if (!langDir->CountEntries()) {
  157.       delete langDir;
  158.       nameDir->RemoveEntry(iNameIdx);
  159.       if (!nameDir->CountEntries()) {
  160.         delete nameDir;
  161.         m_cResDir->RemoveEntry(iTypeIdx);
  162.       }
  163.     }
  164.   }
  165.   else return false;
  166.   return true;
  167. }
  168.  
  169. bool CResourceEditor::UpdateResource(WORD szType, char* szName, LANGID wLanguage, BYTE* lpData, DWORD dwSize) {
  170.   return UpdateResource(MAKEINTRESOURCE(szType), szName, wLanguage, lpData, dwSize);
  171. }
  172.  
  173. bool CResourceEditor::UpdateResource(char* szType, WORD szName, LANGID wLanguage, BYTE* lpData, DWORD dwSize) {
  174.   return UpdateResource(szType, MAKEINTRESOURCE(szName), wLanguage, lpData, dwSize);
  175. }
  176.  
  177. bool CResourceEditor::UpdateResource(WORD szType, WORD szName, LANGID wLanguage, BYTE* lpData, DWORD dwSize) {
  178.   return UpdateResource(MAKEINTRESOURCE(szType), MAKEINTRESOURCE(szName), wLanguage, lpData, dwSize);
  179. }
  180.  
  181. // Returns a copy of the resource requested
  182. // Returns 0 if resource can't be found
  183. BYTE* CResourceEditor::GetResource(char* szType, char* szName, LANGID wLanguage) {
  184.   CResourceDirectory* nameDir = 0;
  185.   CResourceDirectory* langDir = 0;
  186.   CResourceDataEntry* data = 0;
  187.  
  188.   int i = m_cResDir->Find(szType);
  189.   if (i > -1) {
  190.     nameDir = m_cResDir->GetEntry(i)->GetSubDirectory();
  191.     i = nameDir->Find(szName);
  192.     if (i > -1) {
  193.       langDir = nameDir->GetEntry(i)->GetSubDirectory();
  194.       i = langDir->Find(wLanguage);
  195.       if (i > -1) {
  196.         data = langDir->GetEntry(i)->GetDataEntry();
  197.       }
  198.     }
  199.   }
  200.  
  201.   if (data) {
  202.     BYTE* toReturn = new BYTE[data->GetSize()];
  203.     CopyMemory(toReturn, data->GetData(), data->GetSize());
  204.     return toReturn;
  205.   }
  206.   else
  207.     return 0;
  208. }
  209.  
  210. void CResourceEditor::FreeResource(BYTE* pbResource)
  211. {
  212.   if (pbResource)
  213.     delete [] pbResource;
  214. }
  215.  
  216. // Saves the edited PE into a buffer and returns it.
  217. DWORD CResourceEditor::Save(BYTE* pbBuf, DWORD &dwSize) {
  218.   unsigned int i;
  219.   DWORD dwReqSize;
  220.  
  221.   DWORD dwRsrcSize = m_cResDir->GetSize(); // Size of new resource section
  222.   DWORD dwRsrcSizeAligned = RALIGN(dwRsrcSize, m_ntHeaders->OptionalHeader.FileAlignment); // Align it to FileAlignment
  223.  
  224.   // Calculate the total new PE size
  225.   dwReqSize = m_iSize - IMAGE_FIRST_SECTION(m_ntHeaders)[m_dwResourceSectionIndex].SizeOfRawData + dwRsrcSizeAligned;
  226.  
  227.   if (!pbBuf || dwSize < dwReqSize)
  228.     return dwReqSize;
  229.  
  230.   // Use buffer
  231.   BYTE* pbNewPE = pbBuf;
  232.   dwSize = dwReqSize;
  233.   // Fill buffer with zeros
  234.   ZeroMemory(pbNewPE, dwSize);
  235.  
  236.   BYTE* seeker = pbNewPE;
  237.   BYTE* oldSeeker = m_pbPE;
  238.  
  239.   PIMAGE_SECTION_HEADER old_sectionHeadersArray = IMAGE_FIRST_SECTION(m_ntHeaders);
  240.  
  241.   // Copy everything until the resource section (including headers and everything that might come after them)
  242.   // We don't use SizeOfHeaders because sometimes (using VC6) it can extend beyond the first section
  243.   // or (Borland) there could be some more information between the headers and the first section.
  244.   CopyMemory(seeker, oldSeeker, old_sectionHeadersArray[m_dwResourceSectionIndex].PointerToRawData);
  245.  
  246.   // Skip the headers and whatever comes after them
  247.   seeker += old_sectionHeadersArray[m_dwResourceSectionIndex].PointerToRawData;
  248.   oldSeeker += old_sectionHeadersArray[m_dwResourceSectionIndex].PointerToRawData;
  249.  
  250.   // Get new nt headers pointer
  251.   PIMAGE_NT_HEADERS ntHeaders = PIMAGE_NT_HEADERS(pbNewPE + PIMAGE_DOS_HEADER(pbNewPE)->e_lfanew);
  252.   // Get a pointer to the new section headers
  253.   PIMAGE_SECTION_HEADER sectionHeadersArray = IMAGE_FIRST_SECTION(ntHeaders);
  254.  
  255.   // Skip the resource section in the old PE seeker.
  256.   oldSeeker += sectionHeadersArray[m_dwResourceSectionIndex].SizeOfRawData;
  257.  
  258.   // Save the old virtual size of the resource section
  259.   DWORD dwOldVirtualSize = sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize;
  260.  
  261.   // Set the new size of the resource section (size aligned to FileAlignment)
  262.   sectionHeadersArray[m_dwResourceSectionIndex].SizeOfRawData = dwRsrcSizeAligned;
  263.   // Set the virtual size as well (in memory)
  264.   sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize = RALIGN(dwRsrcSize, ntHeaders->OptionalHeader.SectionAlignment);
  265.   ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].Size = sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize;
  266.  
  267.   // Set the new virtual size of the image
  268.   DWORD old = ntHeaders->OptionalHeader.SizeOfImage;
  269.   ntHeaders->OptionalHeader.SizeOfImage = RALIGN(ntHeaders->OptionalHeader.SizeOfHeaders, ntHeaders->OptionalHeader.SectionAlignment);
  270.   for (i = 0; i < ntHeaders->FileHeader.NumberOfSections; i++)
  271.     ntHeaders->OptionalHeader.SizeOfImage += RALIGN(sectionHeadersArray[i].Misc.VirtualSize, ntHeaders->OptionalHeader.SectionAlignment);
  272.  
  273.   // Set the new AddressOfEntryPoint if needed
  274.   if (ntHeaders->OptionalHeader.AddressOfEntryPoint > sectionHeadersArray[m_dwResourceSectionIndex].VirtualAddress)
  275.     ntHeaders->OptionalHeader.AddressOfEntryPoint += sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize - dwOldVirtualSize;
  276.  
  277.   // Set the new BaseOfCode if needed
  278.   if (ntHeaders->OptionalHeader.BaseOfCode > sectionHeadersArray[m_dwResourceSectionIndex].VirtualAddress)
  279.     ntHeaders->OptionalHeader.BaseOfCode += sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize - dwOldVirtualSize;
  280.  
  281.   // Set the new BaseOfData if needed
  282.   if (ntHeaders->OptionalHeader.BaseOfData > sectionHeadersArray[m_dwResourceSectionIndex].VirtualAddress)
  283.     ntHeaders->OptionalHeader.BaseOfData += sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize - dwOldVirtualSize;
  284.  
  285.   // Refresh the headers of the sections that come after the resource section, and the data directory
  286.   for (i = m_dwResourceSectionIndex + 1; i < ntHeaders->FileHeader.NumberOfSections; i++) {
  287.     if (sectionHeadersArray[i].PointerToRawData) {
  288.       sectionHeadersArray[i].PointerToRawData -= IMAGE_FIRST_SECTION(m_ntHeaders)[m_dwResourceSectionIndex].SizeOfRawData;
  289.       sectionHeadersArray[i].PointerToRawData += dwRsrcSizeAligned;
  290.     }
  291.  
  292.     // We must find the right data directory entry before we change the virtual address
  293.     unsigned int uDataDirIdx = 0;
  294.     for (unsigned int j = 0; j < ntHeaders->OptionalHeader.NumberOfRvaAndSizes; j++)
  295.       if (ntHeaders->OptionalHeader.DataDirectory[j].VirtualAddress == sectionHeadersArray[i].VirtualAddress)
  296.         uDataDirIdx = j;
  297.  
  298.     sectionHeadersArray[i].VirtualAddress -= RALIGN(dwOldVirtualSize, ntHeaders->OptionalHeader.SectionAlignment);
  299.     sectionHeadersArray[i].VirtualAddress += RALIGN(sectionHeadersArray[m_dwResourceSectionIndex].Misc.VirtualSize, ntHeaders->OptionalHeader.SectionAlignment);
  300.  
  301.     // Change the virtual address in the data directory too
  302.     if (uDataDirIdx)
  303.       ntHeaders->OptionalHeader.DataDirectory[uDataDirIdx].VirtualAddress = sectionHeadersArray[i].VirtualAddress;
  304.   }
  305.  
  306.   // Write the resource section
  307.   WriteRsrcSec(seeker);
  308.   // Advance the pointer
  309.   seeker += dwRsrcSizeAligned;
  310.  
  311.   // Copy everything that comes after the resource section (other sections and tacked data)
  312.   DWORD dwLeft = m_iSize - (oldSeeker - m_pbPE);
  313.   if (dwLeft)
  314.     CopyMemory(seeker, oldSeeker, dwLeft);
  315.  
  316.   seeker += dwLeft;
  317.   oldSeeker += dwLeft;
  318.  
  319.   /**********************************************************
  320.    * To add checksum to the header use MapFileAndCheckSum
  321.    **********************************************************/
  322.  
  323.   // From now on, we are working on the new PE
  324.   // Freeing the old PE memory is up to the user
  325.   m_pbPE = pbNewPE;
  326.   m_iSize = dwSize;
  327.   m_dosHeader = PIMAGE_DOS_HEADER(m_pbPE);
  328.   m_ntHeaders = ntHeaders;
  329.   // We just wrote the resource section according to m_cResDir, so we don't need to rescan
  330.   // m_dwResourceSectionIndex and m_dwResourceSectionVA have also been left unchanged as
  331.   // we didn't move the resources section
  332.  
  333.   return 0;
  334. }
  335.  
  336. // This function scans exe sections and after find a match with given name
  337. // increments it's virtual size (auto fixes image size based on section alignment, etc)
  338. bool CResourceEditor::AddExtraVirtualSize2PESection(const char* pszSectionName, int addsize)
  339. {
  340.   PIMAGE_SECTION_HEADER sectionHeadersArray = IMAGE_FIRST_SECTION(m_ntHeaders);
  341.  
  342.   // Refresh the headers of the sections that come after the resource section, and the data directory
  343.   for (int i =0; i < m_ntHeaders->FileHeader.NumberOfSections; i++) {
  344.     if ( !strcmp((LPCSTR)sectionHeadersArray[i].Name, pszSectionName) ) {
  345.       sectionHeadersArray[i].Misc.VirtualSize += addsize;
  346.       sectionHeadersArray[i].Characteristics &= ~IMAGE_SCN_MEM_DISCARDABLE;
  347.       sectionHeadersArray[i].Misc.VirtualSize = RALIGN(sectionHeadersArray[i].Misc.VirtualSize, m_ntHeaders->OptionalHeader.SectionAlignment);
  348.       // now fix any section after
  349.       for (int k=i+1; k< m_ntHeaders->FileHeader.NumberOfSections; k++, i++) {
  350.         sectionHeadersArray[k].VirtualAddress = sectionHeadersArray[i].VirtualAddress + sectionHeadersArray[i].Misc.VirtualSize;
  351.         sectionHeadersArray[k].VirtualAddress = RALIGN(sectionHeadersArray[k].VirtualAddress, m_ntHeaders->OptionalHeader.SectionAlignment);
  352.         if ( m_dwResourceSectionIndex == k )
  353.         {
  354.           // fix the resources virtual address if it changed
  355.           m_dwResourceSectionVA = m_ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress = sectionHeadersArray[k].VirtualAddress;
  356.         }
  357.       }
  358.  
  359.       return true;
  360.     }
  361.   }
  362.  
  363.   return false;
  364. }
  365.  
  366. //////////////////////////////////////////////////////////////////////
  367. // Private Methods
  368. //////////////////////////////////////////////////////////////////////
  369.  
  370. // This function scans a give resource directory and return a CResourceDirectory object
  371. // rdRoot must point to the root directory of the resource section
  372. CResourceDirectory* CResourceEditor::ScanDirectory(PRESOURCE_DIRECTORY rdRoot, PRESOURCE_DIRECTORY rdToScan) {
  373.   // Create CResourceDirectory from rdToScan
  374.   CResourceDirectory* rdc = new CResourceDirectory(PIMAGE_RESOURCE_DIRECTORY(rdToScan));
  375.   char* szName;
  376.   PIMAGE_RESOURCE_DATA_ENTRY rde;
  377.  
  378.   // Go through all entries of this resource directory
  379.   for (int i = 0; i < rdToScan->Header.NumberOfNamedEntries + rdToScan->Header.NumberOfIdEntries; i++) {
  380.     // If this entry points to data entry get a pointer to it
  381.     if (!rdToScan->Entries[i].DataIsDirectory)
  382.       rde = PIMAGE_RESOURCE_DATA_ENTRY(rdToScan->Entries[i].OffsetToData + (BYTE*)rdRoot);
  383.  
  384.     // If this entry has a name, translate it from Unicode
  385.     if (rdToScan->Entries[i].NameIsString) {
  386.       PIMAGE_RESOURCE_DIR_STRING_U rds = PIMAGE_RESOURCE_DIR_STRING_U(rdToScan->Entries[i].NameOffset + (char*)rdRoot);
  387.  
  388.       int mbsSize = WideCharToMultiByte(CP_ACP, 0, rds->NameString, rds->Length, 0, 0, 0, 0);
  389.       szName = new char[mbsSize+1];
  390.       WideCharToMultiByte(CP_ACP, 0, rds->NameString, rds->Length, szName, mbsSize, 0, 0);
  391.       szName[mbsSize] = 0;
  392.     }
  393.     // Else, set the name to this entry's id
  394.     else
  395.       szName = MAKEINTRESOURCE(rdToScan->Entries[i].Id);
  396.  
  397.     if (rdToScan->Entries[i].DataIsDirectory)
  398.       rdc->AddEntry(
  399.         new CResourceDirectoryEntry(
  400.           szName,
  401.           ScanDirectory(
  402.             rdRoot,
  403.             PRESOURCE_DIRECTORY(rdToScan->Entries[i].OffsetToDirectory + (BYTE*)rdRoot)
  404.           )
  405.         )
  406.       );
  407.     else
  408.       rdc->AddEntry(
  409.         new CResourceDirectoryEntry(
  410.           szName,
  411.           new CResourceDataEntry(
  412.             (BYTE*)rdRoot + rde->OffsetToData - m_dwResourceSectionVA,
  413.             rde->Size,
  414.             rde->CodePage
  415.           )
  416.         )
  417.       );
  418.  
  419.     // Delete the dynamicly allocated name if it is a name and not an id
  420.     if (!IS_INTRESOURCE(szName))
  421.       delete [] szName;
  422.   }
  423.  
  424.   return rdc;
  425. }
  426.  
  427. // This function writes into a given place in memory (pbRsrcSec) the edited resource section
  428. void CResourceEditor::WriteRsrcSec(BYTE* pbRsrcSec) {
  429.   BYTE* seeker = pbRsrcSec;
  430.  
  431.   queue<CResourceDirectory*> qDirs; // Used to scan the tree by level
  432.   queue<CResourceDataEntry*> qDataEntries; // Used for writing the data entries
  433.   queue<CResourceDataEntry*> qDataEntries2; // Used for writing raw resources data
  434.   queue<CResourceDirectoryEntry*> qStrings; // Used for writing resources' names
  435.  
  436.   qDirs.push(m_cResDir);
  437.  
  438.   while (!qDirs.empty()) {
  439.     CResourceDirectory* crd = qDirs.front();
  440.  
  441.     CopyMemory(seeker, &crd->GetInfo(), sizeof(IMAGE_RESOURCE_DIRECTORY));
  442.     crd->m_dwWrittenAt = DWORD(seeker);
  443.     seeker += sizeof(IMAGE_RESOURCE_DIRECTORY);
  444.  
  445.     for (int i = 0; i < crd->CountEntries(); i++) {
  446.       if (crd->GetEntry(i)->HasName())
  447.         qStrings.push(crd->GetEntry(i));
  448.       if (crd->GetEntry(i)->IsDataDirectory())
  449.         qDirs.push(crd->GetEntry(i)->GetSubDirectory());
  450.       else {
  451.         qDataEntries.push(crd->GetEntry(i)->GetDataEntry());
  452.         qDataEntries2.push(crd->GetEntry(i)->GetDataEntry());
  453.       }
  454.  
  455.       IMAGE_RESOURCE_DIRECTORY_ENTRY rDirE = {0,};
  456.       rDirE.DataIsDirectory = crd->GetEntry(i)->IsDataDirectory();
  457.       rDirE.Id = (crd->GetEntry(i)->HasName()) ? 0 : crd->GetEntry(i)->GetId();
  458.       rDirE.NameIsString = (crd->GetEntry(i)->HasName()) ? 1 : 0;
  459.  
  460.       CopyMemory(seeker, &rDirE, sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY));
  461.       crd->GetEntry(i)->m_dwWrittenAt = DWORD(seeker);
  462.       seeker += sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY);
  463.     }
  464.     qDirs.pop();
  465.   }
  466.  
  467.   /*
  468.    * Write IMAGE_RESOURCE_DATA_ENTRYs.  
  469.    */
  470.   while (!qDataEntries.empty()) {
  471.     CResourceDataEntry* cRDataE = qDataEntries.front();
  472.     IMAGE_RESOURCE_DATA_ENTRY rDataE = {0,};
  473.     rDataE.CodePage = cRDataE->GetCodePage();
  474.     rDataE.Size = cRDataE->GetSize();
  475.  
  476.     CopyMemory(seeker, &rDataE, sizeof(IMAGE_RESOURCE_DATA_ENTRY));
  477.     cRDataE->m_dwWrittenAt = DWORD(seeker);
  478.     seeker += sizeof(IMAGE_RESOURCE_DATA_ENTRY);
  479.  
  480.     qDataEntries.pop();
  481.   }
  482.  
  483.   /*
  484.    * Write strings
  485.    */
  486.   while (!qStrings.empty()) {
  487.     CResourceDirectoryEntry* cRDirE = qStrings.front();
  488.  
  489.     PIMAGE_RESOURCE_DIRECTORY_ENTRY(cRDirE->m_dwWrittenAt)->NameOffset = DWORD(seeker) - DWORD(pbRsrcSec);
  490.  
  491.     char* szName = cRDirE->GetName();
  492.     WORD iLen = lstrlen(szName);
  493.     WCHAR* szwName = new WCHAR[iLen+1];
  494.     // MultiByteToWideChar return value includes the null char, so -1
  495.     iLen = MultiByteToWideChar(CP_ACP, 0, szName, iLen, szwName, iLen) - 1;
  496.     *(WORD*)seeker = iLen;
  497.     seeker += sizeof(WORD);
  498.     CopyMemory(seeker, szwName, iLen*sizeof(WCHAR));
  499.     seeker += iLen*sizeof(WCHAR);
  500.  
  501.     // Even though the number of chars is predefined a null termination is required
  502.     *(WORD*)seeker = 0;
  503.     seeker += sizeof(WORD);
  504.  
  505.     delete [] szName;
  506.     delete [] szwName;
  507.  
  508.     qStrings.pop();
  509.   }
  510.  
  511.   /*
  512.    * Write raw resource data and set offsets in IMAGE_RESOURCE_DATA_ENTRYs.
  513.    */
  514.   while (!qDataEntries2.empty()) {
  515.     CResourceDataEntry* cRDataE = qDataEntries2.front();
  516.     CopyMemory(seeker, cRDataE->GetData(), cRDataE->GetSize());
  517.     PIMAGE_RESOURCE_DATA_ENTRY(cRDataE->m_dwWrittenAt)->OffsetToData = seeker - pbRsrcSec + m_dwResourceSectionVA;
  518.  
  519.     seeker += RALIGN(cRDataE->GetSize(), 8);
  520.  
  521.     qDataEntries2.pop();
  522.   }
  523.  
  524.   /*
  525.    * Set all of the directory entries offsets.
  526.    */
  527.   SetOffsets(m_cResDir, DWORD(pbRsrcSec));
  528. }
  529.  
  530. // Sets the offsets in directory entries
  531. void CResourceEditor::SetOffsets(CResourceDirectory* resDir, DWORD newResDirAt) {
  532.   for (int i = 0; i < resDir->CountEntries(); i++) {
  533.     if (resDir->GetEntry(i)->IsDataDirectory()) {
  534.       PIMAGE_RESOURCE_DIRECTORY_ENTRY(resDir->GetEntry(i)->m_dwWrittenAt)->DataIsDirectory = 1;
  535.       PIMAGE_RESOURCE_DIRECTORY_ENTRY(resDir->GetEntry(i)->m_dwWrittenAt)->OffsetToDirectory = resDir->GetEntry(i)->GetSubDirectory()->m_dwWrittenAt - newResDirAt;
  536.       SetOffsets(resDir->GetEntry(i)->GetSubDirectory(), newResDirAt);
  537.     }
  538.     else {
  539.       PIMAGE_RESOURCE_DIRECTORY_ENTRY(resDir->GetEntry(i)->m_dwWrittenAt)->OffsetToData = resDir->GetEntry(i)->GetDataEntry()->m_dwWrittenAt - newResDirAt;
  540.     }
  541.   }
  542. }
  543.  
  544. //////////////////////////////////////////////////////////////////////
  545. //////////////////////////////////////////////////////////////////////
  546. // CResourceDirectory
  547. //////////////////////////////////////////////////////////////////////
  548. //////////////////////////////////////////////////////////////////////
  549.  
  550. //////////////////////////////////////////////////////////////////////
  551. // Construction/Destruction
  552. //////////////////////////////////////////////////////////////////////
  553.  
  554. CResourceDirectory::CResourceDirectory(PIMAGE_RESOURCE_DIRECTORY prd) {
  555.   m_rdDir = *prd;
  556.   m_rdDir.NumberOfIdEntries = 0;
  557.   m_rdDir.NumberOfNamedEntries = 0;
  558. }
  559.  
  560. CResourceDirectory::~CResourceDirectory() {
  561. }
  562.  
  563. //////////////////////////////////////////////////////////////////////
  564. // Methods
  565. //////////////////////////////////////////////////////////////////////
  566.  
  567. IMAGE_RESOURCE_DIRECTORY CResourceDirectory::GetInfo() {
  568.   return m_rdDir;
  569. }
  570.  
  571. CResourceDirectoryEntry* CResourceDirectory::GetEntry(unsigned int i) {
  572.   if (m_vEntries.size() < i)
  573.     return 0;
  574.   return m_vEntries[i];
  575. }
  576.  
  577. // This function inserts a new directory entry
  578. // It also keeps the directory entries sorted
  579. void CResourceDirectory::AddEntry(CResourceDirectoryEntry* entry) {
  580.   int i = 0;
  581.   if (entry->HasName()) {
  582.     char* szEntName = entry->GetName();
  583.     for (i = 0; i < m_rdDir.NumberOfIdEntries; i++) {
  584.       char* szName = m_vEntries[i]->GetName();
  585.       int cmp = lstrcmp(szName, szEntName);
  586.       delete [] szName;
  587.       if (cmp == 0) {
  588.         delete [] szEntName;
  589.         return;
  590.       }
  591.       if (cmp > 0)
  592.         break;
  593.     }
  594.     delete [] szEntName;
  595.     m_rdDir.NumberOfNamedEntries++;
  596.   }
  597.   else {
  598.     for (i = m_rdDir.NumberOfNamedEntries; i < m_rdDir.NumberOfNamedEntries+m_rdDir.NumberOfIdEntries; i++) {
  599.       if (m_vEntries[i]->GetId() == entry->GetId()) 
  600.         return;
  601.       if (m_vEntries[i]->GetId() > entry->GetId())
  602.         break;
  603.     }
  604.     m_rdDir.NumberOfIdEntries++;
  605.   }
  606.   m_vEntries.insert(m_vEntries.begin() + i, entry);
  607. }
  608.  
  609. void CResourceDirectory::RemoveEntry(int i) {
  610.   if (m_vEntries[i]->HasName())
  611.     m_rdDir.NumberOfNamedEntries--;
  612.   else
  613.     m_rdDir.NumberOfIdEntries--;
  614.   delete m_vEntries[i];
  615.   m_vEntries.erase(m_vEntries.begin() + i);
  616. }
  617.  
  618. int CResourceDirectory::CountEntries() {
  619.   return m_vEntries.size();
  620. }
  621.  
  622. // Returns the index of a directory entry with the specified name
  623. // Name can be a string or an id
  624. // Returns -1 if can not be found
  625. int CResourceDirectory::Find(char* szName) {
  626.   if (IS_INTRESOURCE(szName))
  627.     return Find(WORD(szName));
  628.   else
  629.     if (szName[0] == '#')
  630.       return Find(WORD(atoi(szName+1)));
  631.  
  632.   for (unsigned int i = 0; i < m_vEntries.size(); i++) {
  633.     if (!m_vEntries[i]->HasName())
  634.       continue;
  635.  
  636.     char* szEntName = m_vEntries[i]->GetName();
  637.     int cmp = lstrcmp(szName, szEntName);
  638.     delete [] szEntName;
  639.  
  640.     if (!cmp)
  641.       return i;
  642.   }
  643.  
  644.   return -1;
  645. }
  646.  
  647. // Returns the index of a directory entry with the specified id
  648. // Returns -1 if can not be found
  649. int CResourceDirectory::Find(WORD wId) {
  650.   for (unsigned int i = 0; i < m_vEntries.size(); i++) {
  651.     if (m_vEntries[i]->HasName())
  652.       continue;
  653.  
  654.     if (wId == m_vEntries[i]->GetId())
  655.       return i;
  656.   }
  657.  
  658.   return -1;
  659. }
  660.  
  661. // Get the size of this resource directory (including all of its children)
  662. DWORD CResourceDirectory::GetSize() {
  663.   DWORD dwSize = sizeof(IMAGE_RESOURCE_DIRECTORY);
  664.   for (unsigned int i = 0; i < m_vEntries.size(); i++) {
  665.     dwSize += sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY);
  666.     if (m_vEntries[i]->HasName())
  667.       dwSize += sizeof(IMAGE_RESOURCE_DIR_STRING_U) + (m_vEntries[i]->GetNameLength()+1)*sizeof(WCHAR);
  668.     if (m_vEntries[i]->IsDataDirectory())
  669.       dwSize += m_vEntries[i]->GetSubDirectory()->GetSize();
  670.     else {
  671.       DWORD dwAligned = m_vEntries[i]->GetDataEntry()->GetSize();
  672.       ALIGN(dwAligned, 8);
  673.       dwSize += sizeof(IMAGE_RESOURCE_DATA_ENTRY) + dwAligned;
  674.     }
  675.   }
  676.   return dwSize;
  677. }
  678.  
  679. // Destroys this directory and all of its children
  680. void CResourceDirectory::Destroy() {
  681.   for (unsigned int i = 0; i < m_vEntries.size(); i++) {
  682.     if (m_vEntries[i]->IsDataDirectory()) {
  683.       m_vEntries[i]->GetSubDirectory()->Destroy();
  684.       delete m_vEntries[i]->GetSubDirectory();
  685.     }
  686.     else
  687.       delete m_vEntries[i]->GetDataEntry();
  688.   }
  689. }
  690.  
  691. //////////////////////////////////////////////////////////////////////
  692. //////////////////////////////////////////////////////////////////////
  693. // CResourceDirectoryEntry
  694. //////////////////////////////////////////////////////////////////////
  695. //////////////////////////////////////////////////////////////////////
  696.  
  697. //////////////////////////////////////////////////////////////////////
  698. // Construction/Destruction
  699. //////////////////////////////////////////////////////////////////////
  700.  
  701. CResourceDirectoryEntry::CResourceDirectoryEntry(char* szName, CResourceDirectory* rdSubDir) {
  702.   if (IS_INTRESOURCE(szName)) {
  703.     m_bHasName = false;
  704.     m_szName = 0;
  705.     m_wId = WORD(szName);
  706.   }
  707.   else {
  708.     m_bHasName = true;
  709.     m_szName = new char[lstrlen(szName)+1];
  710.     lstrcpy(m_szName, szName);
  711.   }
  712.   m_bIsDataDirectory = true;
  713.   m_rdSubDir = rdSubDir;
  714. }
  715.  
  716. CResourceDirectoryEntry::CResourceDirectoryEntry(char* szName, CResourceDataEntry* rdeData) {
  717.   if (IS_INTRESOURCE(szName)) {
  718.     m_bHasName = false;
  719.     m_szName = 0;
  720.     m_wId = WORD(szName);
  721.   }
  722.   else {
  723.     m_bHasName = true;
  724.     m_szName = new char[lstrlen(szName)+1];
  725.     lstrcpy(m_szName, szName);
  726.   }
  727.   m_bIsDataDirectory = false;
  728.   m_rdeData = rdeData;
  729. }
  730.  
  731. CResourceDirectoryEntry::~CResourceDirectoryEntry() {
  732.   if (m_szName && m_bHasName)
  733.     delete [] m_szName;
  734. }
  735.  
  736. //////////////////////////////////////////////////////////////////////
  737. // Methods
  738. //////////////////////////////////////////////////////////////////////
  739.  
  740. bool CResourceDirectoryEntry::HasName() {
  741.   return m_bHasName;
  742. }
  743.  
  744. // Don't forget to free the memory used by the string after usage!
  745. char* CResourceDirectoryEntry::GetName() {
  746.   if (!m_bHasName)
  747.     return 0;
  748.   char* szName = 0;
  749.   szName = new char[lstrlen(m_szName)+1];
  750.   lstrcpy(szName, m_szName);
  751.   return szName;
  752. }
  753.  
  754. int CResourceDirectoryEntry::GetNameLength() {
  755.   return lstrlen(m_szName);
  756. }
  757.  
  758. WORD CResourceDirectoryEntry::GetId() {
  759.   if (m_bHasName)
  760.     return 0;
  761.   return m_wId;
  762. }
  763.  
  764. bool CResourceDirectoryEntry::IsDataDirectory() {
  765.   return m_bIsDataDirectory;
  766. }
  767.  
  768. CResourceDirectory* CResourceDirectoryEntry::GetSubDirectory() {
  769.   if (!m_bIsDataDirectory)
  770.     return NULL;
  771.   return m_rdSubDir;
  772. }
  773.  
  774. CResourceDataEntry* CResourceDirectoryEntry::GetDataEntry() {
  775.   if (m_bIsDataDirectory)
  776.     return NULL;
  777.   return m_rdeData;
  778. }
  779.  
  780. //////////////////////////////////////////////////////////////////////
  781. //////////////////////////////////////////////////////////////////////
  782. // CResourceDataEntry
  783. //////////////////////////////////////////////////////////////////////
  784. //////////////////////////////////////////////////////////////////////
  785.  
  786. //////////////////////////////////////////////////////////////////////
  787. // Construction/Destruction
  788. //////////////////////////////////////////////////////////////////////
  789.  
  790. CResourceDataEntry::CResourceDataEntry(BYTE* pbData, DWORD dwSize, DWORD dwCodePage) {
  791.   m_pbData = 0;
  792.   SetData(pbData, dwSize, dwCodePage);
  793. }
  794.  
  795. CResourceDataEntry::~CResourceDataEntry() {
  796.   if (m_pbData)
  797.     delete [] m_pbData;
  798. }
  799.  
  800. //////////////////////////////////////////////////////////////////////
  801. // Methods
  802. //////////////////////////////////////////////////////////////////////
  803.  
  804. // To save memory this function doesn't give you a copy of the data
  805. // Don't mess with the data returned from this function!
  806. BYTE* CResourceDataEntry::GetData() {
  807.   return m_pbData;
  808. }
  809.  
  810. void CResourceDataEntry::SetData(BYTE* pbData, DWORD dwSize) {
  811.   SetData(pbData, dwSize, m_dwCodePage);
  812. }
  813.  
  814. void CResourceDataEntry::SetData(BYTE* pbData, DWORD dwSize, DWORD dwCodePage) {
  815.   if (m_pbData) delete [] m_pbData;
  816.   m_pbData = new BYTE[dwSize];
  817.   CopyMemory(m_pbData, pbData, dwSize);
  818.   m_dwSize = dwSize;
  819.   m_dwCodePage = dwCodePage;
  820. }
  821.  
  822. DWORD CResourceDataEntry::GetSize() {
  823.   return m_dwSize;
  824. }
  825.  
  826. DWORD CResourceDataEntry::GetCodePage() {
  827.   return m_dwCodePage;
  828. }