//http://iphone.fiveforty.net/wiki/index.php/Objective-C_Script_for_IDA

#include <idc.idc>

static fixSelectors()
{
	auto start;
	auto end;
	auto i;
	auto prefix;
	auto name;

	start = SegByBase(SegByName("__message_refs"));
    if (start==BADADDR) return;
	end = SegEnd(start);
	for(i = start; i < end; i = i + 4)
	{
		name = "sel_" + GetString(Dword(i), -1, ASCSTR_C);
		MakeNameEx(i, name, SN_NOWARN);
	}
}

static fixClassRefs()
{
	auto start;
	auto end;
	auto i;
	auto prefix;
	auto name;

	start = SegByBase(SegByName("__cls_refs"));
    if (start==BADADDR) return;
	end = SegEnd(start);
	for(i = start; i < end; i = i + 4)
	{
		name = "clsref_" + GetString(Dword(i), -1, ASCSTR_C);
		MakeNameEx(i, name, SN_NOWARN);
	}
}

static fixCFStrings()
{
	auto start;
	auto end;
	auto i;
	auto prefix;
	auto name;

	start = SegByBase(SegByName("__cfstring"));
    if (start==BADADDR) return;
	end = SegEnd(start);
	for(i = start; i < end; i = i + 16)
	{
		name = GetString(Dword(i + 8), -1, ASCSTR_C);
		MakeNameEx(i, "cfstr_" + name, SN_NOWARN);
		MakeRptCmt(i, name);
		//MakeDword(i);
		//MakeDword(i + 4);
		//MakeDword(i + 8);
		//OpOff(i + 8, 0, 0);
		//MakeDword(i + 12);
	}
}

static makeStructFromIvars(name, ivarAddr, instSize)
{
	auto struc;
	auto count;
	auto i;
	auto size;
	auto offset;
	auto varName;

	struc = AddStrucEx(-1, name, 0);
	count = Dword(ivarAddr);

	for(i = 0; i < count; i++)
	{
		varName = GetString(Dword(ivarAddr + 4 + i * 12), -1, ASCSTR_C);
		offset = Dword(ivarAddr + 4 + i * 12 + 8);

		if(i < count - 1)
		{
			size = Dword(ivarAddr + 4 + (i + 1) * 12 + 8) - offset;
		}
		else
		{
			size = instSize - offset;
		}

		AddStrucMember(struc, varName, offset, 0, 0, size);
	}
}

static fixClasses(meta)
{
	auto start;
	auto end;
	auto i;
	auto j;
	auto name;
	auto methName;
	auto meths;
	auto count;
	auto addr;
	auto codeAddr;
	auto segName;
	auto clsPrefix;
	auto methPrefix;
	auto methChar;

	if(meta == 1)
	{
		segName = "__meta_class";
		clsPrefix = "metadef_";
		methPrefix = "clsmeth_";
		methChar = "+";
	}
	else
	{
		segName = "__class";
		clsPrefix = "clsdef_";
		methPrefix = "instmeth_";
		methChar = "-";
	}

	start = SegByBase(SegByName(segName));
    if (start==BADADDR) return;
	end = SegEnd(start);

	for(i = start; i < end; i = i + 0x40)
	{
		name = GetString(Dword(i + 8), -1, ASCSTR_C);
		MakeNameEx(i, clsPrefix + name, SN_NOWARN);
		MakeNameEx(Dword(i + 28), methPrefix + name, SN_NOWARN);

		if(meta == 0)
		{
			MakeNameEx(Dword(i + 24), "ivars_" + name, SN_NOWARN);
			makeStructFromIvars(name, Dword(i + 24), Dword(i + 20));
		}

		meths = Dword(i + 28);
		count = Dword(meths + 4);
		for(j = 0; j < count; j = j + 1)
		{
			addr = (meths + 8 + j * 3 * 4);
			methName = GetString(Dword(addr), -1, ASCSTR_C);
			codeAddr = Dword(addr + 8);
			MakeNameEx(codeAddr, name + "_" + methName, SN_NOWARN);

			SetFunctionCmt(codeAddr, methChar + " [" + name + " " + methName + "]\n" +
				"Signature:   " + GetString(Dword(addr + 4), -1, ASCSTR_C), 0); 
		}
	}
}

static fixMethodList(name, meths, meta)
{
	auto j;
	auto count;
	auto addr;
	auto methName;
	auto codeAddr;
	auto methChar;
	auto prefix;

	if(meta == 1)
	{
		prefix = "catclsmeth_";
		methChar = "+";
	}
	else
	{
		prefix = "catinstmeth_";
		methChar = "-";
	}

	MakeNameEx(prefix + meths, name, SN_NOWARN);

	count = Dword(meths + 4);
	for(j = 0; j < count; j = j + 1)
	{
		addr = (meths + 8 + j * 3 * 4);
		methName = GetString(Dword(addr), -1, ASCSTR_C);
		codeAddr = Dword(addr + 8);
		MakeNameEx(codeAddr, name + "_" + methName, SN_NOWARN);

		SetFunctionCmt(codeAddr, methChar + " [" + name + " " + methName + "]\n" +
			"Signature:   " + GetString(Dword(addr + 4), -1, ASCSTR_C), 0); 
	}
}

static fixCategories()
{
	auto start;
	auto end;
	auto i;
	auto j;
	auto className;
	auto catName;
	auto methName;
	auto meths;
	auto count;
	auto addr;
	auto codeAddr;
	auto segName;
	auto instMeths;
	auto clsMeths;

	segName = "__category";

	start = SegByBase(SegByName(segName));
    if (start==BADADDR) return;
	end = SegEnd(start);

	for(i = start; i < end; i = i + 0x1C)
	{
		catName = GetString(Dword(i + 0), -1, ASCSTR_C);
		className = GetString(Dword(i + 4), -1, ASCSTR_C);

		MakeNameEx(i, "catdef_" + className + "__" + catName, SN_NOWARN);

		instMeths = Dword(i + 8);
		clsMeths = Dword(i + 12);

		fixMethodList(className + "_" + catName, instMeths, 0);
		fixMethodList(className + "_" + catName, clsMeths, 1);
	}
}

static main()
{
	Message("Fixing selectors...\n");
    fixSelectors();
    Message("Fixing class refs...\n");
	fixClassRefs();
    Message("Fixing CFStrings...\n");
	fixCFStrings();
    Message("Fixing classes 0...\n");
	fixClasses(0);
    Message("Fixing classes 1...\n");
	fixClasses(1);
    Message("Fixing categories...\n");
	fixCategories();
    Message("Done.\n");
}