From 69ddc0df13eba4cb5a9b76150eee19bf47f740a7 Mon Sep 17 00:00:00 2001
From: Force Charlie <charlieio@outlook.com>
Date: Sat, 5 Dec 2015 12:32:15 +0800
Subject: [PATCH] Fix bug

---
 PEAnalyzer/PortableExecutableFile.cpp | 17 ++++++++++++++---
 PEAnalyzer/PortableExecutableFile.h   |  5 -----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/PEAnalyzer/PortableExecutableFile.cpp b/PEAnalyzer/PortableExecutableFile.cpp
index 80582e9..e13286b 100644
--- a/PEAnalyzer/PortableExecutableFile.cpp
+++ b/PEAnalyzer/PortableExecutableFile.cpp
@@ -1,12 +1,18 @@
 /**
 **/
 #include "stdafx.h"
-#include <winnt.h>
+#include <winnt.h>
 #include <Dbghelp.h>
 #include "PortableExecutableFile.h"
 #include <codecvt>
 #pragma comment(lib,"DbgHelp.lib")
 
+#ifndef HINST_THISCOMPONENT
+EXTERN_C IMAGE_DOS_HEADER __ImageBase;
+#define HINST_THISCOMPONENT ((HINSTANCE)&__ImageBase)
+#endif
+
+
 enum FileMappingStatus {
 	kNotOptions,
 	kCreateFile,
@@ -210,7 +216,7 @@ bool PortableExecutableFile::Analyzer()
 	if (!winfile.CreateFileMap(L"PEAnalyzer.Executeable.MAP"))
 		return false;
 	char *baseAddress = (char*)winfile.MapViewOfFile();
-	IMAGE_DOS_HEADER *pDOSHeader = reinterpret_cast<IMAGE_DOS_HEADER *>(baseAddress);
+	IMAGE_DOS_HEADER *pDOSHeader = reinterpret_cast<IMAGE_DOS_HEADER *>(baseAddress);
 	IMAGE_NT_HEADERS *pNTHeader = reinterpret_cast<IMAGE_NT_HEADERS *>(baseAddress + pDOSHeader->e_lfanew);
 	if (largeFile.QuadPart < (LONGLONG)(sizeof(IMAGE_DOS_HEADER) + pDOSHeader->e_lfanew + sizeof(IMAGE_NT_HEADERS)))
 		return false;
@@ -256,7 +262,12 @@ bool PortableExecutableFile::Analyzer()
 	wsprintfW(imageVersion, L"%d.%d",
 			  pNTHeader->OptionalHeader.MajorImageVersion,
 			  pNTHeader->OptionalHeader.MinorImageVersion);
-	wsprintfW(entryPoint, L"0x%x", pNTHeader->OptionalHeader.AddressOfEntryPoint);
+	auto self = &__ImageBase;
+	auto selfNTHeader = reinterpret_cast<IMAGE_NT_HEADERS *>((char*)self + self->e_lfanew);
+	if (selfNTHeader->FileHeader.Machine != pNTHeader->FileHeader.Machine) {
+		clrMessage = L"PEAnaylzer's Architecture is different with Input PE File,cannot check CLR";
+		return true;
+	}
 	IMAGE_DATA_DIRECTORY *entry = &(pNTHeader->OptionalHeader).DataDirectory[IMAGE_DIRECTORY_ENTRY_COMHEADER];
 	if (entry->Size != sizeof(IMAGE_COR20_HEADER)) {
 		return true;
diff --git a/PEAnalyzer/PortableExecutableFile.h b/PEAnalyzer/PortableExecutableFile.h
index 09d2173..2f029b2 100644
--- a/PEAnalyzer/PortableExecutableFile.h
+++ b/PEAnalyzer/PortableExecutableFile.h
@@ -32,7 +32,6 @@ class PortableExecutableFile {
 	wchar_t osVersion[16];
 	wchar_t subsystemVersion[16];
 	wchar_t imageVersion[16];
-	wchar_t entryPoint[16];
 public:
 	PortableExecutableFile(const std::wstring &mPath);
 	bool Analyzer();
@@ -72,10 +71,6 @@ class PortableExecutableFile {
 	{
 		return imageVersion;
 	}
-	const wchar_t *GetEntryPoint()
-	{
-		return entryPoint;
-	}
 };
 
 #endif