/*
 * scrabble.idc, Finds refactorable code parts that could be used during exploitation
 * 
 * For detailed information about the usage of refactorable code within exploits
 * and based exploitation methods. Refer to [ http://www.tty64.org/doc/smackthestack.txt ] 
 * 
 * - izik@tty64.org
 */

#include <idc.idc>

//
// Configuration options
//
//  - FIND_ALL       : Finds all of the above
//  - FIND_RET       : Finds RET instruction
//  - FIND_POP_RET   : Finds POPs follow by RET sequance
//  - FIND_CALL_REGS : Finds CALLs to register instruction
//  - FIND_JMP_REGS  : Finds JMPs to register instruction
//

//

//

#define FIND_ALL 1

// #define FIND_RET 1
// #define FIND_POP_RET 1
// #define FIND_CALL_REGS 1
// #define FIND_JMP_REGS 1
//

//

#ifdef FIND_ALL
  #define FIND_RET 1
  #define FIND_POP_RET 1
  #define FIND_CALL_REGS 1
  #define FIND_JMP_REGS 1
#endif

//
// Internal definitions
//

#define SCRABBLE(addr,pattern) FindBinary(addr, SEARCH_DOWN|SEARCH_NEXT, pattern);

//
// Scrabble patterns
//

#define JMP_EAX_INST "FF E0"
#define JMP_EBX_INST "FF E3" 
#define JMP_ECX_INST "FF E1"
#define JMP_EDX_INST "FF E2"
#define JMP_ESI_INST "FF E6"
#define JMP_EDI_INST "FF E7"
#define JMP_EBP_INST "FF E5"
#define JMP_ESP_INST "FF E4"

#define CALL_EAX_INST "FF D0"
#define CALL_EBX_INST "FF D3"
#define CALL_ECX_INST "FF D1"
#define CALL_EDX_INST "FF D2"
#define CALL_ESI_INST "FF D6"
#define CALL_EDI_INST "FF D7"
#define CALL_EBP_INST "FF D5"
#define CALL_ESP_INST "FF D4"

#define RET_INST "C3"

#define POP_EAX_INST "58"
#define POP_EBX_INST "5B"
#define POP_ECX_INST "59"
#define POP_EDX_INST "5A"
#define POP_ESI_INST "5E"
#define POP_EDI_INST "5F"
#define POP_EBP_INST "5D"
#define POP_ESP_INST "5C"

//

/*
 * strtok, improvised strtok() a-like function
 * * buffer, given buffer
 * * token, given token
 */

static strtok(buffer, token) {

	auto i;

	for (i = strlen(buffer) - 1; i > 0; i--) {

		if (substr(buffer, i, i+1) == token) {
			break;
		}
	}

	return substr(buffer, 0, i+1);
}

/*
 * find_single_binary_pattern, Finds a single binary pattern
 * * fp, given file pointer (for logging)
 * * binary_pattern, given binary pattern to look for
 * * pattern_desc, given pattern description (for logging)
 * * 
 */

static find_single_binary_pattern(fp, binary_pattern, pattern_desc) {

	auto retval, addr, total_patterns, i;
		
	// Append pattern to log file

	fprintf(fp, "%s:\n", pattern_desc);

	for (i = strlen(pattern_desc) + 1; i > 0; i--) {

		fputc('-', fp);
	}

	fputc('\n', fp);
	fputc('\n', fp);

	// Look for the binary pattern

	addr = MinEA();
	retval = 1;
	total_patterns = 0;

	do {

		retval = SCRABBLE(addr, binary_pattern);

		if (retval != BADADDR) {

			// Found!, log the result!

			fprintf(fp, "%02x %02x %02x %02x : %a | %s\n", 
				Byte(retval), Byte(retval+1), Byte(retval+2), Byte(retval+3), 
					retval, GetDisasm(retval));

			addr = retval;

			total_patterns++;

		}

	} while (retval != BADADDR); 
	
	// Update log file with results

	if (total_patterns == 0) {

		fprintf(fp, "* None!\n\n");
	
	} else {

		fprintf(fp, "\n* Found %d occurrences of %s\n\n", 
			total_patterns , pattern_desc);
	}

	return total_patterns;
}

/*
 * find_pair_binary_pattern, Finds a pair of binary patterns
 * * fp, given file pointer (for logging)
 * * first_binary_pattern, given leading binary pattern to look for
 * * last_binary_pattern, given least binary pattern to look for
 * * pair_desc, given pair description (for logging)
 * * 
 */

static find_pair_binary_pattern(fp, first_binary_pattern, last_binary_pattern, pattern_desc) {

	auto mix;

	mix = form("%s %s", first_binary_pattern, last_binary_pattern);

	return find_single_binary_pattern(fp, mix, pattern_desc);
}

//

static main() {

	auto fp, log_filename;

	// Wait for IDA to finish its business first.

	Batch(0);
	Wait();

	// Create a log file	

	log_filename = form("%sscrabble_%stxt",
		strtok(GetInputFilePath(), "\\"), strtok(GetInputFile(), "."));

	fp = fopen(log_filename, "a+t");

	// We've got file pointer?

	if (fp == 0) {

		Message("ERROR: Unable to open [%s] for writing\n", log_filename);

		return 0;
	}

	// Lets play some scrabble!

	Message("Playing scrabble ... ");
	
	#ifdef FIND_RET 

	 	find_single_binary_pattern(fp, RET_INST, "RET instruction");

	#endif

	#ifdef FIND_JMP_REGS 

		find_single_binary_pattern(fp, JMP_EAX_INST, "JMP to EAX register");
		find_single_binary_pattern(fp, JMP_EBX_INST, "JMP to EBX register");
		find_single_binary_pattern(fp, JMP_ECX_INST, "JMP to ECX register");
		find_single_binary_pattern(fp, JMP_EDX_INST, "JMP to EDX register");
		find_single_binary_pattern(fp, JMP_ESI_INST, "JMP to ESI register");
		find_single_binary_pattern(fp, JMP_EDI_INST, "JMP to EDI register");
		find_single_binary_pattern(fp, JMP_EBP_INST, "JMP to EBP register");
		find_single_binary_pattern(fp, JMP_ESP_INST, "JMP to ESP register");

	#endif

	#ifdef FIND_CALL_REGS 

		find_single_binary_pattern(fp, CALL_EAX_INST, "CALL to EAX register");
		find_single_binary_pattern(fp, CALL_EBX_INST, "CALL to EBX register");
		find_single_binary_pattern(fp, CALL_ECX_INST, "CALL to ECX register");
		find_single_binary_pattern(fp, CALL_EDX_INST, "CALL to EDX register");
		find_single_binary_pattern(fp, CALL_ESI_INST, "CALL to ESI register");
		find_single_binary_pattern(fp, CALL_EDI_INST, "CALL to EDI register");
		find_single_binary_pattern(fp, CALL_EBP_INST, "CALL to EBP register");
		find_single_binary_pattern(fp, CALL_ESP_INST, "CALL to ESP register");

	#endif

	#ifdef FIND_POP_RET   

		find_pair_binary_pattern(fp, POP_EAX_INST, RET_INST, "POP into EAX followed by RET");
		find_pair_binary_pattern(fp, POP_EBX_INST, RET_INST, "POP into EBX followed by RET");
		find_pair_binary_pattern(fp, POP_ECX_INST, RET_INST, "POP into ECX followed by RET");
		find_pair_binary_pattern(fp, POP_EDX_INST, RET_INST, "POP into EDX followed by RET");
		find_pair_binary_pattern(fp, POP_ESI_INST, RET_INST, "POP into ESI followed by RET");
		find_pair_binary_pattern(fp, POP_EDI_INST, RET_INST, "POP into EDI followed by RET");
		find_pair_binary_pattern(fp, POP_EBP_INST, RET_INST, "POP into EBP followed by RET");
		find_pair_binary_pattern(fp, POP_ESP_INST, RET_INST, "POP into ESP followed by RET");		

	#endif

	// Close the log file and call it a day

	fclose(fp);

	Message("Done!\n");

	Message("Log saved [%s]\n", log_filename);
	
	return 1;
}