#include "PEOptions.h"

#include "StrLib.h"

#if !defined(PE_NO_PEFILEEX)
#include "PEFileEx.h"
#endif // PE_NO_PEFILEEX

#include "PEImpExp.h"

#if !defined(PE_NO_LOGFILE)
#include "LogFile.h"
#endif // PE_NO_LOGFILE

////////////
// Import //
////////////
DWORD dwInternalGetImportedFunctionEntry(DWORD dwImageBase,PIMAGE_IMPORT_MODULE_DIRECTORY pid,LPCSTR lpProcName)
{
	char *lpBuffer=(char*) dwImageBase;
  if ( (lpBuffer) && (pid) )
  {
		// walk through import table name list (DLL)
		while(pid->dwRVAFunctionNameList)
		{
			char *DllNamePtr=(char*) &lpBuffer[pid->dwRVAModuleName];
			DWORD dwRVAFunctionNameList=pid->dwRVAFunctionNameList;
			DWORD dwRVAFunctionAddressList=pid->dwRVAFunctionAddressList;
			DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[dwRVAFunctionNameList];
			DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[dwRVAFunctionAddressList];
			
			// walk through DLL func
			while(*FunctionNamePtr)
			{
				// ordinal or not ?
        if (!IMAGE_SNAP_BY_ORDINAL(*FunctionNamePtr))
				{
          PIMAGE_IMPORT_BY_NAME ImportPtr=(PIMAGE_IMPORT_BY_NAME) &lpBuffer[*FunctionNamePtr];
          WORD Hint=ImportPtr->Hint;
          char *NamePtr=(char*) &ImportPtr->Name[0];

          //if (!StringCompare(NamePtr,lpProcName))
		  if (!StringCompareToCryptKey(NamePtr,lpProcName))
					{
						// found what we want
						return((DWORD) FunctionAdressPtr);
					}
				}
				FunctionNamePtr++;
				FunctionAdressPtr++;
			}
			pid++;
		}
	}
	return(0);
}

BOOL bInternalGetImportedFunctionName(DWORD dwImageBase,PIMAGE_IMPORT_MODULE_DIRECTORY pid,DWORD dwAddr,char **szDllName,char **szFuncName)
{
	char *lpBuffer=(char*) dwImageBase;
  
  *szDllName=NULL;
  *szFuncName=NULL;

  if ( (lpBuffer) && (pid) )
  {
		// walk through import table name list (DLL)
		while(pid->dwRVAFunctionNameList)
		{
			char *DllNamePtr=(char*) &lpBuffer[pid->dwRVAModuleName];
			DWORD dwRVAFunctionNameList=pid->dwRVAFunctionNameList;
			DWORD dwRVAFunctionAddressList=pid->dwRVAFunctionAddressList;
			DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[dwRVAFunctionNameList];
			DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[dwRVAFunctionAddressList];
			
			// walk through DLL func
			while(*FunctionNamePtr)
			{
        if (((DWORD) *FunctionAdressPtr)==dwAddr)
        {
          if (!IMAGE_SNAP_BY_ORDINAL(*FunctionNamePtr))
	  			{
            PIMAGE_IMPORT_BY_NAME ImportPtr=(PIMAGE_IMPORT_BY_NAME) &lpBuffer[*FunctionNamePtr];
            WORD Hint=ImportPtr->Hint;
            char *FuncNamePtr=(char*) &ImportPtr->Name[0];

            // found what we want
            *szDllName=DllNamePtr;
            *szFuncName=FuncNamePtr;
            return(TRUE);
          }
          else
          {
            *szDllName=DllNamePtr;
            *szFuncName=NULL;
            // don't handle ordinal
            return(TRUE);
          }
        }
				FunctionNamePtr++;
				FunctionAdressPtr++;
			}
			pid++;
		}
	}
	return(FALSE);
}


DWORD dwGetImportedFunctionEntryFromImageBase(DWORD dwImageBase,LPCSTR lpProcName)
{
	char *lpBuffer=(char*) dwImageBase;
	PIMAGE_OPTIONAL_HEADER pPEOptionalHeader=(PIMAGE_OPTIONAL_HEADER) OPTHDROFFSET(lpBuffer); 
	DWORD ImportTableRVA=pPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
	if (ImportTableRVA)
	{
		PIMAGE_IMPORT_MODULE_DIRECTORY pid=(PIMAGE_IMPORT_MODULE_DIRECTORY) &lpBuffer[ImportTableRVA];
    return(dwInternalGetImportedFunctionEntry(dwImageBase,pid,lpProcName));
  }
  return(0);
}

BOOL bGetImportedFunctionNameFromImageBase(DWORD dwImageBase,DWORD dwAddr,char **szDllName,char **szFuncName)
{
	char *lpBuffer=(char*) dwImageBase;
	PIMAGE_OPTIONAL_HEADER pPEOptionalHeader=(PIMAGE_OPTIONAL_HEADER) OPTHDROFFSET(lpBuffer); 
	DWORD ImportTableRVA=pPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
	if (ImportTableRVA)
	{
		PIMAGE_IMPORT_MODULE_DIRECTORY pid=(PIMAGE_IMPORT_MODULE_DIRECTORY) &lpBuffer[ImportTableRVA];
    return(bInternalGetImportedFunctionName(dwImageBase,pid,dwAddr,szDllName,szFuncName));
  }
  return(FALSE);
}

#if !defined(PE_NO_PEFILEEX)
DWORD dwGetImportedFunctionEntryFromPEFile(LPPEFILE *lpPEFile,LPCSTR lpProcName)
{
	DWORD ImportTableRVA=(*lpPEFile)->PEFileHeader.lpPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
  if (ImportTableRVA)
  {
    LPPESECTION lpImportSection=lpFindSectionWithRVA(lpPEFile,ImportTableRVA);
    if (lpImportSection)
    {
      DWORD dwOffset=(DWORD) ImportTableRVA - lpImportSection->lpSectionHeader->VirtualAddress;
      PIMAGE_IMPORT_MODULE_DIRECTORY pid=(PIMAGE_IMPORT_MODULE_DIRECTORY) &lpImportSection->lpBuffer[dwOffset];    
      DWORD dwImageBase=((DWORD) lpImportSection->lpBuffer) - lpImportSection->lpSectionHeader->VirtualAddress;
      return(dwInternalGetImportedFunctionEntry(dwImageBase,pid,lpProcName));
    }
  }
  return(0);
}
#endif // !PE_NO_PEFILEEX

#if !defined(PE_NO_PEFILEEX)
BOOL bGetImportedFunctionNameFromPEFile(LPPEFILE *lpPEFile,DWORD dwAddr,char **szDllName,char **szFuncName)
{
	DWORD ImportTableRVA=(*lpPEFile)->PEFileHeader.lpPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
  if (ImportTableRVA)
  {
    LPPESECTION lpImportSection=lpFindSectionWithRVA(lpPEFile,ImportTableRVA);
    if (lpImportSection)
    {
      DWORD dwOffset=(DWORD) ImportTableRVA - lpImportSection->lpSectionHeader->VirtualAddress;
      PIMAGE_IMPORT_MODULE_DIRECTORY pid=(PIMAGE_IMPORT_MODULE_DIRECTORY) &lpImportSection->lpBuffer[dwOffset];    
      DWORD dwImageBase=((DWORD) lpImportSection->lpBuffer) - lpImportSection->lpSectionHeader->VirtualAddress;
      return(bInternalGetImportedFunctionName(dwImageBase,pid,dwAddr,szDllName,szFuncName));
    }
  }
  return(FALSE);
}
#endif // !PE_NO_PEFILEEX

#if !defined(PE_NO_LOGFILE) && !defined(PE_NO_PEFILEEX)
void LogImportedFunctionFromPEFile(LPPEFILE *lpPEFile)
{
	DWORD ImportTableRVA=(*lpPEFile)->PEFileHeader.lpPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
    
  WriteLogFile("\nImport Table : RVA=%08X",ImportTableRVA);  

  if (ImportTableRVA)
  {
    LPPESECTION lpImportSection=lpFindSectionWithRVA(lpPEFile,ImportTableRVA);
    if (lpImportSection)
    {
      DWORD dwOffset=(DWORD) ImportTableRVA - lpImportSection->lpSectionHeader->VirtualAddress;
      PIMAGE_IMPORT_MODULE_DIRECTORY pid=(PIMAGE_IMPORT_MODULE_DIRECTORY) &lpImportSection->lpBuffer[dwOffset];    
      DWORD dwImageBase=((DWORD) lpImportSection->lpBuffer) - lpImportSection->lpSectionHeader->VirtualAddress;
      DWORD dwRealImageBase=(*lpPEFile)->PEFileHeader.lpPEOptionalHeader->ImageBase;

	    char *lpBuffer=(char*) dwImageBase;
      if ( (lpBuffer) && (pid) )
      {
		    // walk through import table name list (DLL)
		    while(pid->dwRVAFunctionNameList)
		    {
			    char *DllNamePtr=(char*) &lpBuffer[pid->dwRVAModuleName];
			    DWORD dwRVAFunctionNameList=pid->dwRVAFunctionNameList;
			    DWORD dwRVAFunctionAddressList=pid->dwRVAFunctionAddressList;
			    DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[dwRVAFunctionNameList];
			    DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[dwRVAFunctionAddressList];
          WORD Hint=0;
          char DllName[MAX_PATH];

          StringCopy(DllName,DllNamePtr);
          StringUpper(DllName);
          WriteLogFile("  Imported functions from %s",DllName);
          *strrchr(DllName,'.')=0;
			
			    // walk through DLL func
			    while(*FunctionNamePtr)
			    {
   			    // ordinal or not ?
            if (!IMAGE_SNAP_BY_ORDINAL(*FunctionNamePtr))
	  			  {
              PIMAGE_IMPORT_BY_NAME ImportPtr=(PIMAGE_IMPORT_BY_NAME) &lpBuffer[*FunctionNamePtr];
              char *FuncNamePtr=(char*) &ImportPtr->Name[0];            
              Hint=ImportPtr->Hint;

              WriteLogFile("    Addr:(%08X) hint(%04X) Name:%s",*FunctionAdressPtr,Hint,FuncNamePtr);
            }
            else
            {
              Hint=(WORD) IMAGE_ORDINAL(*FunctionNamePtr);
              WriteLogFile("    Addr:(%08X) hint(%04X) Name:???",*FunctionAdressPtr,Hint);
            }
				    FunctionNamePtr++;
				    FunctionAdressPtr++;
			    }
			    pid++;
		    }
      }
    }
  }
}
#endif // !PE_NO_LOGFILE




////////////
// export //
////////////
/*
DWORD dwGetInternalExportedFunctionEntry(DWORD dwImageBase,PIMAGE_EXPORT_DIRECTORY ped,LPCSTR lpProcName)
{
	char *lpBuffer=(char*) dwImageBase;
  if ( (lpBuffer) && (ped) )
  {
    DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNames];
    DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfFunctions];
    DWORD *FunctionOrdinalPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNameOrdinals];
    int i;
    for (i=0; i<(int)ped->NumberOfNames; i++)
    {
      char *NamePtr=(char*) &lpBuffer[*FunctionNamePtr];
      if (!StringCompare(NamePtr,lpProcName))
      {
        return( (DWORD) FunctionAdressPtr);
      }
      FunctionNamePtr++;
      FunctionAdressPtr++;
      FunctionOrdinalPtr++;
    }
	}
	return(0);
}
*/

DWORD dwGetInternalExportedFunctionEntry(DWORD dwImageBase,PIMAGE_EXPORT_DIRECTORY ped,LPCSTR lpProcName)
{
	char *lpBuffer=(char*) dwImageBase;
  if ( (lpBuffer) && (ped) )
  {
    DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNames];
    DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfFunctions];
    WORD *FunctionOrdinalPtr=(WORD*) &lpBuffer[(DWORD)ped->AddressOfNameOrdinals];
    int i;
    for (i=0; i<(int)ped->NumberOfNames; i++)
    {
      char *NamePtr=(char*) &lpBuffer[*FunctionNamePtr];
			DWORD AddrPtr=*(FunctionAdressPtr + *FunctionOrdinalPtr);
      //if (!StringCompare(NamePtr,lpProcName))
	  if (!StringCompareToCryptKey(NamePtr,lpProcName))
      {
        return(AddrPtr);
      }
      FunctionNamePtr++;
      FunctionOrdinalPtr++;
    }
	}
	return(0);
}


DWORD dwGetExportedFunctionEntryFromImageBase(DWORD dwImageBase,LPCSTR lpProcName)
{
	char *lpBuffer=(char*) dwImageBase;
	PIMAGE_OPTIONAL_HEADER pPEOptionalHeader=(PIMAGE_OPTIONAL_HEADER) OPTHDROFFSET(lpBuffer); 
	DWORD ExportTableRVA=pPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;
	if (ExportTableRVA)
	{
		PIMAGE_EXPORT_DIRECTORY ped=(PIMAGE_EXPORT_DIRECTORY) &lpBuffer[ExportTableRVA];
    return(dwGetInternalExportedFunctionEntry(dwImageBase,ped,lpProcName));
	}
	return(0);
}

#if !defined(PE_NO_PEFILEEX)
DWORD dwGetExportedFunctionEntryFromPEFile(LPPEFILE *lpPEFile,LPCSTR lpProcName)
{
	DWORD ExportTableRVA=(*lpPEFile)->PEFileHeader.lpPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;
  if (ExportTableRVA)
  {
    LPPESECTION lpExportSection=lpFindSectionWithRVA(lpPEFile,ExportTableRVA);
    if (lpExportSection)
    {
      DWORD dwOffset=(DWORD) ExportTableRVA - lpExportSection->lpSectionHeader->VirtualAddress;
      PIMAGE_EXPORT_DIRECTORY ped=(PIMAGE_EXPORT_DIRECTORY) &lpExportSection->lpBuffer[dwOffset];    
      DWORD dwImageBase=((DWORD) lpExportSection->lpBuffer) - lpExportSection->lpSectionHeader->VirtualAddress;     
      return(dwGetInternalExportedFunctionEntry(dwImageBase,ped,lpProcName));
    }
  }
  return(0);
}
#endif // !PE_NO_PEFILEEX

BOOL bFindExportedFunctionNameWithOrdinalFromImageBase(DWORD dwImageBase,WORD uwOrdinal,char *szFuncName)
{
  BOOL bRet=FALSE;
	char *lpBuffer=(char*) dwImageBase;
	PIMAGE_OPTIONAL_HEADER pPEOptionalHeader=(PIMAGE_OPTIONAL_HEADER) OPTHDROFFSET(lpBuffer); 
	DWORD ExportTableRVA=pPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;
	if (ExportTableRVA)
	{
		PIMAGE_EXPORT_DIRECTORY ped=(PIMAGE_EXPORT_DIRECTORY) &lpBuffer[ExportTableRVA];
    if (ped)
    {
      DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNames];
      DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfFunctions];
      WORD *FunctionOrdinalPtr=(WORD*) &lpBuffer[(DWORD)ped->AddressOfNameOrdinals];
      int i;       

      uwOrdinal=(WORD) (uwOrdinal-ped->Base);

      for (i=0;i<(int)ped->NumberOfNames; i++)
      {
        char *NamePtr=(char*) &lpBuffer[*FunctionNamePtr];
   		  DWORD AddrPtr=*(FunctionAdressPtr + *FunctionOrdinalPtr);
        WORD Hint=*FunctionOrdinalPtr;

        if (Hint==uwOrdinal)
        {
          StringCopy(szFuncName,NamePtr);
          bRet=TRUE;
          break;
        }       

        FunctionNamePtr++;
        FunctionOrdinalPtr++;
      }     
    }
  }
  bRet=TRUE;
  return(bRet);
}

#if !defined(PE_NO_LOGFILE) && !defined(PE_NO_PEFILEEX)
void LogExportedFunctionFromPEFile(LPPEFILE *lpPEFile)
{
	DWORD ExportTableRVA=(*lpPEFile)->PEFileHeader.lpPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;

  WriteLogFile("\nExport Table : RVA=%08X",ExportTableRVA);  
  
  if (ExportTableRVA)
  {
    LPPESECTION lpExportSection=lpFindSectionWithRVA(lpPEFile,ExportTableRVA);
    if (lpExportSection)
    {
      DWORD dwOffset=(DWORD) ExportTableRVA - lpExportSection->lpSectionHeader->VirtualAddress;
      PIMAGE_EXPORT_DIRECTORY ped=(PIMAGE_EXPORT_DIRECTORY) &lpExportSection->lpBuffer[dwOffset];    
      DWORD dwImageBase=((DWORD) lpExportSection->lpBuffer) - lpExportSection->lpSectionHeader->VirtualAddress;     

    	char *lpBuffer=(char*) dwImageBase;
      if ( (lpBuffer) && (ped) )
      {

        DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNames];
        DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfFunctions];
        WORD *FunctionOrdinalPtr=(WORD*) &lpBuffer[(DWORD)ped->AddressOfNameOrdinals];
        int i;       
        for (i=0; i<(int)ped->NumberOfNames; i++)
        {
           char *NamePtr=(char*) &lpBuffer[*FunctionNamePtr];
       		 DWORD AddrPtr=*(FunctionAdressPtr + *FunctionOrdinalPtr);
           WORD Hint=*FunctionOrdinalPtr;

           WriteLogFile("    Addr:(%08X) hint(%04X) ordinal(%04X) Name:%s",AddrPtr,Hint,Hint+ped->Base,NamePtr);

           FunctionNamePtr++;
           FunctionOrdinalPtr++;
        }            
     
      }
    }
	}
}
#endif // !PE_NO_LOGFILE  && !PE_NO_PEFILEEX

/*
DWORD ASMdwGetExportedFunctionEntryFromImageBase(DWORD dwImageBase,LPCSTR lpProcName)
{
	char *lpBuffer=(char*) dwImageBase;
	PIMAGE_OPTIONAL_HEADER pPEOptionalHeader=(PIMAGE_OPTIONAL_HEADER) OPTHDROFFSET(lpBuffer); 
	DWORD ExportTableRVA=pPEOptionalHeader->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;
	if (ExportTableRVA)
	{
		PIMAGE_EXPORT_DIRECTORY ped=(PIMAGE_EXPORT_DIRECTORY) &lpBuffer[ExportTableRVA];

    DWORD *FunctionNamePtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNames];
    DWORD *FunctionAdressPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfFunctions];
    DWORD *FunctionOrdinalPtr=(DWORD*) &lpBuffer[(DWORD)ped->AddressOfNameOrdinals];
    int i;
    for (i=0; i<(int)ped->NumberOfNames; i++)
    {
      char *NamePtr=(char*) &lpBuffer[*FunctionNamePtr];
      if (!StringCompare(NamePtr,lpProcName))
      {
        return( (DWORD) FunctionAdressPtr+dwImageBase);
      }
      FunctionNamePtr++;
      FunctionAdressPtr++;
      FunctionOrdinalPtr++;
    }
	}
	return(0);
}


/*
DWORD dwGetExportedFunctionAddrFromImageBase(DWORD dwImageBase,LPCSTR lpProcName)
{
  DWORD dwAddr=dwGetExportedFunctionEntryFromImageBase(dwImageBase,lpProcName);
  if (dwAddr)
  {
    dwAddr= dwImageBase + (*((DWORD*) dwAddr));
  }
  return(dwAddr);
}
*/

/*
//
// Export Format
//

typedef struct _IMAGE_EXPORT_DIRECTORY {
    DWORD   Characteristics;
    DWORD   TimeDateStamp;
    WORD    MajorVersion;
    WORD    MinorVersion;
    DWORD   Name;
    DWORD   Base;
    DWORD   NumberOfFunctions;
    DWORD   NumberOfNames;
    PDWORD  *AddressOfFunctions;
    PDWORD  *AddressOfNames;
    PWORD   *AddressOfNameOrdinals;
} IMAGE_EXPORT_DIRECTORY, *PIMAGE_EXPORT_DIRECTORY;

//
// Import Format
//

typedef struct _IMAGE_IMPORT_BY_NAME {
    WORD    Hint;
    BYTE    Name[1];
} IMAGE_IMPORT_BY_NAME, *PIMAGE_IMPORT_BY_NAME;

typedef struct _IMAGE_THUNK_DATA {
    union {
        PBYTE  ForwarderString;
        PDWORD Function;
        DWORD Ordinal;
        PIMAGE_IMPORT_BY_NAME AddressOfData;
    } u1;
} IMAGE_THUNK_DATA;
typedef IMAGE_THUNK_DATA * PIMAGE_THUNK_DATA;

#define IMAGE_ORDINAL_FLAG 0x80000000
#define IMAGE_SNAP_BY_ORDINAL(Ordinal) ((Ordinal & IMAGE_ORDINAL_FLAG) != 0)
#define IMAGE_ORDINAL(Ordinal) (Ordinal & 0xffff)

typedef struct _IMAGE_IMPORT_DESCRIPTOR {
    union {
        DWORD   Characteristics;                // 0 for terminating null import descriptor
        PIMAGE_THUNK_DATA OriginalFirstThunk;   // RVA to original unbound IAT
    };
    DWORD   TimeDateStamp;                  // 0 if not bound,
                                            // -1 if bound, and real date\time stamp
                                            //     in IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT (new BIND)
                                            // O.W. date/time stamp of DLL bound to (Old BIND)

    DWORD   ForwarderChain;                 // -1 if no forwarders
    DWORD   Name;
    PIMAGE_THUNK_DATA FirstThunk;           // RVA to IAT (if bound this IAT has actual addresses)
} IMAGE_IMPORT_DESCRIPTOR;
typedef IMAGE_IMPORT_DESCRIPTOR UNALIGNED *PIMAGE_IMPORT_DESCRIPTOR;

*/