/* 
Saffron Dynamic Instrumentation Code
Version 0.2a 
Released 8/2/2007 at Blackhat USA 2007
Updates:
3/28/2008 DQ Fixed dumping size, validated with PIN 2.3-17236, IA32

Written by Danny Quist
Copyright (C) 2008 Offensive Computing, LLC
http://www.offensivecomputing.net/

DOWNLOAD PIN

http://rogue.colorado.edu/pin/

This was tested on PIN 2.1-12211 on Windows XP SP2 and Vista both in 32-bit mode.

I recommend getting the Visual Studio targetted release of PIN, as I've had good
luck with.  Once you have installed it, put this file into the "SimpleExamples" 
directory and recompile.

BUILD PIN

In Visual Studio go to Build->Build Solution.  Everything should build correctly.

USAGE

Change to the SimpleExamples directory inside of the pin folder.  Have the executable 
you are interested in ready to run.  To run Saffron, execute this command:

C:\pin-2.1...\SimpleExamples>..\bin\pin.exe -t saffron-di -- yourexehere.exe

Saffron will then deposit a dumped version of the executable into the directory.  Inspect
each one till you find the one that yields the most information.  A good choice tends
to be the last one.  

A log will be made in saffron-di.out in the SimpleExamples directory which contains some
useful (maybe) information.

Caveats:
 * 32-bit only!  This could be ported to IA-32e and Itanium
 * These executables probably will not run, they are meant for analysis only
 * The entry point in each of the dumps is a guess based on the current value of EIP at
   the time of the dump

Thanks to Lorie Liebrock, Valsmith, Ty Bodell, and Houdini for their help

This software is available as-is, with no warranty expressed or implied.
It is free for non-commercial and educational use only.  No other use 
is permitted without express written permission from 
Offensive Computing, LLC
*/

#include "pin.H"
#include "portability.H"
#include <iostream>
#include <map>

namespace WINDOWS
{
#include <windows.h>
#include <tlhelp32.h>
}
#define MakePtr( cast, ptr, addValue ) (cast)( (WINDOWS::DWORD)(ptr) + (WINDOWS::DWORD)(addValue))

UINT64 ins_count = 0;
FILE *fout;
#define FILENAME_SIZE 128
char filename[FILENAME_SIZE];
map<unsigned int, unsigned int> writtenAddrs;
unsigned int dumpnum = 0;

#define OCRET_OK 0
#define OCRET_ERROR 1

namespace WINDOWS 
{
typedef struct {
  WORD  woNumOfSect;
  DWORD dwImageBase;
  DWORD dwSizeOfImage;
  DWORD dwAddrOfEP;
  DWORD dwBaseOfCode;
  DWORD dwBaseOfData;
} PEFILEINFO, *LPPEFILEINFO;

typedef struct {
  BYTE  byName[IMAGE_SIZEOF_SHORT_NAME];
  DWORD dwVSize;
  DWORD dwVOffset;
  DWORD dwRSize;
  DWORD dwROffset;
  DWORD dwCharacteristics;
} SECTIONINFO, *LPSECTIONINFO;

class pedump
{
public:
	pedump(char *infile, unsigned char *newoep);
	void saveDumpFile(char *outfile);
	~pedump(void);

private:
	HANDLE hFile;
	HANDLE hHeap;
	PEFILEINFO PEFileInfo;
	LPSECTIONINFO lpSectInfo;
	LPBYTE lpDumpData;
	DWORD dwFileSize;
	unsigned char *oep;
	int getPEDataFromFile(char *infile);
	int populateNewPEDump();
};

void pedump::saveDumpFile(char *outfile)
{
	HANDLE hOutFile = CreateFile(outfile, GENERIC_READ | GENERIC_WRITE, 0, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	DWORD dwAccBytes = 0;

	if(hFile != INVALID_HANDLE_VALUE) 
	{
		SetFilePointer(hFile, 0, 0, FILE_BEGIN);
		fprintf(fout, "Dumping %u bytes from 0x%x\n", dwFileSize, lpDumpData);
		if (!WriteFile(hFile, lpDumpData, PEFileInfo.dwSizeOfImage, &dwAccBytes, NULL))
		{
			fprintf(fout, "Error writing file (%u)\n", GetLastError());
		}
		fprintf(fout, "Dumped %u bytes\n", dwAccBytes);
		CloseHandle(hFile);
    }
	else
	{
		fprintf(fout, "Could not write to %s\n", outfile);
	}

	return;
}

pedump::pedump(char *infile, unsigned char *newoep)
{
	lpSectInfo = NULL;
	oep = newoep;
	if ( getPEDataFromFile(infile) != OCRET_OK )
	{
		printf("getPEDataFromFile() failed\n");
		return;
	}

	if ( populateNewPEDump() != OCRET_OK )
	{
		printf("populateNewPEDump() failed\n");
		return;
	}
}

int pedump::getPEDataFromFile(char *infile)
{
	DWORD dwReadSize = 0;
	HANDLE hHeap = 0;
	LPBYTE fbuf = NULL;
	PIMAGE_DOS_HEADER inDosHeader = NULL;
	PIMAGE_NT_HEADERS inPeHeader = NULL;
	PIMAGE_SECTION_HEADER inSectionHeader = NULL;

	hFile = CreateFile(infile, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

	if (hFile == INVALID_HANDLE_VALUE)
	{
		printf("Error reading file\n");
		return OCRET_ERROR;
	}

	dwFileSize = GetFileSize(hFile, NULL);
	hHeap = HeapCreate(HEAP_NO_SERIALIZE, 1, 0);
	fbuf = (LPBYTE) HeapAlloc(hHeap, 0, dwFileSize);

	if (fbuf == NULL)
	{
		printf("Error allocating memory\n");
		return OCRET_ERROR;
	}

	if(ReadFile(hFile,fbuf,dwFileSize,&dwReadSize,NULL) == 0) 
	{
		printf("Error reading from file\n");
		HeapFree(hHeap,HEAP_NO_SERIALIZE,fbuf);
		return OCRET_ERROR;
	}

	CloseHandle(hFile);

	inDosHeader = (PIMAGE_DOS_HEADER) fbuf;
	
	if (inDosHeader->e_magic != IMAGE_DOS_SIGNATURE)
	{
		printf("File is not a PE file (invalid DOS signature\n");
		HeapFree(hHeap,HEAP_NO_SERIALIZE,fbuf);
		return OCRET_ERROR;
	}

	inPeHeader = (PIMAGE_NT_HEADERS)(fbuf + inDosHeader->e_lfanew);
	if(inPeHeader->Signature != IMAGE_NT_SIGNATURE) 
	{
		printf("File is not a PE file (invalid NT signature)\n");
		HeapFree(hHeap,HEAP_NO_SERIALIZE,fbuf);
		return OCRET_ERROR;
	}

	PEFileInfo.woNumOfSect   = inPeHeader->FileHeader.NumberOfSections;
	PEFileInfo.dwImageBase   = inPeHeader->OptionalHeader.ImageBase;
	PEFileInfo.dwSizeOfImage = inPeHeader->OptionalHeader.SizeOfImage;
	PEFileInfo.dwBaseOfCode  = inPeHeader->OptionalHeader.BaseOfCode ;
	PEFileInfo.dwBaseOfData  = inPeHeader->OptionalHeader.BaseOfData ;
	PEFileInfo.dwAddrOfEP    = (DWORD) oep - inPeHeader->OptionalHeader.ImageBase;

	lpSectInfo = (LPSECTIONINFO)malloc(sizeof(SECTIONINFO)*(PEFileInfo.woNumOfSect+1));

	if (!lpSectInfo)
	{
		printf("malloc failed\n");
		HeapFree(hHeap, HEAP_NO_SERIALIZE, fbuf);
		return OCRET_ERROR;
	}

	ZeroMemory(lpSectInfo,sizeof(SECTIONINFO)*(PEFileInfo.woNumOfSect+1));
	inSectionHeader = IMAGE_FIRST_SECTION(inPeHeader);
	for(int i=0; i<(int)PEFileInfo.woNumOfSect; i++) 
	{
		memcpy((lpSectInfo+i)->byName, (inSectionHeader+i)->Name, IMAGE_SIZEOF_SHORT_NAME);
		(lpSectInfo+i)->dwVSize            = (inSectionHeader+i)->Misc.VirtualSize;
		(lpSectInfo+i)->dwVOffset          = (inSectionHeader+i)->VirtualAddress;
		(lpSectInfo+i)->dwRSize            = (inSectionHeader+i)->SizeOfRawData;
		(lpSectInfo+i)->dwROffset          = (inSectionHeader+i)->PointerToRawData;
		(lpSectInfo+i)->dwCharacteristics  = (inSectionHeader+i)->Characteristics;
	}

	HeapFree(hHeap, HEAP_NO_SERIALIZE, fbuf);

	return OCRET_OK;
}

int pedump::populateNewPEDump()
{
	DWORD dwFrom = 0;
	DWORD dwSize = 0;
	DWORD dwAccBytes = 0;

	PIMAGE_DOS_HEADER inDosHeader = NULL;
	PIMAGE_NT_HEADERS inPeHeader = NULL;
	PIMAGE_SECTION_HEADER inSectionHeader = NULL;

	dwFrom = PEFileInfo.dwImageBase;
	dwSize = PEFileInfo.dwSizeOfImage;

	hHeap = HeapCreate(HEAP_NO_SERIALIZE,1,0);
	lpDumpData = (LPBYTE) HeapAlloc(hHeap,HEAP_NO_SERIALIZE | HEAP_ZERO_MEMORY,dwSize);
	memcpy(lpDumpData, (void *) dwFrom, dwSize);

	inDosHeader = (PIMAGE_DOS_HEADER) lpDumpData;
	
	if (inDosHeader->e_magic != IMAGE_DOS_SIGNATURE)
	{
		printf("File is not a PE file (invalid DOS signature\n");
		HeapFree(hHeap,HEAP_NO_SERIALIZE,lpDumpData);
		return OCRET_ERROR;
	}

	inPeHeader = (PIMAGE_NT_HEADERS)(lpDumpData + inDosHeader->e_lfanew);
	if(inPeHeader->Signature != IMAGE_NT_SIGNATURE) 
	{
		printf("File is not a PE file (invalid NT signature)\n");
		HeapFree(hHeap,HEAP_NO_SERIALIZE,lpDumpData);
		return OCRET_ERROR;
	}

	inPeHeader->FileHeader.NumberOfSections        = PEFileInfo.woNumOfSect;
	inPeHeader->OptionalHeader.ImageBase           = PEFileInfo.dwImageBase;
	inPeHeader->OptionalHeader.SizeOfImage         = PEFileInfo.dwSizeOfImage;
	inPeHeader->OptionalHeader.BaseOfCode          = PEFileInfo.dwBaseOfCode;
	inPeHeader->OptionalHeader.BaseOfData          = PEFileInfo.dwBaseOfData;
	inPeHeader->OptionalHeader.AddressOfEntryPoint = PEFileInfo.dwAddrOfEP;

	inSectionHeader = IMAGE_FIRST_SECTION(inPeHeader);
	for(int i=0; i<(int)PEFileInfo.woNumOfSect; i++) 
	{
		memcpy((inSectionHeader+i)->Name,(lpSectInfo+i)->byName, IMAGE_SIZEOF_SHORT_NAME);
		(inSectionHeader+i)->Misc.VirtualSize = (lpSectInfo+i)->dwVSize;
		(inSectionHeader+i)->VirtualAddress   = (lpSectInfo+i)->dwVOffset;
		(inSectionHeader+i)->SizeOfRawData    = (((lpSectInfo+i)->dwRSize) < ((lpSectInfo+i)->dwVSize) ? ((lpSectInfo+i)->dwVSize) : ((lpSectInfo+i)->dwRSize));
		(inSectionHeader+i)->PointerToRawData = (lpSectInfo+i)->dwROffset;
		(inSectionHeader+i)->Characteristics  = (lpSectInfo+i)->dwCharacteristics;
		fprintf(fout, "Name: %s VirtualSize: %u VirtualAddress: 0x%8.8x SizeOfRawData: %u PointerToRawData: 0x%8.8x Characteristics %8.8x\n",
			(inSectionHeader+i)->Name,
			(inSectionHeader+i)->Misc.VirtualSize,
			(inSectionHeader+i)->VirtualAddress,
			(inSectionHeader+i)->SizeOfRawData,
			(inSectionHeader+i)->PointerToRawData,
			(inSectionHeader+i)->Characteristics);
	}

	return OCRET_OK;
}

pedump::~pedump(void)
{
	if (lpSectInfo != NULL)
		free (lpSectInfo);

	if (lpDumpData != NULL)
		HeapFree(hHeap, HEAP_NO_SERIALIZE, lpDumpData);

	lpSectInfo = NULL;
	lpDumpData = NULL;
}

} // end namespace WINDOWS
/* ===================================================================== */

INT32 Usage()
{
    cerr <<
        "This tool prints out the number of dynamic instructions executed to stderr.\n"
        "\n";

    cerr << KNOB_BASE::StringKnobSummary();

    cerr << endl;

    return -1;
}

// Read the PE Optional Header in order to get ImageBase and ImageSize needed for dumper
int ReadPEHeader(unsigned int *ImageBase, unsigned int *ImageSize, WINDOWS::LPSTR fname)
{
	WINDOWS::HANDLE hFile = NULL;
    WINDOWS::HANDLE hFileMapping = NULL;
    WINDOWS::LPVOID lpFileBase = NULL;
    WINDOWS::PIMAGE_DOS_HEADER dosHeader = NULL;
    WINDOWS::PIMAGE_NT_HEADERS pNTHeader = NULL;
	
    hFile = WINDOWS::CreateFile(fname, GENERIC_READ, FILE_SHARE_READ, NULL,
                        OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
                    
    hFileMapping = WINDOWS::CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
   
    lpFileBase = WINDOWS::MapViewOfFile(hFileMapping, FILE_MAP_READ, 0, 0, 0);
   
    dosHeader = (WINDOWS::PIMAGE_DOS_HEADER)lpFileBase;

    if ( dosHeader->e_magic == IMAGE_DOS_SIGNATURE )
    {
        pNTHeader = MakePtr( WINDOWS::PIMAGE_NT_HEADERS, dosHeader, dosHeader->e_lfanew );
		
		if ( pNTHeader->Signature == IMAGE_NT_SIGNATURE )
        {
			WINDOWS::PIMAGE_OPTIONAL_HEADER pOptHeader = (WINDOWS::PIMAGE_OPTIONAL_HEADER)&pNTHeader->OptionalHeader;

			*ImageBase = (unsigned int) pOptHeader->ImageBase;
			*ImageSize = (unsigned int) pOptHeader->SizeOfImage;
        }
    }
	else 
	{
        printf("unrecognized file format\n");
	}

    WINDOWS::UnmapViewOfFile(lpFileBase);
    WINDOWS::CloseHandle(hFileMapping);
    WINDOWS::CloseHandle(hFile);

	return(1);

} // End ReadPEHeader

// Dump PE process memory to a file. IAT and file still needs to be rebuilt
void DumpProcMem(unsigned int ImageBase, unsigned int ImageSize, unsigned int newoep, WINDOWS::DWORD pid) 
{
	char foutname[FILENAME_SIZE];

	WINDOWS::SIZE_T ReadBytes = 0;
	WINDOWS::SIZE_T WriteBytes = 0;
	WINDOWS::PIMAGE_DOS_HEADER dosHeader = NULL;
    WINDOWS::PIMAGE_NT_HEADERS pNTHeader = NULL;
	WINDOWS::PIMAGE_FILE_HEADER pImgFileHdr = NULL;
	WINDOWS::PIMAGE_OPTIONAL_HEADER pOptHeader = NULL;

	unsigned char * buffer = (unsigned char *) malloc( ImageSize );

	memset(foutname, 0, sizeof(foutname));
	memcpy(buffer, (void *)ImageBase, ImageSize);

	sprintf(foutname, "dump-%u-%s", dumpnum++, filename);

	WINDOWS::HANDLE hFile = WINDOWS::CreateFile(foutname,   
                    GENERIC_READ|GENERIC_WRITE,         
                    0,      
                    NULL,                 
                    OPEN_ALWAYS,        
                    FILE_ATTRIBUTE_NORMAL,
                    NULL);                

	// Find the Entry Point
	dosHeader = (WINDOWS::PIMAGE_DOS_HEADER) buffer;
	if (dosHeader->e_magic == IMAGE_DOS_SIGNATURE)
	{
        pNTHeader = MakePtr( WINDOWS::PIMAGE_NT_HEADERS, dosHeader, dosHeader->e_lfanew );
		
		if ( pNTHeader->Signature == IMAGE_NT_SIGNATURE )
        {
			WINDOWS::PIMAGE_OPTIONAL_HEADER pOptHeader = (WINDOWS::PIMAGE_OPTIONAL_HEADER)&pNTHeader->OptionalHeader;
			
			fprintf(fout, "Current AddressOfEntryPoint 0x%8.8x setting to newoep (0x%x) RVA (0x%x)\n", 
				pOptHeader->AddressOfEntryPoint,
				newoep,
				newoep - ImageBase);
			pOptHeader->AddressOfEntryPoint = newoep - ImageBase;
        }
	}

    WINDOWS::WriteFile(hFile, buffer, ImageSize, &WriteBytes, NULL); 
    WINDOWS::CloseHandle(hFile);

	free(buffer);
}

VOID doDetach(void *ip)
{
	unsigned int ImageBase = 0;
	unsigned int ImageSize = 0;
	char outfilename[FILENAME_SIZE];

	memset(outfilename, 0, sizeof(outfilename));

	fprintf(fout, "ImageBase 0x%8.8x Halting EIP Addr: 0x%8.8x (Diff: 0x%x) ImageSize %x\n", 
		ImageBase,
		(unsigned int) ip,
		(unsigned int) ip - ImageBase,
		ImageSize
		);

	WINDOWS::pedump ocpe(filename, (unsigned char *) ip);
	sprintf(outfilename, "dump-%u-%s", dumpnum++, filename);
	ocpe.saveDumpFile(outfilename);
	
	PIN_Detach();
}

VOID measureip(void *ip)
{
	if (writtenAddrs.find((unsigned int) ip) != writtenAddrs.end())
	{
		doDetach(ip);
	}
}

VOID MemoryWrite(void *ip, void *addr)
{
	if (writtenAddrs.find((unsigned int) addr) == writtenAddrs.end())
	{
		writtenAddrs[(unsigned int) addr] = (unsigned int)ip;
	}

	if (writtenAddrs.find((unsigned int) ip) != writtenAddrs.end())
	{
		doDetach(ip);
	}
}

VOID Instruction(INS ins, VOID *v)
{
  INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)measureip, IARG_INST_PTR, IARG_END);
  if (INS_IsMemoryWrite(ins)) {
    INS_InsertPredicatedCall(ins,
			     IPOINT_BEFORE, (AFUNPTR) MemoryWrite,
			     IARG_INST_PTR,
			     IARG_MEMORYWRITE_EA, 
			     IARG_END);
			     }
			     
}

VOID Fini(INT32 code, VOID *v)
{
  fprintf(fout, "Count: %lu\n", ins_count);
  fclose(fout);
}

VOID detachCallback(VOID *v)
{
	
}

int main(int argc, char *argv[])
{
    if( PIN_Init(argc,argv) )
    {
        return Usage();
    }
    
    fout = fopen("saffron-di.out", "w");

	memset(filename, 0, sizeof(filename));
	strncpy(filename, argv[argc-1], FILENAME_SIZE-1);
	fprintf(fout, "Debugging %s\n", argv[argc-1]);
    INS_AddInstrumentFunction(Instruction, 0);
	PIN_AddDetachFunction(detachCallback, 0);

    // Never returns
    PIN_StartProgram();
    
    return 0;
}