/*
	XDSLOAD.CPP
	2002, Lorenzo Lutti (lorenzo.lutti@tiscalinet.it)

	This program loads a process and handles its I/O by debugging it and trapping
	the PRIVILEGED_INSTRUCTION exception.

	Note 1: I've tested XDSLOAD only under Windows 2000; probably it runs under XP as well.
	The OpenThread() function is supported by Win2000, WinXP and WinME; I suppose
	there are other undocumented ways to get the same thing (retrieving a HANDLE from a TID) under
	the other OSes.

	Note 2: XDSLOAD works only for programs not designed for NT/2000/XP, that normally
	would crash	with the "Illegal instruction" debug message. It doesn't work with programs that
	access the I/O space in a "legal" way (kernel-mode driver, "unlocking" the I/O instructions
	by using Ke386SetIoAccessMap(), intercepting the I/O with a VDD and so on), simply because
	they don't throw any debugging exception when they access the I/O space.

	VERSION HISTORY

	08-oct-2002 (Lorenzo Lutti): first release.
*/


#include "stdafx.h"
#include "xdsload.h"
#include "xdsloadDlg.h"

#include "xds510\\xds510.h"

#define CHAR_LIMIT 10000000

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

typedef HANDLE (WINAPI *ptrOpenThread)(DWORD,BOOL,DWORD);
ptrOpenThread	OpenThread;

typedef struct
{
	CXdsloadDlg *dlg;
	PROCESS_INFORMATION pi;
} THREAD_DATA;

/////////////////////////////////////////////////////////////////////////////
// CXdsloadDlg dialog

CXdsloadDlg::CXdsloadDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CXdsloadDlg::IDD, pParent)
{
	//{{AFX_DATA_INIT(CXdsloadDlg)
	//}}AFX_DATA_INIT
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
	m_thread = NULL;
}

void CXdsloadDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	//{{AFX_DATA_MAP(CXdsloadDlg)
	DDX_Control(pDX, IDC_CHECK1, m_verbose);
	DDX_Control(pDX, IDC_RICHEDIT1, m_monitor);
	DDX_Control(pDX, IDOK, m_run);
	DDX_Control(pDX, IDC_EDIT1, m_path);
	//}}AFX_DATA_MAP
}

BEGIN_MESSAGE_MAP(CXdsloadDlg, CDialog)
	//{{AFX_MSG_MAP(CXdsloadDlg)
	ON_BN_CLICKED(IDC_BUTTON1, OnBrowse)
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()

/////////////////////////////////////////////////////////////////////////////
// CXdsloadDlg message handlers

BOOL CXdsloadDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon

	// Retrieve OpenThread() address (not necessary if you have an updated PSDK)
	HMODULE hMod = LoadLibrary( "kernel32.dll");
	if( hMod == NULL)
	{
		AfxMessageBox( "LoadLibrary");
		EndDialog( IDCANCEL);
		return FALSE;
	}
	OpenThread = (ptrOpenThread)GetProcAddress( hMod, "OpenThread");
	if( OpenThread == NULL)
	{
		AfxMessageBox( "GetProcAddress");
		EndDialog( IDCANCEL);
		return FALSE;
	}

	// I/O grant
	m_wrapper.GiveIO();

	// XDS init
	XDS510_init();

	return TRUE;  // return TRUE  unless you set the focus to a control
}

void CXdsloadDlg::OnBrowse() 
{
	CFileDialog dlg( TRUE, NULL, NULL, OFN_HIDEREADONLY | OFN_PATHMUSTEXIST |
		OFN_FILEMUSTEXIST, "Executable files (*.exe)|*.exe||", this);

	if( dlg.DoModal() == IDOK)
		m_path.SetWindowText( dlg.GetPathName());
}

void CXdsloadDlg::OnOK() 
{
	m_run.EnableWindow( FALSE);

	m_monitor.SetWindowText( "");
	m_monitor.SetFocus();

	// Execute process (initially suspended)
	CString path;
	m_path.GetWindowText( path);
	PROCESS_INFORMATION pi;
	STARTUPINFO si;
	::FillMemory( &si, sizeof(si), 0);
	si.cb = sizeof( si);
	if( !CreateProcess( NULL, (char*)(LPCTSTR)path, NULL, NULL, FALSE, CREATE_SUSPENDED, NULL, NULL, &si, &pi))
	{
		m_run.EnableWindow( TRUE);
		AfxMessageBox( "CreateProcess");
		return;
	}

	// Create the debug thread
	THREAD_DATA *td = new THREAD_DATA;
	td->pi = pi;
	td->dlg = this;
	if( (m_thread = AfxBeginThread( DebugThread, (LPVOID)td  )) == NULL)
	{
		delete td;
		m_run.EnableWindow( TRUE);
		TerminateProcess( pi.hProcess, 0);
		AfxMessageBox( "CreateThread");
		return;		
	}
}

UINT CXdsloadDlg::DebugThread( LPVOID pParam)
{
	PROCESS_INFORMATION pi = ((THREAD_DATA*)pParam)->pi;
	CXdsloadDlg* dlg = ((THREAD_DATA*)pParam)->dlg;
	delete ((THREAD_DATA*)pParam);
	BOOL bVerbose = dlg->m_verbose.GetCheck() ? TRUE : FALSE;

	// Debug the process
	DEBUG_EVENT de;
	if( !DebugActiveProcess( pi.dwProcessId))
	{
		TerminateProcess( pi.hProcess, 0);
		AfxMessageBox( "DebugActiveProcess");
		dlg->m_run.EnableWindow( TRUE);
		return 0;
	}
	
	// Resume the process
	ResumeThread( pi.hThread);
	
	// Debugging loop
	while( TRUE)
	{
		// Wait indefinitely for a debug event
		if( !WaitForDebugEvent( &de, INFINITE))
		{
			TerminateProcess( pi.hProcess, 0);
			AfxMessageBox( "WaitForDebugEvent");
			dlg->m_run.EnableWindow( TRUE);
			return 0;
		}

		CString str;
		CString oldtxt;

		// Handle only the EXCEPTION_DEBUG_EVENT event
		if( de.dwDebugEventCode != EXCEPTION_DEBUG_EVENT)
		{
			// Event dump
			str = DebugEventString( de.dwDebugEventCode) + "\r\n";
			dlg->m_monitor.GetWindowText( oldtxt);
			int nChars = oldtxt.GetLength();
			if( nChars < CHAR_LIMIT)
			{
				dlg->m_monitor.SetSel( nChars, -1);
				dlg->m_monitor.ReplaceSel( str);				
			}
			
			// If the event is EXIT_PROCESS_DEBUG_EVENT, exit
			if( de.dwDebugEventCode == EXIT_PROCESS_DEBUG_EVENT)
			{
				dlg->m_run.EnableWindow( TRUE);
				return 0;				
			}

			// Always consider continuables these exceptions (TODO: probably wrong)
			ContinueDebugEvent( de.dwProcessId, de.dwThreadId, DBG_CONTINUE);
			continue;
		}
		else
		{
			// Handle only the STATUS_PRIVILEGED_INSTRUCTION exception
			if( de.u.Exception.ExceptionRecord.ExceptionCode != STATUS_PRIVILEGED_INSTRUCTION)
			{
				// Exception dump
				BOOL known;
				CString str = DebugEventString( de.dwDebugEventCode) +
					" " + ExceptionString( de.u.Exception.ExceptionRecord.ExceptionCode, &known) + "\r\n";
				dlg->m_monitor.GetWindowText( oldtxt);
				int nChars = oldtxt.GetLength();
				if( nChars < CHAR_LIMIT)
				{
					dlg->m_monitor.SetSel( nChars, -1);
					dlg->m_monitor.ReplaceSel( str);				
				}

				// Don't handle the unknown exceptions
				// (frequently the programs use them for their business)
				if( !known)
				{
					ContinueDebugEvent( de.dwProcessId, de.dwThreadId, DBG_EXCEPTION_NOT_HANDLED);
				}
				else
				// Continue for first chance (maybe someone can handle it...) or continuable exceptions
				if( de.u.Exception.dwFirstChance ||
					!(de.u.Exception.ExceptionRecord.ExceptionFlags & EXCEPTION_NONCONTINUABLE))
				{
					ContinueDebugEvent( de.dwProcessId, de.dwThreadId, DBG_CONTINUE);
				}
				else
				// Otherwise, game over.
				{
					TerminateProcess( pi.hProcess, 0);
					AfxMessageBox( "Second-chance exception");
					dlg->m_run.EnableWindow( TRUE);
					return 0;
				}
				
				continue;				
			}

			// Read the opcode that caused the exception
			HANDLE hProcess = OpenProcess( PROCESS_ALL_ACCESS, FALSE, de.dwProcessId);
			if( hProcess == NULL)
			{
				TerminateProcess( pi.hProcess, 0);
				AfxMessageBox( "OpenProcess");
				dlg->m_run.EnableWindow( TRUE);
				return 0;
			}
			unsigned char opcodes[2];
			DWORD dwBytesRead;
			if( !ReadProcessMemory( hProcess, de.u.Exception.ExceptionRecord.ExceptionAddress,
				opcodes, 2, &dwBytesRead))
			{
				TerminateProcess( pi.hProcess, 0);
				AfxMessageBox( "ReadProcessMemory");
				dlg->m_run.EnableWindow( TRUE);
				return 0;
			}