/*
Frhed - Free hex editor
Copyright (C) 2000 Raihan Kibria

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

Last change: 2013-04-10 by Jochen Neubeck
*/
#include "precomp.h"
#include "DllProxies.h"
#include <imagehlp.h>
#include <ctype.h>
#include "resource.h"
#include "hexwnd.h"
#include "hexwdlg.h"
#include "gktools.h"
#include "simparr.h"
#include "StringTable.h"
#include "AnsiConvert.h"

static PList PartitionInfoList;

BOOL WINAPI GetDllExportNames(LPCTSTR pszFilename, ULONG_PTR* lpulOffset, ULONG_PTR* lpulSize)
{
	struct IMAGEHLP *IMAGEHLP = ::IMAGEHLP;
	if (IMAGEHLP == 0)
		return FALSE;

	LOADED_IMAGE li;
	if (!IMAGEHLP->MapAndLoad(const_cast<LPSTR>(&*static_cast<T2A>(pszFilename)), NULL, &li, TRUE, TRUE))
		return FALSE;

	BOOL bDone = FALSE;
	if (DWORD dw = li.FileHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress)
	{
		PIMAGE_EXPORT_DIRECTORY pExpDir = (PIMAGE_EXPORT_DIRECTORY)
			IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress, dw, 0);
		if (pExpDir)
		{
			DWORD firstName = pExpDir->Name;
			DWORD lastName = pExpDir->Name;
			if (pExpDir->NumberOfNames)
			{
				DWORD *pExpNames = (DWORD *)
					IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress, pExpDir->AddressOfNames, 0);
				if (pExpNames)
				{
					pExpNames += pExpDir->NumberOfNames - 1;
					lastName = *pExpNames;
				}
			}
			char *pszFirstName = static_cast<char *>(
				IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress, firstName, 0));
			char *pszLastName = static_cast<char *>(
				IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress, lastName, 0));
			if (pszFirstName && pszLastName)
			{
				*lpulOffset = pszFirstName - reinterpret_cast<char *>(li.MappedAddress);
				*lpulSize = pszLastName + strlen(pszLastName) + 1 - pszFirstName;
			}
		}
		bDone = TRUE;
	}
	IMAGEHLP->UnMapAndLoad(&li);
	return bDone;
}

BOOL WINAPI GetDllImportNames(LPCTSTR pszFilename, ULONG_PTR* lpulOffset, ULONG_PTR* lpulSize)
{
	struct IMAGEHLP *IMAGEHLP = ::IMAGEHLP;
	if (IMAGEHLP == 0)
		return FALSE;

	LOADED_IMAGE li;
	if (!IMAGEHLP->MapAndLoad(const_cast<LPSTR>(&*static_cast<T2A>(pszFilename)), NULL, &li, TRUE, TRUE))
		return FALSE;

	BOOL bDone = FALSE;
	if (DWORD dw = li.FileHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress)
	{
		IMAGE_IMPORT_DESCRIPTOR *pDescriptor = (IMAGE_IMPORT_DESCRIPTOR *)
			IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress, dw, 0);
		char *upper = reinterpret_cast<char *>(li.MappedAddress);
		char *lower = upper + li.SizeOfImage;
		while (pDescriptor->FirstThunk)
		{
			if (char *name = static_cast<char *>(
				IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress,
				pDescriptor->Name, 0)))
			{
				if (lower > name)
					lower = name;
				char *end = name + strlen(name) + 1;
				if (upper < end)
					upper = end;
				if (IMAGE_THUNK_DATA *pThunk = static_cast<IMAGE_THUNK_DATA *>(
					IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress,
					pDescriptor->OriginalFirstThunk, 0)))
				{
					// https://github.com/adoxa/cmdkey/blob/master/edit.c
					while (pThunk->u1.AddressOfData)
					{
						IMAGE_IMPORT_BY_NAME *pImport = static_cast<IMAGE_IMPORT_BY_NAME *>(
							IMAGEHLP->ImageRvaToVa(li.FileHeader, li.MappedAddress, static_cast<ULONG>(pThunk->u1.AddressOfData), 0));
						if (pImport)
						{
							char *name = reinterpret_cast<char *>(pImport->Name);
							char *end = name + strlen(name);
							if (lower > name)
								lower = name;
							if (upper < end)
								upper = end;
						}
						++pThunk;
					}
				}
			}
			++pDescriptor;
		}
		if (lower < upper)
		{
			*lpulOffset = lower - reinterpret_cast<char *>(li.MappedAddress);
			*lpulSize = upper - lower;
			bDone = TRUE;
		}
	}
	IMAGEHLP->UnMapAndLoad(&li);
	return bDone;
}

static MEMORY_CODING_DESCRIPTION const BuiltinEncoders[] =
{
	{ "ROT-13", Rot13Encoder },
	{ "XOR -1", XorEncoder },
	{ 0, 0 }
};

static void AddEncoders(HListBox *pListbox, MEMORY_CODING_DESCRIPTION const *lpEncoders)
{
	for ( ; lpEncoders->lpszDescription ; ++lpEncoders)
	{
		int i = pListbox->AddString(static_cast<A2T>(lpEncoders->lpszDescription));
		pListbox->SetItemData(i, reinterpret_cast<DWORD_PTR>(lpEncoders));
	}
}

INT_PTR EncodeDecodeDialog::DlgProc(HWindow *pDlg, UINT uMsg, WPARAM wParam, LPARAM)
{
	switch (uMsg)
	{
	case WM_INITDIALOG:
		{
			HListBox *pListbox = static_cast<HListBox *>(pDlg->GetDlgItem(IDC_ENCODE_LIST));
			AddEncoders(pListbox, BuiltinEncoders);
			TCHAR path[MAX_PATH];
			GetModuleFileName(NULL, path, MAX_PATH);
			PathRemoveFileSpec(path);
			PathAppend(path, _T("Encoders\\*.dll"));
			WIN32_FIND_DATA ff;
			HANDLE h = FindFirstFile(path, &ff);
			if (h != INVALID_HANDLE_VALUE)
			{
				do
				{
					PathRemoveFileSpec(path);
					PathAppend(path, ff.cFileName);
					// Add encoder dll name with subfolder name
					HMODULE hLibrary = GetModuleHandle(path);
					if (hLibrary == 0)
						hLibrary = LoadLibrary(path);
					if (hLibrary)
					{
						if (LPFNGetMemoryCodings callback = (LPFNGetMemoryCodings)
							GetProcAddress(hLibrary, "GetMemoryCodings"))
						{
							AddEncoders(pListbox, callback());
						}
					}
				} while (FindNextFile(h, &ff));
				FindClose(h);
			}
			pListbox->SetCurSel(0);
			pDlg->CheckDlgButton(IDC_ENCODE_ENC, BST_CHECKED);
		}
		return TRUE;

	case WM_COMMAND:
		switch (wParam)
		{
		case IDOK:
			{
				MEMORY_CODING mc;
				TString sBuffer;
				pDlg->GetDlgItemText(IDC_ENCODE_ARGS, sBuffer);
				mc.bEncode = pDlg->IsDlgButtonChecked(IDC_ENCODE_ENC);
				mc.lpszArguments = sBuffer.c_str();
				HListBox *pListbox = static_cast<HListBox *>(pDlg->GetDlgItem(IDC_ENCODE_LIST));
				int nCurSel = pListbox->GetCurSel();
				if (nCurSel < 0)
					return TRUE;
				mc.fpMcd = reinterpret_cast<MEMORY_CODING_DESCRIPTION const *>(pListbox->GetItemData(nCurSel));
				int lower = 0;
				int upper = m_dataArray.size() - 1;
				if (bSelected)
				{
					lower = iGetStartOfSelection();
					upper = iGetEndOfSelection();
				}
				mc.lpbMemory = &m_dataArray[lower];
				mc.dwSize = upper - lower + 1;
				UndoRecord::Data *olddata = UndoRecord::alloc(mc.lpbMemory, mc.dwSize);
				mc.fpMcd->fpEncodeFunc(&mc);
				push_undorecord(lower, mc.dwSize, olddata);
				bFilestatusChanged = true;
				repaint();
			}
			// fall through
		case IDCANCEL:
			pDlg->EndDialog(wParam);
			return TRUE;
		}
		break;
	}
	return FALSE;
}

INT_PTR OpenDriveDialog::DlgProc(HWindow *pDlg, UINT uMsg, WPARAM wParam, LPARAM)
{
	switch (uMsg)
	{
	case WM_INITDIALOG:
		{
			pDlg->ShowWindow(SW_SHOW);
			if (PartitionInfoList.IsEmpty())
			{
				if (IPhysicalDrive *Drive = CreatePhysicalDriveInstance())
				{
					Drive->GetPartitionInfo(&PartitionInfoList);
					delete Drive;
				}
			}
			HListBox *list = static_cast<HListBox *>(pDlg->GetDlgItem(IDC_DRIVE_LIST));
			LIST_ENTRY *Flink = &PartitionInfoList;
			while ((Flink = Flink->Flink) != &PartitionInfoList)
			{
				PartitionInfo *pi = static_cast<PartitionInfo *>(Flink);
				int i = list->AddString(pi->GetNameAsString());
				list->SetItemData(i, reinterpret_cast<DWORD_PTR>(pi));
			}
			list->SetCurSel(0);
		}
		return TRUE;

	case WM_COMMAND:
		switch (wParam)
		{
		case IDOK:
			{
				HListBox *list = static_cast<HListBox *>(pDlg->GetDlgItem(IDC_DRIVE_LIST));
				int nCurSel = list->GetCurSel();
				if (nCurSel < 0)
					return TRUE;
				PartitionInfo *SelectedPartitionInfo =
					reinterpret_cast<PartitionInfo *>(list->GetItemData(nCurSel));

				IPhysicalDrive *Drive = CreatePhysicalDriveInstance();
				if (Drive == 0 || !Drive->Open(SelectedPartitionInfo->m_dwDrive))
				{
					MessageBox(pwnd, GetLangString(IDS_DRIVES_ERR_OPEN), MB_ICONERROR);
					delete Drive;
					return TRUE;
				}

				delete HexEditorWindow::Drive;
				HexEditorWindow::Drive = Drive;
				HexEditorWindow::SelectedPartitionInfo = SelectedPartitionInfo;
				EnableDriveButtons(TRUE);
				if (Track.GetObjectSize() != SelectedPartitionInfo->m_dwBytesPerSector)
					Track.Create(SelectedPartitionInfo->m_dwBytesPerSector);
				CurrentSectorNumber = 0;
				RefreshCurrentTrack();
			}
			// fall through
		case IDCANCEL:
			pDlg->EndDialog(wParam);
			return TRUE;
		}
		break;
	}
	return FALSE;
}

INT_PTR GotoTrackDialog::DlgProc(HWindow *pDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	switch (uMsg)
	{
	case WM_INITDIALOG:
		{
			TCHAR szTempBuffer[10240];

			_stprintf(szTempBuffer, _T("%I64d"), CurrentSectorNumber);
			pDlg->SetDlgItemText(IDC_DRIVE_TRACK, szTempBuffer);

			DISK_GEOMETRY dg;
			Drive->GetDriveGeometry(&dg);

			INT64 TotalSizeInBytes = dg.SectorsPerTrack;
			TotalSizeInBytes *= dg.BytesPerSector;
			TotalSizeInBytes *= dg.TracksPerCylinder;
			TotalSizeInBytes *= dg.Cylinders.QuadPart;

			_stprintf(szTempBuffer,
				_T("%s = %I64d\r\n")
				_T("%s = %I64d\r\n")
				_T("%s = %ld\r\n")
				_T("%s = %ld\r\n")
				_T("%s = %ld\r\n")
				_T("%s = %I64d\r\n"),
				GetLangString(IDS_DRIVES_CYLINDERS),
				dg.Cylinders.QuadPart,
				GetLangString(IDS_DRIVES_SECTORS),
				SelectedPartitionInfo->m_NumberOfSectors,
				GetLangString(IDS_DRIVES_TRACSPERCYL),
				dg.TracksPerCylinder,
				GetLangString(IDS_DRIVES_SECTPERTRACK),
				dg.SectorsPerTrack,
				GetLangString(IDS_DRIVES_BYTESPERSECT),
				dg.BytesPerSector,
				GetLangString(IDS_DRIVES_TOTALBYTES),
				TotalSizeInBytes);

			pDlg->SetDlgItemText(IDC_DRIVE_INFO, szTempBuffer);
		}
		return FALSE;

	case WM_COMMAND:
		switch (LOWORD (wParam))
		{
		case IDOK:
			{
				TCHAR szBuffer[256];
				pDlg->GetDlgItemText(IDC_DRIVE_TRACK, szBuffer, RTL_NUMBER_OF(szBuffer));

				INT64 TempCurrentSectorNumber = 0;
				_stscanf(szBuffer, _T("%I64d"), &TempCurrentSectorNumber);
				if (TempCurrentSectorNumber < 0 ||
					TempCurrentSectorNumber >= SelectedPartitionInfo->m_NumberOfSectors)
				{
					return TRUE;
				}
				CurrentSectorNumber = TempCurrentSectorNumber;
				RefreshCurrentTrack();
			}
			// fall through
		case IDCANCEL:
			pDlg->EndDialog(wParam);
			return TRUE;
		}
		break;
	}
	return FALSE;
}
