'''
********************************************************************************
Name:		HighlightMnemsAndLocs.py
Author:		case b <cbarnes@accuvant.com>
Version:	1

[Description]

This script can be run from IDA's 'Script File' menu item or included in
idapythonrc.py. When run it will add or remove a menu item under
'Edit->Other->Highlight Special Lines' with a hot key of alt-h.

This script will highlight lines that contain jumps, calls and calls to
user named locations with special prefixes. Highlights will need to be 
reapplied if changes to location names occur and you would like the highlighting
to reflect this. This involves hitting Alt-h twice.


[Notes]

Prefix matching compares the prefix of the location with those in
self.__CONFIG['locs']*['Prefix'] and dumbly iterates over them. This means if
you have a prefix of 'pre' before a prefix of 'prepre' in the configuration,
the color of 'pre' will be applied to a location with the name of 'prepreLoc' 
and not the color of 'prepre'.

To have this script run at startup, place lines like the following in your
idapythonrc.py file:

from HighlightMnemsAndLocs import Highlighter
hl = Highlighter(); hl.Run()

This script has only been tested in IDA 6.3 on MacOS. YMMV.

Happy hunting. If you have any questions, comments, rants, etc. please send
them my way.

********************************************************************************
'''
import idaapi, idc, idautils

class Highlighter( object):
	'''Highlights lines with jumps or calls and applies special highlighting to
	calls with specially named locations.
	'''
	def __init__( self):
		self.__ERRORS = {
			'SUCCESS' 	: 0, # Everything's great
		}

		self.__CONFIG = {
			'HIGHLIGHT_CALLS'	: True,
			'HIGHLIGHT_LOCS'	: True,
			'HIGHLIGHT_JUMPS'	: True,
			'calls' :	{
						'Mnems': ['call'],
						'Color': 0xFF222244
					},
			'jumps' :	{
						'Mnems': ['jo','jno','js','jns','je','jz','jne','jnz','jb','jnae','jc',\
	 						'jnb','jae','jnc','jbe','jna','ja','jnbe','jl','jnge','jge',\
			 				'jnl','jle','jng','jg','jpe','jnp','jpo','jcxz','jecxz'],
						'Color' : 0xFF442222
					},
			'locs' :	[
						{
							'Prefix' : '____?',
							'Color': 0xFFDD6644
						},
						{
							'Prefix' : '__?',
							'Color': 0xFFCC5544
						},
						{
							'Prefix' : '_?',
							'Color': 0xFF994422
						}
					]
		}
		self.__HighLighted = False
	def DoHighLight( self, resetColors = False):
		from idautils import * # Ensure imports since we may been started from idapythonrc.py
		from idc import *
		for segment in Segments():
			for head in Heads( segment, SegEnd( segment)):
				if isCode( GetFlags( head)):
					mnemonic = GetMnem( head)
					if self.__CONFIG['HIGHLIGHT_CALLS'] and mnemonic in self.__CONFIG['calls']['Mnems']:
						op = GetOpnd( head, 0)
						opEA = LocByName( op)
						of = GetFlags( opEA)
						if resetColors:
							SetColor( head, CIC_ITEM, 0xFFFFFF)
						elif self.__CONFIG['HIGHLIGHT_LOCS'] and hasUserName( of):
							for locs in self.__CONFIG['locs']:
								if op.startswith( locs['Prefix']):
									if resetColors:
										SetColor( head, CIC_ITEM, 0xFFFFFF)
									else:
										SetColor( head, CIC_ITEM, locs['Color'])
									break
						else:
							SetColor( head, CIC_ITEM, self.__CONFIG['calls']['Color'])
					elif self.__CONFIG['HIGHLIGHT_JUMPS'] and mnemonic in self.__CONFIG['jumps']['Mnems']: 
						if resetColors:
							SetColor( head, CIC_ITEM, 0xFFFFFF)
						else:
							SetColor( head, CIC_ITEM, self.__CONFIG['jumps']['Color'])
	def MenuItemSelected( self):
		if self.__HighLighted:
			self.DoHighLight( True)
		else:
			self.DoHighLight()
		self.__HighLighted = not self.__HighLighted
	def AddMenuItemAndKey( self):
		return idaapi.add_menu_item( "Edit/Other/", "Highlight Special Lines", \
			'Alt-h', 0, self.MenuItemSelected, tuple())
	def RemoveMenuItemAndKey( self, menuItem):
		idaapi.del_menu_item( menuItem)
		del menuItem
	def Run( self):
		global HighlighterSet
		try:
			HighlighterSet
			self.RemoveMenuItemAndKey( HighlighterSet)
			del HighlighterSet
		except:
			HighlighterSet = self.AddMenuItemAndKey()
		return self.__ERRORS['SUCCESS']
	
if __name__ == "__main__":
	script = Highlighter()
	script.Run()