#include <mach-o/dyld.h>
#include <syslog.h>
#include <pthread.h>
#include <CoreFoundation/CoreFoundation.h>
#include "SCPatchPrivate.h"
#include "SCPatchCommon.h"
#include "SCPatchController.h"
#include "mach_inject.h"

typedef list<SCPatchRecord>::iterator	SCPatchRecordIterator;

//-------------------------------------------------------------------------------------------------------------
mach_error_t		SCmac_err_FromOSErr(OSErr err) {
	return err ? (err_mac|err) : err_none;
}

//-------------------------------------------------------------------------------------------------------------
OSErr				SCOSErrFrom_mac_err(mach_error_t error) {
	return (error & err_mac) ? (err_mac) : noErr;
}

//-------------------------------------------------------------------------------------------------------------
#pragma mark -
//-------------------------------------------------------------------------------------------------------------
SCPatchController::SCPatchController()
{
	mBundle = CFBundleGetMainBundle();
	CFRetain(mBundle);
	
	mApplicationBundleIdentifier = CFBundleGetIdentifier(mBundle);
	CFRetain(mApplicationBundleIdentifier);
}

//-------------------------------------------------------------------------------------------------------------
SCPatchController::SCPatchController(CFStringRef bundleIdentifier)
{
    mBundle = CFBundleGetBundleWithIdentifier(bundleIdentifier);
	CFRetain(mBundle);

	mApplicationBundleIdentifier = bundleIdentifier;
	CFRetain(mApplicationBundleIdentifier);
}

//-------------------------------------------------------------------------------------------------------------
SCPatchController::~SCPatchController(void)
{
	if(mApplicationBundleIdentifier)
		CFRelease(mApplicationBundleIdentifier);
	
	if(mBundle)
		CFRelease(mBundle);
}
  	
//-------------------------------------------------------------------------------------------------------------
#pragma mark -
//-------------------------------------------------------------------------------------------------------------
void SCPatchController::AddPatch(CFStringRef bundleIdentifier,
								CFStringRef subPath,
								CFStringRef name)
{
	SCPatchRecord	patch(mBundle, bundleIdentifier, subPath, name);
	
	mPatchList.push_back(patch);
}

//-------------------------------------------------------------------------------------------------------------
// Inject the patches into running applications and start watching app launches   
OSErr SCPatchController::InstallPatches(void)
{
	OSErr 					procErr = noErr, err = noErr;
	EventTypeSpec			appSpec[]  = { { kEventClassApplication, kEventAppLaunched },
										   { kEventClassApplication, kEventAppTerminated } };
	
	// Start listening for messages from patches
	SetMessageRetryInterval(0.1);	// Try resending messages to patches every 1/10 second.
	SetMessageRetryLimit(600);		// Give up after 60 seconds.
	
	if((err = StartListening(mApplicationBundleIdentifier, false, true)) != noErr)
	{
		printf("SCPC: controller could not open port for listening\n");
	}
	
	if(err == noErr)
	{
		InstallApplicationEventHandler(NewEventHandlerUPP(ApplicationEventHandler),  2, appSpec, this, NULL);
		
	    ProcessSerialNumber psn = { 0, kNoProcess };
	    
	    while(!procErr)
	    {
	        if((procErr = GetNextProcess(&psn)) == noErr)
	        {
				if(!IsProcessPatched(&psn))
				{
	        		InstallPatchesInProcess(&psn);
				}
				else
				{
					ProcessInfoRec 			info;
					CFStringRef				str;
					Str255					pStr;
					    
				    info.processInfoLength = sizeof(ProcessInfoRec);
				    info.processName = pStr;
				    info.processAppSpec = nil;
					    
					if((err = GetProcessInformation(&psn, &info)) == noErr && 
					   (str = CFStringCreateWithPascalString(NULL, pStr, kCFStringEncodingMacRoman)) != NULL)
					{
						PatchNotification(&psn, info.processSignature, info.processType, str, info.processMode);
						CFRelease(str);
					}

				}
			}
		}
	}
	return err;
}

// Patch a single process (ignoring return value of ShouldPatchProcess)
mach_error_t SCPatchController::PatchProcess(ProcessSerialNumber *psn)
{
	return InstallPatchesInProcess(psn, true);
}

//-------------------------------------------------------------------------------------------------------------
// Get info and keep tabs on patched processes
Boolean SCPatchController::IsProcessPatched(ProcessSerialNumber *inPSN)
{
	SCPatchRecordIterator	iter;

	if(mPatchContextList.GetContext(inPSN))
	{
		// fprintf(stderr, "Process patched (found context)\n");
		return true;
	}
	
	for(iter = mPatchList.begin(); iter != mPatchList.end(); iter++)
	{
		OSErr	err;
		
		err = SendPing(iter->GetIdentifier(), inPSN);
		
		// If we get a destPortErr, no one's listening on the desired port.
		// If we get noErr or noResponseErr, there's someone listening.
		if(err == noErr || err == noResponseErr)
		{
			pid_t			 		pid;
			ProcessInfoRec 			info;
			Str255					pStr;
			    
		    info.processInfoLength = sizeof(ProcessInfoRec);
		    info.processName = pStr;
		    info.processAppSpec = nil;
			    
			if((err = GetProcessInformation(inPSN, &info)) == noErr && 
			   (err = GetProcessPID(inPSN, &pid)) == noErr)
			{
				mPatchContextList.NewContext(inPSN, pid, info.processSignature);
			}
			// fprintf(stderr, "Process patched (ping succeeded)\n");
			return true;
		}
		else
		{
			// fprintf(stderr, "Process not patched (ping failed)\n");
		}
	}
	
	return false;
}

//-------------------------------------------------------------------------------------------------------------
UInt32 SCPatchController::GetPatchFlags(ProcessSerialNumber *inPSN)
{
	SCPatchContext	*context = mPatchContextList.GetContext(inPSN);
	
	if(context)
		return context->flags;
	else
		return 0;
}

//-------------------------------------------------------------------------------------------------------------
void SCPatchController::SetPatchFlags(ProcessSerialNumber *inPSN, UInt32 flags, UInt32 whichFlags)
{
	SCPatchContext	*context = mPatchContextList.GetContext(inPSN);
	
	if(context)
		context->flags = (context->flags & ~whichFlags) | flags;
}

//-------------------------------------------------------------------------------------------------------------
OSErr SCPatchController::ForEachPatchedProcess(SCPatchIterationProc proc, void *data)
{
	SCPatchContextIterator	iter;
	OSErr					err = noErr;
	
	for(iter = mPatchContextList.begin(); iter != mPatchContextList.end(); iter++)
		if((err = proc(&iter->second.psn, iter->second.creator, iter->second.flags, data)) != noErr)
			break;
	return err;
}

//-------------------------------------------------------------------------------------------------------------
#pragma mark -
//-------------------------------------------------------------------------------------------------------------
OSErr SCPatchController::HandleMessage(const AppleEvent *theAE)
{
	OSErr				err = eventNotHandledErr;
	OSType				eventClass, eventID;
	Size				actualSize;
	DescType			actualType;
	ProcessSerialNumber	psn;

	if((err = AEGetAttributePtr(theAE, keyEventIDAttr, typeType, &actualType, &eventID, sizeof(OSType), &actualSize)) != noErr ||
	   (err = AEGetAttributePtr(theAE, keyEventClassAttr, typeType, &actualType, &eventClass, sizeof(OSType), &actualSize)) != noErr)
	{
		return err;
	}
	
	if(eventClass == kSCMessageClass && eventID == kSCPatchSuccess &&
	   (err = AEGetParamPtr(theAE, keyPSN, typeProcessSerialNumber, &actualType, &psn, sizeof(ProcessSerialNumber), &actualSize)) == noErr)
	{
		RecordPatchAndNotify(&psn, NULL);
	}
	
	// Return eventNotHandledErr even though we've done our 
	// thing so that the subclass gets a chance at it too.
	return eventNotHandledErr;
}

//-------------------------------------------------------------------------------------------------------------
pascal OSStatus SCPatchController::ApplicationEventHandler(EventHandlerCallRef handlerRef, 
														   EventRef event, void *userData)
{
	#pragma	unused (handlerRef)
	
	SCPatchController	*self = (SCPatchController *)userData;
	UInt32				kind = GetEventKind(event);
	OSStatus			err;
	ProcessSerialNumber	psn;
	
	err = GetEventParameter(event, kEventParamProcessID, typeProcessSerialNumber, 
						    NULL, sizeof(ProcessSerialNumber), NULL, &psn);
	
    fprintf(stderr, "ApplicationEventHandler called\n");
	if(err == noErr)
	{
		if(kind == kEventAppLaunched)
		{
			// fprintf(stderr, "SCPC: App launched\n");
			if(!self->IsProcessPatched(&psn))
				self->InstallPatchesInProcess(&psn);
		}
		else if(kind == kEventAppTerminated)
		{
			pid_t			 		pid;
			ProcessInfoRec 			info;
			CFStringRef				str;
			Str255					pStr;
			    
		    info.processInfoLength = sizeof(ProcessInfoRec);
		    info.processName = pStr;
		    info.processAppSpec = nil;
			    
			if((err = GetProcessInformation(&psn, &info)) == noErr && 
			   (err = GetProcessPID(&psn, &pid)) == noErr &&
			   (str = CFStringCreateWithPascalString(NULL, pStr, kCFStringEncodingMacRoman)) != NULL)
			{
				self->UnpatchNotification(&psn, info.processSignature, info.processType, str, info.processMode);
				CFRelease(str);
			}
			// fprintf(stderr, "SCPC: App died\n");
			self->mPatchContextList.DeleteContext(&psn);
		}
	}
	
	// Always pass the event down
	return eventNotHandledErr;
}

//-------------------------------------------------------------------------------------------------------------
mach_error_t SCPatchController::InstallPatchesInProcess(ProcessSerialNumber *psn, bool onDemand)
{
	SCPatchRecordIterator	iter;
	pid_t			 		pid;
	ProcessInfoRec 			info;
	CFStringRef				str;
	Str255					pStr;
	OSErr					err = noErr;
	mach_error_t			error = err_none;
	SCPatchLoaderParams		*params = NULL;
	    
    info.processInfoLength = sizeof(ProcessInfoRec);
    info.processName = pStr;
    info.processAppSpec = nil;
	    
	if((err = GetProcessInformation(psn, &info)) == noErr && 
	   (err = GetProcessPID(psn, &pid)) == noErr &&
	   (str = CFStringCreateWithPascalString(NULL, pStr, kCFStringEncodingMacRoman)) != NULL)
	{
		// fprintf(stderr, "SCPC: examining application (sig = %.4s, type = %.4s, flags = 0x%x, name = %s)\n", 
				// &info.processSignature, &info.processType, info.processMode, 
				// CFStringGetCStringPtr(str, kCFStringEncodingMacRoman));
				
		if((onDemand ||
			ShouldPatchProcess(psn, info.processSignature, info.processType, str, info.processMode)) &&
		   (params = (SCPatchLoaderParams *)malloc(sizeof(SCPatchLoaderParams))) != NULL)
		{
			// fprintf(stderr, "SCPC: patching application (sig = %.4s, type = %.4s, flags = 0x%x, name = %s)\n", 
					// &info.processSignature, &info.processType, info.processMode, 
					// CFStringGetCStringPtr(str, kCFStringEncodingMacRoman));
			
			params->version = 1;
			params->size = sizeof(SCPatchLoaderParams);
			params->patchCount = 0;
			
			for(iter = mPatchList.begin(); iter != mPatchList.end() && err == noErr; iter++)
				err = AddPatchToParams(iter->GetIdentifier(), iter->GetURL(), &params);
			
			if(err == noErr)
				error = InjectPatches(psn, params);
			free(params);
		}
		CFRelease(str);
	}

	if (error != err_none) return error;
	else return mac_err(err);
}

//-------------------------------------------------------------------------------------------------------------
OSErr SCPatchController::AddPatchToParams(CFStringRef bundleIdentifier, CFURLRef url, SCPatchLoaderParams **params)
{
	OSErr		err = err_couldnt_find_patch_bundle;
	CFStringRef	patchPath = NULL;
	size_t		newSize;
	
	if(bundleIdentifier == NULL || url == NULL || params == NULL)
		return paramErr;
	
#if 0
	if((patchPath = CFURLGetString(url)) != NULL)
		err = noErr;
#else
	CFBundleRef	patchBundle;
	CFURLRef	patchURL;

	// Find out where the patch lives
	if((patchBundle = CFBundleCreate(kCFAllocatorDefault, url)) != NULL &&
	   (patchURL = CFBundleCopyExecutableURL(patchBundle)) != NULL)
	{
		CFURLRef	absoluteURL;

		if((absoluteURL = CFURLCopyAbsoluteURL(patchURL)) != NULL)
		{
			CFRelease(patchURL);
			patchURL = absoluteURL;
		}
		
		if((patchPath = CFURLCopyFileSystemPath(patchURL, kCFURLPOSIXPathStyle)) != NULL)
			err = noErr;

		CFRelease(patchURL);
		CFRelease(patchBundle);
	}
#endif

	// Increase the size of params.  This allocation may have waste - CFStringGetLength * 2 is the max _possible_ size.
	if(!err)
	{
		newSize = (*params)->size + sizeof(SCPatchLoaderData) + CFStringGetLength(patchPath) * 2;
		
		if((*params = (SCPatchLoaderParams *)realloc(*params, newSize)) != NULL)
			(*params)->size = newSize;
		else
			err = ENOMEM;
	}
	
	// Fill in all the params
	if(!err)
	{
		HFSUniStr255	*bundleID = SCPatchGetHFSUniStrPointer(*params, (*params)->patchCount);
		char			*urlData =SCPatchGetStringPointer(*params, (*params)->patchCount);
		
		if((bundleID->length = CFStringGetLength(bundleIdentifier)) > 256)
		{
			fprintf(stderr, "SCPC: patch bundleIdentifier is too long.  It must be 256 chars or less.\n");
			err = err_couldnt_load_injection_bundle;
		}
		else
		{
			CFStringGetCharacters(bundleIdentifier, CFRangeMake(0, bundleID->length), bundleID->unicode);
			if(CFStringGetCString(patchPath, urlData, CFStringGetLength(patchPath) * 2, kCFStringEncodingUTF8))
				(*params)->patchCount++;
			else
				err = err_couldnt_load_injection_bundle;
		}
	}
#if 1
	if(patchPath)
		CFRelease(patchPath);
#endif

	return err;
}

//-------------------------------------------------------------------------------------------------------------
mach_error_t SCPatchController::InjectPatches(ProcessSerialNumber *psn, SCPatchLoaderParams *params)
{
	mach_error_t		err = err_none;
	
	//	Convert PSN to PID.
	pid_t pid;
	if(!err)
		err = mac_err(GetProcessPID(psn, &pid));
	
	// Fill in all the params
	if(!err)
	{
		// Make copies of the bundle identifiers and set up other params
		if((params->parentBundleID.length = CFStringGetLength(mApplicationBundleIdentifier)) > 256)
		{
			fprintf(stderr, "SCPC: app bundleIdentifier is too long.  It must be 256 chars or less.\n");
			err = err_couldnt_load_injection_bundle;
		}
		else
		{
			params->parent = pid;
			CFStringGetCharacters(mApplicationBundleIdentifier, CFRangeMake(0, params->parentBundleID.length), params->parentBundleID.unicode);
		}
	}
	
	// Find the loader
	CFURLRef loaderURL = NULL;
	if(!err)
		if((loaderURL = CFBundleCopyResourceURL(mBundle, CFSTR("SCPatchLoader"), CFSTR("bundle"), NULL)) == NULL)
			err = err_couldnt_find_injection_bundle;

#if 0
	// If we get "file not found", the user must have moved our application bundle.  Try finding it
	// with Launch Services and going on that way.
	if(err == err_couldnt_load_injection_bundle)  // REVISIT
	{
		CFURLRef	appURL, patchURL;
		CFStringRef	patchPath;
		
		if(LSFindApplicationForInfo(NULL, sApplicationBundleIdentifier, NULL, NULL, &appURL) == noErr)
		{
			if((patchURL = CFURLCreateCopyAppendingPathComponent(NULL, appURL, subPath, false)) != NULL &&
			   (patchPath = CFURLCopyFileSystemPath(patchURL, kCFURLPOSIXPathStyle)) != NULL)
			{
				err = PatchControllerLoadPatchForProcess(pid, bundleID, patchPath, true);
				CFRelease(patchPath);
				CFRelease(patchURL);
			}
			CFRelease(appURL);
		}
	}
#endif

	//	Create injection bundle instance.
	CFBundleRef injectionBundle = NULL;
	if(!err) 
		if((injectionBundle = CFBundleCreate(kCFAllocatorDefault, loaderURL)) == NULL)
			err = err_couldnt_load_injection_bundle;
	
	//	Load the thread code injection.
	void *injectionCode = NULL;
	if(!err) 
		if((injectionCode = CFBundleGetFunctionPointerForName(injectionBundle, CFSTR(INJECT_ENTRY_SYMBOL))) == NULL)
			err = err_couldnt_find_injectedThread_symbol;

	//	Inject the code.
	if(!err)
	{
		err = ::mach_inject((mach_inject_entry)injectionCode, params, params->size, pid, 0);
	}
	else
	{
		fprintf(stderr, "Patcher got error %d\n", err);
	}

	//	Clean up.
	if(loaderURL)
		CFRelease(loaderURL);
	if(injectionBundle)
		CFRelease(injectionBundle);
	
	return err;
}

//-------------------------------------------------------------------------------------------------------------
OSErr SCPatchController::RecordPatchAndNotify(ProcessSerialNumber *psn, ProcessInfoRec *info)
{
	OSErr					err = noErr;
	ProcessInfoRec 			localInfo;
	Str255					pStr;
	CFStringRef				str;
	pid_t			 		pid;
	
	// Make sure we haven't already added this process
	if(mPatchContextList.GetContext(psn))
		return noErr;

	// If no info was given, get it ourselves
	if(!info)
	{
		info = &localInfo;
	    info->processInfoLength = sizeof(ProcessInfoRec);
	    info->processName = pStr;
	    info->processAppSpec = nil;
		err = GetProcessInformation(psn, info);
	}

	// Now add it to our context list and let the subclass know	
	if(err == noErr && (err = GetProcessPID(psn, &pid)) == noErr &&
	   (str = CFStringCreateWithPascalString(NULL, info->processName, kCFStringEncodingMacRoman)) != NULL)
	{
		mPatchContextList.NewContext(psn, pid, info->processSignature);
		PatchNotification(psn, info->processSignature, info->processType, str, info->processMode);
		CFRelease(str);
	}
	return err;
}