ImportDirectory.h

/*
* ImportDirectory.h - Part of the PeLib library.
*
* Copyright (c) 2004 Sebastian Porst (webmaster@the-interweb.com)
* All rights reserved.
*
* This software is licensed under the zlib/libpng License.
* For more details see http://www.opensource.org/licenses/zlib-license.php
* or the license information file (license.htm) in the root directory 
* of PeLib.
*/

#ifndef IMPORTDIRECTORY_H
#define IMPORTDIRECTORY_H

#include "PeLibAux.h"
#include "PeHeader.h"

namespace PeLib
{
        enum currdir {OLDDIR = 1, NEWDIR};
                
        class PeLibException;


        template<int bits>
        class ImportDirectory
        {
                typedef typename std::vector<PELIB_IMAGE_IMPORT_DIRECTORY<bits> >::iterator ImpDirFileIterator;
                typedef typename std::vector<PELIB_IMAGE_IMPORT_DIRECTORY<bits> >::const_iterator ConstImpDirFileIterator;
        
                private:
                  std::vector<PELIB_IMAGE_IMPORT_DIRECTORY<bits> > m_vOldiid;
                  std::vector<PELIB_IMAGE_IMPORT_DIRECTORY<bits> > m_vNewiid;

                // I can't convince Borland C++ to compile the function outside of the class declaration.
                // That's why the function definition is here.
                template<typename T> bool hasFunction(std::string strFilename, T value, bool(PELIB_THUNK_DATA<bits>::* comp)(T) const) const
            {
                        ConstImpDirFileIterator FileIter = m_vOldiid.begin();
                        ConstImpDirFileIterator EndIter = m_vOldiid.end();
                        
                        for (int i=0;i<=1;i++) // Loop once for m_vOldiid and once for m_vNewiid
                        {
                                do
                                {
                                        FileIter = std::find_if(FileIter, EndIter, std::bind2nd(std::mem_fun_ref(&PELIB_IMAGE_IMPORT_DIRECTORY<bits>::operator==), strFilename));
                                        
                                        if (FileIter != EndIter)
                                        {
                                                typename std::vector<PELIB_THUNK_DATA<bits> >::const_iterator Iter = std::find_if(FileIter->originalfirstthunk.begin(), FileIter->originalfirstthunk.end(), std::bind2nd(std::mem_fun_ref(comp), value));
                                                if (Iter != FileIter->originalfirstthunk.end())
                                                {
                                                        return true;
                                                }
                                                ++FileIter;
                                        }
                                }
                                while (FileIter != EndIter);
                                
                                FileIter = m_vNewiid.begin();
                                EndIter = m_vNewiid.end();
                        }
                        
                        return false;
                }
        

                public:

                  int addFunction(const std::string& strFilename, word wHint);
                  int addFunction(const std::string& strFilename, const std::string& strFuncname);

                  unsigned int getFileIndex(const std::string& strFilename, currdir cdDir) const;
                  unsigned int getFunctionIndex(const std::string& strFilename, const std::string& strFuncname, currdir cdDir) const;

                  std::string getFileName(dword dwFilenr, currdir cdDir) const;
                  
                  void setFileName(dword filenr, currdir dir, const std::string& name);
                  
                  word getFunctionHint(dword dwFilenr, dword dwFuncnr, currdir cdDir) const;
                  void setFunctionHint(dword dwFilenr, dword dwFuncnr, currdir cdDir, word value);
                  std::string getFunctionName(dword dwFilenr, dword dwFuncnr, currdir cdDir) const;
                  void setFunctionName(dword dwFilenr, dword dwFuncnr, currdir cdDir, const std::string& functionName);
                  dword getNumberOfFiles(currdir cdDir) const;
                  dword getNumberOfFunctions(dword dwFilenr, currdir cdDir) const;
                  int read(const std::string& strFilename, unsigned int uiOffset, unsigned int uiSize, const PeHeaderT<bits>& pehHeader);
                  void rebuild(std::vector<byte>& vBuffer, dword dwRva, bool fixEntries = true) const;
                  int removeFile(const std::string& strFilename);
                  int removeFunction(const std::string& strFilename, const std::string& strFuncname);
                  int removeFunction(const std::string& strFilename, word wHint);
                  unsigned int size() const;
                  int write(const std::string& strFilename, unsigned int uiOffset, unsigned int uiRva);

                  dword getFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir) const;
                  void setFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir, dword value);
                  dword getOriginalFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir) const;
                  void setOriginalFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir, dword value);

//                dword getFirstThunk(const std::string& strFilename, const std::string& strFuncname, currdir cdDir) const throw (PeLibException);
//                dword getOriginalFirstThunk(const std::string& strFilename, const std::string& strFuncname, currdir cdDir) const throw (PeLibException);

                  dword getFirstThunk(const std::string& strFilename, currdir cdDir) const;
                  dword getOriginalFirstThunk(const std::string& strFilename, currdir cdDir) const;
                  dword getForwarderChain(const std::string& strFilename, currdir cdDir) const;
                  dword getRvaOfName(const std::string& strFilename, currdir cdDir) const;
                  dword getTimeDateStamp(const std::string& strFilename, currdir cdDir) const;

                  dword getFirstThunk(dword dwFilenr, currdir cdDir) const;
                  void setFirstThunk(dword dwFilenr, currdir cdDir, dword value);
                  dword getOriginalFirstThunk(dword dwFilenr, currdir cdDir) const;
                  void setOriginalFirstThunk(dword dwFilenr, currdir cdDir, dword value);
                  dword getForwarderChain(dword dwFilenr, currdir cdDir) const;
                  void setForwarderChain(dword dwFilenr, currdir cdDir, dword value);
                  dword getRvaOfName(dword dwFilenr, currdir cdDir) const;
                  void setRvaOfName(dword dwFilenr, currdir cdDir, dword value);
                  dword getTimeDateStamp(dword dwFilenr, currdir cdDir) const;
                  void setTimeDateStamp(dword dwFilenr, currdir cdDir, dword value);

//                word getFunctionHint(const std::string& strFilename, const std::string& strFuncname, currdir cdDir) const throw (PeLibException);
        };

        template<int bits>
        int ImportDirectory<bits>::addFunction(const std::string& strFilename, word wHint)
        {
                if (hasFunction(strFilename, wHint, &PELIB_THUNK_DATA<bits>::equalHint))
                {
                        return 1;
//              throw Exceptions::EntryAlreadyExists(ImportDirectoryId, __LINE__);
                }
                
                // Find the imported file.
                ImpDirFileIterator FileIter = std::find_if(m_vNewiid.begin(), m_vNewiid.end(), std::bind2nd(std::mem_fun_ref(&PELIB_IMAGE_IMPORT_DIRECTORY<bits>::operator==), strFilename));

                PELIB_IMAGE_IMPORT_DIRECTORY<bits> iid;
                PELIB_THUNK_DATA<bits> td;
                td.hint = wHint;
                td.itd.Ordinal = wHint | IMAGE_ORDINAL_FLAGS<bits>::IMAGE_ORDINAL_FLAG;
                iid.name = strFilename;
                if (FileIter == m_vNewiid.end())
                {
                        iid.originalfirstthunk.push_back(td);
                        iid.firstthunk.push_back(td);
                        m_vNewiid.push_back(iid);
                }
                else
                {
                        FileIter->originalfirstthunk.push_back(td);
                        FileIter->firstthunk.push_back(td);
                }
                
                return 0;
        }

        template<int bits>
        int ImportDirectory<bits>::addFunction(const std::string& strFilename, const std::string& strFuncname)
        {
                if (hasFunction(strFilename, strFuncname, &PELIB_THUNK_DATA<bits>::equalFunctionName))
                {
                        return 1;
                        // throw Exceptions::EntryAlreadyExists(ImportDirectoryId, __LINE__);
                }
                
                // Find the imported file.
                ImpDirFileIterator FileIter = std::find_if(m_vNewiid.begin(), m_vNewiid.end(), std::bind2nd(std::mem_fun_ref(&PELIB_IMAGE_IMPORT_DIRECTORY<bits>::operator==), strFilename));

                PELIB_IMAGE_IMPORT_DIRECTORY<bits> iid;
                PELIB_THUNK_DATA<bits> td;
                td.fname = strFuncname;
                iid.name = strFilename;
                if (FileIter == m_vNewiid.end())
                {
                        iid.originalfirstthunk.push_back(td);
                        iid.firstthunk.push_back(td);
                        m_vNewiid.push_back(iid);
                }
                else
                {
                        FileIter->originalfirstthunk.push_back(td);
                        FileIter->firstthunk.push_back(td);
                }
                
                return 0;
        }

        template<int bits>
        unsigned int ImportDirectory<bits>::getFileIndex(const std::string& strFilename, currdir cdDir) const
        {
                const std::vector<PELIB_IMAGE_IMPORT_DIRECTORY<bits> >* currDir;
                
                if (cdDir == OLDDIR)
                {
                         currDir = &m_vOldiid;
                }
                else
                {
                         currDir = &m_vNewiid;
                }

                ConstImpDirFileIterator FileIter = std::find_if(currDir->begin(), currDir->end(), std::bind2nd(std::mem_fun_ref(&PELIB_IMAGE_IMPORT_DIRECTORY<bits>::operator==), strFilename));

                if (FileIter != currDir->end())
                {
                        return static_cast<unsigned int>(std::distance(currDir->begin(), FileIter));
                }
                else
                {
                        return -1;
                        // throw Exceptions::InvalidName(ImportDirectoryId, __LINE__);
                }

                return 0;
        }

        template<int bits>
        unsigned int ImportDirectory<bits>::getFunctionIndex(const std::string& strFilename, const std::string& strFuncname, currdir cdDir) const
        {
                unsigned int uiFile = getFileIndex(strFilename, cdDir);

                for (unsigned int i=0;i<getNumberOfFunctions(uiFile, cdDir);i++)
                {
                        if (getFunctionName(uiFile, i, cdDir) == strFuncname) return i;
                }

                return -1;
        }

        template<int bits>
        std::string ImportDirectory<bits>::getFileName(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR) return m_vOldiid[dwFilenr].name;
                else return m_vNewiid[dwFilenr].name;
        }

        template<int bits>
        void ImportDirectory<bits>::setFileName(dword filenr, currdir dir, const std::string& name)
        {
                if (dir == OLDDIR) m_vOldiid[filenr].name = name;
                else m_vNewiid[filenr].name = name;
        }

        template<int bits>
        std::string ImportDirectory<bits>::getFunctionName(dword dwFilenr, dword dwFuncnr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        // Unsafe
                        if (m_vOldiid[dwFilenr].impdesc.OriginalFirstThunk)
                        {
                                return m_vOldiid[dwFilenr].originalfirstthunk[dwFuncnr].fname;
                        }
                        else
                        {
                                return m_vOldiid[dwFilenr].firstthunk[dwFuncnr].fname;
                        }
                }
                else
                {
                        if (m_vNewiid[dwFilenr].impdesc.OriginalFirstThunk)
                        {
                                return m_vNewiid[dwFilenr].originalfirstthunk[dwFuncnr].fname;
                        }
                        else
                        {
                                return m_vNewiid[dwFilenr].firstthunk[dwFuncnr].fname;
                        }
                }
        }

        template<int bits>
        void ImportDirectory<bits>::setFunctionName(dword dwFilenr, dword dwFuncnr, currdir cdDir, const std::string& functionName)
        {
                if (cdDir == OLDDIR)
                {
                        // Unsafe
                        if (m_vOldiid[dwFilenr].impdesc.OriginalFirstThunk)
                        {
                                m_vOldiid[dwFilenr].originalfirstthunk[dwFuncnr].fname = functionName;
                        }
                        else
                        {
                                m_vOldiid[dwFilenr].firstthunk[dwFuncnr].fname = functionName;
                        }
                }
                else
                {
                        if (m_vNewiid[dwFilenr].impdesc.OriginalFirstThunk)
                        {
                                m_vNewiid[dwFilenr].originalfirstthunk[dwFuncnr].fname = functionName;
                        }
                        else
                        {
                                m_vNewiid[dwFilenr].firstthunk[dwFuncnr].fname = functionName;
                        }
                }
        }

        template<int bits>
        word ImportDirectory<bits>::getFunctionHint(dword dwFilenr, dword dwFuncnr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        if (m_vOldiid[dwFilenr].impdesc.OriginalFirstThunk)
                        {
                                return m_vOldiid[dwFilenr].originalfirstthunk[dwFuncnr].hint;
                        }
                        else
                        {
                                return m_vOldiid[dwFilenr].firstthunk[dwFuncnr].hint;
                        }
                }
                else return m_vNewiid[dwFilenr].originalfirstthunk[dwFuncnr].hint;
        }

        template<int bits>
        void ImportDirectory<bits>::setFunctionHint(dword dwFilenr, dword dwFuncnr, currdir cdDir, word value)
        {
                if (cdDir == OLDDIR)
                {
                        if (m_vOldiid[dwFilenr].impdesc.OriginalFirstThunk)
                        {
                                m_vOldiid[dwFilenr].originalfirstthunk[dwFuncnr].hint = value;
                        }
                        else
                        {
                                m_vOldiid[dwFilenr].firstthunk[dwFuncnr].hint = value;
                        }
                }
                else m_vNewiid[dwFilenr].originalfirstthunk[dwFuncnr].hint = value;
        }

        template<int bits>
        dword ImportDirectory<bits>::getNumberOfFiles(currdir cdDir) const
        {
                if (cdDir == OLDDIR) return static_cast<dword>(m_vOldiid.size());
                else return static_cast<dword>(m_vNewiid.size());
        }

        template<int bits>
        dword ImportDirectory<bits>::getNumberOfFunctions(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR) return static_cast<unsigned int>(m_vOldiid[dwFilenr].firstthunk.size());
                else return static_cast<unsigned int>(m_vNewiid[dwFilenr].firstthunk.size());
        }

        template<int bits>
        int ImportDirectory<bits>::read(const std::string& strFilename, unsigned int uiOffset, unsigned int uiSize, const PeHeaderT<bits>& pehHeader)
        {
                if (uiSize < PELIB_IMAGE_IMPORT_DESCRIPTOR::size())
                {
//                      throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                        return 1;
                }

                std::ifstream ifFile(strFilename.c_str(), std::ios_base::binary);
                unsigned long ulFileSize = fileSize(ifFile);

                if (!ifFile)
                {
//                      throw Exceptions::CannotOpenFile(ImportDirectoryId, __LINE__);
                        return 1;
                }
                
                if (ulFileSize < uiOffset + uiSize)
                {
//                      throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                        return 1;
                }

                ifFile.seekg(uiOffset, std::ios_base::beg);

                std::vector<unsigned char> vImportdirectory(uiSize);
                ifFile.read(reinterpret_cast<char*>(&vImportdirectory[0]), uiSize);

                PELIB_IMAGE_IMPORT_DIRECTORY<bits> iidCurr;
                unsigned int uiDesccounter = 0;

                InputBuffer inpBuffer(vImportdirectory);

                std::vector<PELIB_IMAGE_IMPORT_DIRECTORY<bits> > vOldIidCurr;

                do // Read and store all descriptors
                {
                        inpBuffer >> iidCurr.impdesc.OriginalFirstThunk;
                        inpBuffer >> iidCurr.impdesc.TimeDateStamp;
                        inpBuffer >> iidCurr.impdesc.ForwarderChain;
                        inpBuffer >> iidCurr.impdesc.Name;
                        inpBuffer >> iidCurr.impdesc.FirstThunk;

                        if (iidCurr.impdesc.OriginalFirstThunk != 0 || iidCurr.impdesc.TimeDateStamp != 0 || iidCurr.impdesc.ForwarderChain != 0 ||
                        iidCurr.impdesc.Name != 0 || iidCurr.impdesc.FirstThunk != 0)
                        {
                                vOldIidCurr.push_back(iidCurr);
                        }

                        uiDesccounter++;

                        if (uiSize < (uiDesccounter + 1) * PELIB_IMAGE_IMPORT_DESCRIPTOR::size()) break;
                } while (iidCurr.impdesc.OriginalFirstThunk != 0 || iidCurr.impdesc.TimeDateStamp != 0 || iidCurr.impdesc.ForwarderChain != 0 ||
                                iidCurr.impdesc.Name != 0 || iidCurr.impdesc.FirstThunk != 0);

                const unsigned int MAXNAMESIZE = 100;
                char namebuffer[MAXNAMESIZE + 1] = {0};
                        
                // Name
                for (unsigned int i=0;i<vOldIidCurr.size();i++)
                {
                        try
                        {
                                ifFile.seekg(static_cast<unsigned int>(pehHeader.rvaToOffset(vOldIidCurr[i].impdesc.Name)), std::ios_base::beg);
                        }
                        catch (...)
                        {
                                // Invalid offset
//                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                return 1;
                        }

                        std::string dllname = "";
                        
                        do
                        {
                                ifFile.read(namebuffer, MAXNAMESIZE);
                                dllname += namebuffer;
                        } while (strlen(namebuffer) == MAXNAMESIZE && !ifFile.fail());
                        
                        vOldIidCurr[i].name = dllname;
                        
                }

                // OriginalFirstThunk
                for (unsigned int i=0;i<vOldIidCurr.size();i++)
                {
                        PELIB_THUNK_DATA<bits> tdCurr;
                        dword uiVaoft = vOldIidCurr[i].impdesc.OriginalFirstThunk;

                        if (!uiVaoft)
                        {
                                continue;
                        }

                        try
                        {
                                ifFile.seekg(static_cast<unsigned int>(pehHeader.rvaToOffset(uiVaoft)), std::ios_base::beg);
                        }
                        catch (...)
                        {
                                // Invalid offset
//                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                return 1;
                        }

                        do
                        {
                                if (ulFileSize < pehHeader.rvaToOffset(uiVaoft) + sizeof(tdCurr.itd.Ordinal))
                                {
                                        return 1;
//                                      throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                }
                                uiVaoft += sizeof(tdCurr.itd.Ordinal);
                                
                                ifFile.read(reinterpret_cast<char*>(&tdCurr.itd.Ordinal), sizeof(tdCurr.itd.Ordinal));
                                if (tdCurr.itd.Ordinal) vOldIidCurr[i].originalfirstthunk.push_back(tdCurr);
                        } while (tdCurr.itd.Ordinal);
                }

                // FirstThunk
                for (unsigned int i=0;i<vOldIidCurr.size();i++)
                {
                        dword uiVaoft = vOldIidCurr[i].impdesc.FirstThunk;
                        PELIB_THUNK_DATA<bits> tdCurr;

                        try
                        {
                                ifFile.seekg(static_cast<unsigned int>(pehHeader.rvaToOffset(uiVaoft)), std::ios_base::beg);
                        }
                        catch (...)
                        {
                                // Invalid offset
//                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                return 1;
                        }

                        do
                        {
                                if (ulFileSize < pehHeader.rvaToOffset(uiVaoft) + sizeof(tdCurr.itd.Ordinal))
                                {
//                                      throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                        return 1;
                                }
                                uiVaoft += sizeof(tdCurr.itd.Ordinal);
                                
                                ifFile.read(reinterpret_cast<char*>(&tdCurr.itd.Ordinal), sizeof(tdCurr.itd.Ordinal));
                                if (tdCurr.itd.Ordinal) vOldIidCurr[i].firstthunk.push_back(tdCurr);
                        } while (tdCurr.itd.Ordinal);
                }

                // Names
                for (unsigned int i=0;i<vOldIidCurr.size();i++)
                {
                        if (vOldIidCurr[i].impdesc.OriginalFirstThunk)
                        {
                                for (unsigned int j=0;j<vOldIidCurr[i].originalfirstthunk.size();j++)
                                {
                                        if (vOldIidCurr[i].originalfirstthunk[j].itd.Ordinal & IMAGE_ORDINAL_FLAGS<bits>::IMAGE_ORDINAL_FLAG)
                                        {
                                                vOldIidCurr[i].originalfirstthunk[j].hint = 0;
                                                continue;
                                        }

                                        try
                                        {
                                                ifFile.seekg(static_cast<unsigned int>(pehHeader.rvaToOffset(vOldIidCurr[i].originalfirstthunk[j].itd.Ordinal)), std::ios_base::beg);
                                        }
                                        catch (...)
                                        {
                                                // Invalid offset
                                                //throw PeLibException(PeLibException::INVALID_IMPORT_DIRECTORY, __FILE__, __LINE__);
//                                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                                return 1;
                                         }

                                        try
                                        {
                                                ifFile.read(reinterpret_cast<char*>(&vOldIidCurr[i].originalfirstthunk[j].hint), sizeof(vOldIidCurr[i].originalfirstthunk[j].hint));
                                        }
                                        catch(...)
                                        {
//                                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                                return 1;
                                        }

                                        std::string funcname = "";
                                        do
                                        {
                                                ifFile.read(namebuffer, MAXNAMESIZE);
                                                funcname += namebuffer;
                                        } while (strlen(namebuffer) == MAXNAMESIZE && !ifFile.fail());
                        
                                        vOldIidCurr[i].originalfirstthunk[j].fname = funcname;
                                }
                        }
                        else
                        {
                                for (unsigned int j=0;j<vOldIidCurr[i].firstthunk.size();j++)
                                {
                                        if (vOldIidCurr[i].firstthunk[j].itd.Ordinal & IMAGE_ORDINAL_FLAGS<bits>::IMAGE_ORDINAL_FLAG)
                                        {
                                                continue;
                                        }

                                        try
                                        {
                                                ifFile.seekg(static_cast<unsigned int>(pehHeader.rvaToOffset(vOldIidCurr[i].firstthunk[j].itd.Ordinal)), std::ios_base::beg);
                                        }
                                        catch (...)
                                        {
                                                // Invalid offset
//                                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                                return 1;
                                        }

                                        try
                                        {
                                                ifFile.read(reinterpret_cast<char*>(&vOldIidCurr[i].firstthunk[j].hint), sizeof(vOldIidCurr[i].firstthunk[j].hint));
                                        }
                                        catch(...)
                                        {
//                                              throw Exceptions::InvalidFormat(ImportDirectoryId, __LINE__);
                                                return 1;
                                        }

                                        std::string funcname = "";
                                        do
                                        {
                                                ifFile.read(namebuffer, MAXNAMESIZE);
                                                funcname += namebuffer;
                                        } while (strlen(namebuffer) == MAXNAMESIZE && !ifFile.fail());
                        
                                        vOldIidCurr[i].firstthunk[j].fname = funcname;
                                }
                        }
                }
                std::swap(vOldIidCurr, m_vOldiid);
                return 0;
        }

        template<int bits>
        void ImportDirectory<bits>::rebuild(std::vector<byte>& vBuffer, dword dwRva, bool fixEntries) const
        {
                unsigned int uiImprva = dwRva;
                unsigned int uiSizeofdescriptors = (static_cast<unsigned int>(m_vNewiid.size() + m_vOldiid.size()) + 1) * PELIB_IMAGE_IMPORT_DESCRIPTOR::size();
                
                unsigned int uiSizeofdllnames = 0, uiSizeoffuncnames = 0;
                unsigned int uiSizeofoft = 0;

                for (unsigned int i=0;i<m_vNewiid.size();i++)
                {
                        uiSizeofdllnames += static_cast<unsigned int>(m_vNewiid[i].name.size()) + 1;
                        uiSizeofoft += (static_cast<unsigned int>(m_vNewiid[i].originalfirstthunk.size())+1) * PELIB_IMAGE_THUNK_DATA<bits>::size();

                        for(unsigned int j=0;j<m_vNewiid[i].originalfirstthunk.size();j++)
                        {
                                // +3 for hint (word) and 00-byte
                                uiSizeoffuncnames += (static_cast<unsigned int>(m_vNewiid[i].originalfirstthunk[j].fname.size()) + 3);
                        }
                }

//              for (unsigned int i=0;i<m_vNewiid.size();i++)
//              {
//                      uiSizeofoft += (static_cast<unsigned int>(m_vNewiid[i].originalfirstthunk.size())+1) * PELIB_IMAGE_THUNK_DATA<bits>::size();
//              }
                
                OutputBuffer obBuffer(vBuffer);

                // Rebuild IMAGE_IMPORT_DESCRIPTORS
                for (unsigned int i=0;i<m_vOldiid.size();i++)
                {
                        obBuffer << m_vOldiid[i].impdesc.OriginalFirstThunk;
                        obBuffer << m_vOldiid[i].impdesc.TimeDateStamp;
                        obBuffer << m_vOldiid[i].impdesc.ForwarderChain;
                        obBuffer << m_vOldiid[i].impdesc.Name;
                        obBuffer << m_vOldiid[i].impdesc.FirstThunk;
                }

                unsigned int dllsize = 0;

                for (unsigned int i=0;i<m_vNewiid.size();i++)
                {
                        dword dwPoft = uiSizeofdescriptors + uiImprva;

                        for (unsigned int j=1;j<=i;j++)
                        {
                                dwPoft += (static_cast<unsigned int>(m_vNewiid[j-1].originalfirstthunk.size()) + 1) * PELIB_IMAGE_THUNK_DATA<bits>::size();
                        }

                        obBuffer << (fixEntries ? dwPoft : m_vNewiid[i].impdesc.OriginalFirstThunk);
                        obBuffer << m_vNewiid[i].impdesc.TimeDateStamp;
                        obBuffer << m_vNewiid[i].impdesc.ForwarderChain;
                        dword dwPdll = uiSizeofdescriptors + uiSizeofoft + uiImprva + dllsize;
                        obBuffer << (fixEntries ? dwPdll : m_vNewiid[i].impdesc.Name);
                        obBuffer << (fixEntries ? dwPoft : m_vNewiid[i].impdesc.FirstThunk);

                        dllsize += static_cast<unsigned int>(m_vNewiid[i].name.size()) + 1;
                }

                obBuffer << (dword)0;
                obBuffer << (dword)0;
                obBuffer << (dword)0;
                obBuffer << (dword)0;
                obBuffer << (dword)0;

                unsigned int uiPfunc = uiSizeofdescriptors + uiSizeofoft + uiSizeofdllnames + uiImprva;

                // Rebuild original first thunk
                for (unsigned int i=0;i<m_vNewiid.size();i++)
                {
                        for (unsigned int j=0;j<m_vNewiid[i].originalfirstthunk.size();j++)
                        {
                                if (m_vNewiid[i].originalfirstthunk[j].itd.Ordinal & IMAGE_ORDINAL_FLAGS<bits>::IMAGE_ORDINAL_FLAG
                                        || fixEntries == false)
                                {
                                        obBuffer << m_vNewiid[i].originalfirstthunk[j].itd.Ordinal;
                                }
                                else
                                {
                                        obBuffer << uiPfunc;
                                }
                                uiPfunc += static_cast<unsigned int>(m_vNewiid[i].originalfirstthunk[j].fname.size()) + 3;
                        }
                        obBuffer << (dword)0;
                }

                // Write dllnames into import directory
                for (unsigned int i=0;i<m_vNewiid.size();i++)
                {
                        obBuffer.add(m_vNewiid[i].name.c_str(), static_cast<unsigned int>(m_vNewiid[i].name.size())+1);
                }

                // Write function names into directory
                for (unsigned int i=0;i<m_vNewiid.size();i++)
                {
                        for (unsigned int j=0;j<m_vNewiid[i].originalfirstthunk.size();j++)
                        {
                                obBuffer << m_vNewiid[i].originalfirstthunk[j].hint;
                                obBuffer.add(m_vNewiid[i].originalfirstthunk[j].fname.c_str(), static_cast<unsigned int>(m_vNewiid[i].originalfirstthunk[j].fname.size()) + 1);
                        }
                }
        }

        template<int bits>
        int ImportDirectory<bits>::removeFile(const std::string& strFilename)
        {
                unsigned int oldSize = static_cast<unsigned int>(m_vNewiid.size());

                m_vNewiid.erase(std::remove_if(m_vNewiid.begin(), m_vNewiid.end(), std::bind2nd(std::mem_fun_ref(&PELIB_IMAGE_IMPORT_DIRECTORY<bits>::operator==), strFilename)), m_vNewiid.end());
                
                return oldSize == m_vNewiid.size() ? 1 : 0;
        }

        template<int bits>
        int ImportDirectory<bits>::removeFunction(const std::string& strFilename, const std::string& strFuncname)
        {
                ImpDirFileIterator viPos = m_vNewiid.begin();

                int notFound = 1;

                while (viPos != m_vNewiid.end())
                {
                        if (isEqualNc(viPos->name, strFilename))
                        {
                                unsigned int oldSize = static_cast<unsigned int>(viPos->originalfirstthunk.size());
                                viPos->originalfirstthunk.erase(std::remove_if(viPos->originalfirstthunk.begin(), viPos->originalfirstthunk.end(), std::bind2nd(std::mem_fun_ref(&PELIB_THUNK_DATA<bits>::equalFunctionName), strFuncname)), viPos->originalfirstthunk.end());
                                //viPos->originalfirstthunk.erase(std::remove_if(viPos->originalfirstthunk.begin(), viPos->originalfirstthunk.end(), std::bind2nd(CompPolicy<PELIB_THUNK_DATA, std::string>(), strFuncname)));
                                if (viPos->originalfirstthunk.size() != oldSize) notFound = 0;
                        }
                        ++viPos;
                }
                
                return notFound;
        }

        template<int bits>
        int ImportDirectory<bits>::removeFunction(const std::string& strFilename, word wHint)
        {
                ImpDirFileIterator viPos = m_vNewiid.begin();
                int notFound = 1;

                while (viPos != m_vNewiid.end())
                {
                        if (isEqualNc(viPos->name, strFilename))
                        {
                                unsigned int oldSize = static_cast<unsigned int>(viPos->originalfirstthunk.size());
                                viPos->originalfirstthunk.erase(std::remove_if(viPos->originalfirstthunk.begin(), viPos->originalfirstthunk.end(), std::bind2nd(std::mem_fun_ref(&PELIB_THUNK_DATA<bits>::equalHint), wHint)), viPos->originalfirstthunk.end());
                                unsigned int newPos = static_cast<unsigned int>(viPos->originalfirstthunk.size());
                                if (viPos->originalfirstthunk.size() != oldSize) notFound = 0;
                        }
                        ++viPos;
                }
                
                return notFound;
        }

        template<int bits>
        int ImportDirectory<bits>::write(const std::string& strFilename, unsigned int uiOffset, unsigned int uiRva)
        {
                std::fstream ofFile(strFilename.c_str(), std::ios_base::in | std::ios_base::out | std::ios_base::binary);

                if (!ofFile)
                {
                        ofFile.clear();
                        return 1;
                }

                ofFile.seekp(uiOffset, std::ios_base::beg);

                std::vector<byte> vBuffer;

                rebuild(vBuffer, uiRva);

                ofFile.write(reinterpret_cast<const char*>(&vBuffer[0]), vBuffer.size());
                ofFile.close();

                std::copy(m_vNewiid.begin(), m_vNewiid.end(), std::back_inserter(m_vOldiid));
                m_vNewiid.clear();
                
                return 0;
        }

        template<int bits>
        unsigned int ImportDirectory<bits>::size() const
        {
                // Only the descriptors of m_vOldiid must be rebuilt, not the data they point to.
                return std::accumulate(m_vNewiid.begin(), m_vNewiid.end(), 0, accumulate<PELIB_IMAGE_IMPORT_DIRECTORY<bits> >)
                + (m_vOldiid.size() + 1) * PELIB_IMAGE_IMPORT_DESCRIPTOR::size();
        }

        template<int bits>
        dword ImportDirectory<bits>::getFirstThunk(const std::string& strFilename, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[getFileIndex(strFilename, cdDir)].impdesc.FirstThunk;
                }
                else
                {
                        return m_vNewiid[getFileIndex(strFilename, cdDir)].impdesc.FirstThunk;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getOriginalFirstThunk(const std::string& strFilename, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[getFileIndex(strFilename, cdDir)].impdesc.OriginalFirstThunk;
                }
                else
                {
                        return m_vNewiid[getFileIndex(strFilename, cdDir)].impdesc.OriginalFirstThunk;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getForwarderChain(const std::string& strFilename, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[getFileIndex(strFilename, cdDir)].impdesc.ForwarderChain;
                }
                else
                {
                        return m_vNewiid[getFileIndex(strFilename, cdDir)].impdesc.ForwarderChain;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getTimeDateStamp(const std::string& strFilename, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[getFileIndex(strFilename, cdDir)].impdesc.TimeDateStamp;
                }
                else
                {
                        return m_vNewiid[getFileIndex(strFilename, cdDir)].impdesc.TimeDateStamp;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getRvaOfName(const std::string& strFilename, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[getFileIndex(strFilename, cdDir)].impdesc.Name;
                }
                else
                {
                        return m_vNewiid[getFileIndex(strFilename, cdDir)].impdesc.Name;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getFirstThunk(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[dwFilenr].impdesc.FirstThunk;
                }
                else
                {
                        return m_vNewiid[dwFilenr].impdesc.FirstThunk;
                }
        }

        template<int bits>
        void ImportDirectory<bits>::setFirstThunk(dword dwFilenr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR)
                {
                        m_vOldiid[dwFilenr].impdesc.FirstThunk = value;
                }
                else
                {
                        m_vNewiid[dwFilenr].impdesc.FirstThunk = value;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getOriginalFirstThunk(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[dwFilenr].impdesc.OriginalFirstThunk;
                }
                else
                {
                        return m_vNewiid[dwFilenr].impdesc.OriginalFirstThunk;
                }
        }

        template<int bits>
        void ImportDirectory<bits>::setOriginalFirstThunk(dword dwFilenr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR)
                {
                        m_vOldiid[dwFilenr].impdesc.OriginalFirstThunk = value;
                }
                else
                {
                        m_vNewiid[dwFilenr].impdesc.OriginalFirstThunk = value;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getForwarderChain(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[dwFilenr].impdesc.ForwarderChain;
                }
                else
                {
                        return m_vNewiid[dwFilenr].impdesc.ForwarderChain;
                }
        }

        template<int bits>
        void ImportDirectory<bits>::setForwarderChain(dword dwFilenr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR)
                {
                        m_vOldiid[dwFilenr].impdesc.ForwarderChain = value;
                }
                else
                {
                        m_vNewiid[dwFilenr].impdesc.ForwarderChain = value;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getTimeDateStamp(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[dwFilenr].impdesc.TimeDateStamp;
                }
                else
                {
                        return m_vNewiid[dwFilenr].impdesc.TimeDateStamp;
                }
        }

        template<int bits>
        void ImportDirectory<bits>::setTimeDateStamp(dword dwFilenr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR)
                {
                        m_vOldiid[dwFilenr].impdesc.TimeDateStamp = value;
                }
                else
                {
                        m_vNewiid[dwFilenr].impdesc.TimeDateStamp = value;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getRvaOfName(dword dwFilenr, currdir cdDir) const
        {
                if (cdDir == OLDDIR)
                {
                        return m_vOldiid[dwFilenr].impdesc.Name;
                }
                else
                {
                        return m_vNewiid[dwFilenr].impdesc.Name;
                }
        }

        template<int bits>
        void ImportDirectory<bits>::setRvaOfName(dword dwFilenr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR)
                {
                        m_vOldiid[dwFilenr].impdesc.Name = value;
                }
                else
                {
                        m_vNewiid[dwFilenr].impdesc.Name = value;
                }
        }

        template<int bits>
        dword ImportDirectory<bits>::getFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir) const
        {
                if (cdDir == OLDDIR) return m_vOldiid[dwFilenr].firstthunk[dwFuncnr].itd.Ordinal;
                else return m_vNewiid[dwFilenr].firstthunk[dwFuncnr].itd.Ordinal;
        }
        
        template<int bits>
        void ImportDirectory<bits>::setFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR) m_vOldiid[dwFilenr].firstthunk[dwFuncnr].itd.Ordinal = value;
                else m_vNewiid[dwFilenr].firstthunk[dwFuncnr].itd.Ordinal = value;
        }
        
        template<int bits>
        dword ImportDirectory<bits>::getOriginalFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir) const
        {
                if (cdDir == OLDDIR) return m_vOldiid[dwFilenr].originalfirstthunk[dwFuncnr].itd.Ordinal;
                else return m_vNewiid[dwFilenr].originalfirstthunk[dwFuncnr].itd.Ordinal;
        }

        template<int bits>
        void ImportDirectory<bits>::setOriginalFirstThunk(dword dwFilenr, dword dwFuncnr, currdir cdDir, dword value)
        {
                if (cdDir == OLDDIR) m_vOldiid[dwFilenr].originalfirstthunk[dwFuncnr].itd.Ordinal = value;
                else m_vNewiid[dwFilenr].originalfirstthunk[dwFuncnr].itd.Ordinal = value;
        }

        typedef ImportDirectory<32> ImportDirectory32;
        typedef ImportDirectory<64> ImportDirectory64;
}

#endif

---