//
// Copyright (c) 2002
// Ronald Kevin Burton
//
// Z poniższym kodem nie jest związana żadna gwarancja poprawności działania.
// Program został dołączony do książki ".NET CLR. Księga eksperta" w celu
// ilustracji koncepcji i zasad przedstawionych w tej książce. Program może być 
// używany na własne ryzyko.
//
// Przyznaje się prawo do użycia lub kopiowania tego oprogramowania do dowolnego celu
// bez konieczności ponoszenia żadnych opłat pod warunkiem, że powyższe uwagi zostaną 
// zachowane we wszystkich kopiach. Przyznaje się także prawo do modyfikacji kodu
// i dystrybucji zmodyfikowanego kodu pod warunkiem zachowania powyższych uwag
// oraz dołączenia informacji mówiącej o modyfikacji kodu.
//
// 
// ImportPage.cpp : plik implementacji
//

#include "stdafx.h"
#include "AssemblyDoc.h"
#include "AssemblyView.h"
#include "ImportPage.h"


// Okno dialogowe CImportPage

IMPLEMENT_DYNAMIC(CImportPage, CPropertyPage)
CImportPage::CImportPage()
	: CPropertyPage(CImportPage::IDD)
{
}

CImportPage::~CImportPage()
{
}

void CImportPage::DoDataExchange(CDataExchange* pDX)
{
	CPropertyPage::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_IMPORTTABLETREE, m_ctrlImportTableTree);
}

BEGIN_MESSAGE_MAP(CImportPage, CPropertyPage)
END_MESSAGE_MAP()

// Procedury obsługi wiadomości CImportPage
BOOL CImportPage::OnInitDialog()
{
	CPropertyPage::OnInitDialog();

	// Wypełnienie danymi
	CPropertySheet *pPropertySheet = STATIC_DOWNCAST(CPropertySheet, GetParent());
	CAssemblyView* pView = STATIC_DOWNCAST(CAssemblyView, pPropertySheet->GetParent());
	CAssemblyDoc* pDoc = pView->GetDocument();
	PBYTE pAssembly = pDoc->FileData();
	PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)pAssembly;
    if ( dosHeader->e_magic == IMAGE_DOS_SIGNATURE )
    {
		PIMAGE_NT_HEADERS pNTHeader;
		PIMAGE_NT_HEADERS64 pNTHeader64;

		// Utworzenie wskaźników do 32- i 64-bitowej wersji nagłówka.
		pNTHeader = MakePtr( PIMAGE_NT_HEADERS,
			                 dosHeader,
							 dosHeader->e_lfanew );

		pNTHeader64 = (PIMAGE_NT_HEADERS64)pNTHeader;

		// Najpierw weryfikacja, czy pole e_lfanew przekazało właściwy wskaźnik,
		// a następnie weryfikacja sygnatury PE.
		if ( IsBadReadPtr( pNTHeader, sizeof(pNTHeader->Signature) ) )
		{
			ATLTRACE(L"To nie jest plik typu Portable Executable (PE) EXE\n");
			return FALSE;
		}

		if ( pNTHeader->Signature != IMAGE_NT_SIGNATURE )
		{
			ATLTRACE(L"To nie jest plik typu Portable Executable (PE) EXE\n");
			return FALSE;
		}

		bool bIs64Bit = ( pNTHeader->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC );

		if(bIs64Bit)
		{
			DumpImportsSection(pAssembly, pNTHeader64 );
		}
		else
		{
			DumpImportsSection(pAssembly, pNTHeader);
		}
    }

	return TRUE;  // zwrócenie TRUE, chyba że ognisko ustawiono na element sterujący
}
template <class T, class U, class V>
void CImportPage::DumpImportsOfOneModule(	T* pINT, U* pIAT, V * pNTHeader,		// 'T', 'U' = IMAGE_THUNK_DATA, 'v' = IMAGE_NT_HEADERS
								PIMAGE_IMPORT_DESCRIPTOR pImportDesc,
								PBYTE pImageBase,
								HTREEITEM treeRowItem)
{
    PIMAGE_IMPORT_BY_NAME pOrdinalName;

 	WCHAR lBuffer[256];  
 	WCHAR tBuffer[256];  
	LPCSTR tableString;
	m_ctrlImportTableTree.InsertItem(L"Ordn  Name", treeRowItem);

	bool bIs64Bit = ( pNTHeader->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC );

	while ( 1 ) // Nieskończona pętla
	{
		if ( pINT->u1.AddressOfData == 0 )
			break;

		ULONGLONG ordinal = (ULONGLONG)-1;

		if ( bIs64Bit )
		{
			if ( IMAGE_SNAP_BY_ORDINAL64(pINT->u1.Ordinal) )
				ordinal = IMAGE_ORDINAL64(pINT->u1.Ordinal);
		}
		else
		{
			if ( IMAGE_SNAP_BY_ORDINAL32(pINT->u1.Ordinal) )
				ordinal = IMAGE_ORDINAL32(pINT->u1.Ordinal);			
		}

		if ( ordinal != -1 )
		{
			wsprintf(lBuffer, L"%4u", ordinal );
		}
		else
		{
			// pINT->u1.AddressOfData ma teoretycznie 32 lub 64 bity, ale w przypadku reprezentacji dyskowej możemy założyć,
			// że jest to RVA. Z tego powodu rzutowanie jest wykonywane do DWORD.
			pOrdinalName = (PIMAGE_IMPORT_BY_NAME)GetPtrFromRVA(static_cast<DWORD>(pINT->u1.AddressOfData), pNTHeader, pImageBase);
			
			tableString = (LPSTR)pOrdinalName->Name;
 			mbstowcs(tBuffer, tableString, strlen(tableString) + 1);
			wsprintf(lBuffer, L"%4u  %ls ", pOrdinalName->Hint, tBuffer);
		}
		
		// Jeśli obraz został powiązany, należy dołączyć
		// adres powiązania
		if ( pImportDesc->TimeDateStamp )
		{
			wsprintf( tBuffer, L" (Bound to: %08X)", pIAT->u1.Function );
			wcscat(lBuffer, tBuffer);
		}

		m_ctrlImportTableTree.InsertItem(lBuffer, treeRowItem);

		pINT++;         // Przejście do kolejnej pozycji
		pIAT++;         // Przejście do kolejnej pozycji
	}	
}

//
// Zrzut tablicy importowej (sekcja .idata) pliku PE.
//
template <class T> void CImportPage::DumpImportsSection(PBYTE pImageBase, T * pNTHeader)	// 'T' = PIMAGE_NT_HEADERS 
{
    PIMAGE_IMPORT_DESCRIPTOR pImportDesc;
    DWORD importsStartRVA;

    // Wyszukanie położenia sekcji importowej (zwykle, ale nie zawsze, w sekcji .idata).
    // Z tego powodu RVA jest pobierane z katalogu danych.
    importsStartRVA = GetImgDirEntryRVA(pNTHeader,IMAGE_DIRECTORY_ENTRY_IMPORT);
    if ( !importsStartRVA )
        return;

    pImportDesc = (PIMAGE_IMPORT_DESCRIPTOR)GetPtrFromRVA(importsStartRVA,pNTHeader,pImageBase);
	if ( !pImportDesc )
		return;
            
	bool bIs64Bit = ( pNTHeader->FileHeader.SizeOfOptionalHeader == IMAGE_SIZEOF_NT_OPTIONAL64_HEADER );

	 
 	WCHAR lBuffer[256];  
	HTREEITEM treeRowItem;
	LPCSTR tableString;
    while ( 1 )
    {
        // Sprawdzenie, czy osiągnięto pusty IMAGE_IMPORT_DESCRIPTOR
        if ( (pImportDesc->TimeDateStamp==0 ) && (pImportDesc->Name==0) )
            break;
        
		tableString = (LPCSTR)GetPtrFromRVA(pImportDesc->Name, pNTHeader, pImageBase);
 		mbstowcs(lBuffer, tableString, strlen(tableString) + 1);
		treeRowItem = m_ctrlImportTableTree.InsertItem(lBuffer);

        wsprintf(lBuffer, L"Import Lookup Table RVA   : 0x%08X",
      			 pImportDesc->Characteristics);
		m_ctrlImportTableTree.InsertItem(lBuffer, treeRowItem);

        wsprintf(lBuffer, L"TimeDateStamp                  : 0x%08X", pImportDesc->TimeDateStamp);
		if ( (pImportDesc->TimeDateStamp != 0) &&
			 (pImportDesc->TimeDateStamp != -1) )
		{
			wcscat(lBuffer, L" -> ");
			time_t timeStamp = pImportDesc->TimeDateStamp;
			wcscat(lBuffer, _wctime( &timeStamp ));
		}
		m_ctrlImportTableTree.InsertItem(lBuffer, treeRowItem);

        wsprintf(lBuffer, L"ForwarderChain                  : 0x%08X", pImportDesc->ForwarderChain);
		m_ctrlImportTableTree.InsertItem(lBuffer, treeRowItem);
        wsprintf(lBuffer, L"DLL Name RVA                    : 0x%08X", pImportDesc->Name);
		m_ctrlImportTableTree.InsertItem(lBuffer, treeRowItem);
        wsprintf(lBuffer, L"Import Address Table RVA : 0x%08X", pImportDesc->FirstThunk);
		m_ctrlImportTableTree.InsertItem(lBuffer, treeRowItem);
    
        DWORD rvaINT = pImportDesc->OriginalFirstThunk;
        DWORD rvaIAT = pImportDesc->FirstThunk;

        if ( rvaINT == 0 )   // Brak pola Characteristics?
        {
            // Tak! Należy uzyskać niezerowe pole FirstThunk.
            rvaINT = rvaIAT;
            
			// Brak pola FirstThunk? Niedobrze!!!
            if ( rvaINT == 0 )
                return;
        }
        
        // Dopasowanie wskaźnika do położenia tablic w pliku mapowania pamięci.
        PIMAGE_THUNK_DATA pINT = (PIMAGE_THUNK_DATA)GetPtrFromRVA(rvaINT, pNTHeader, pImageBase);
		if (!pINT )
			return;

        PIMAGE_THUNK_DATA pIAT = (PIMAGE_THUNK_DATA)GetPtrFromRVA(rvaIAT, pNTHeader, pImageBase);
    
    
		if(bIs64Bit)
			DumpImportsOfOneModule( (PIMAGE_THUNK_DATA64)pINT, (PIMAGE_THUNK_DATA64)pIAT, pNTHeader, pImportDesc, pImageBase, treeRowItem);
		else
			DumpImportsOfOneModule( (PIMAGE_THUNK_DATA32)pINT, (PIMAGE_THUNK_DATA32)pIAT, pNTHeader, pImportDesc, pImageBase, treeRowItem);

		// Przejście do kolejnego IMAGE_IMPORT_DESCRIPTOR.
        pImportDesc++;
        ATLTRACE("\n");
    }

    printf("\n");
}

