x86 switch scanner

#include <idc.idc>
#define TRUE 1
#define FALSE 0
#define true 1
#define false 0
#define NULL 0

/*
    Project name: x86 switch scanner

    Author: fastman92

    Description: this script finds all switches and displays a list of cases for each.
*/


static IsSwitchJmp(ea)
{
    if(GetMnem(ea) != "jmp" || GetOpType(ea, 0) != o_mem)
        return false;

    auto op = GetOpnd(ea, 0);

    auto last = substr(op, strlen(op) - 3, -1);

    return last == "*4]";
}

static main()
{
    Message("Start of x86 switch scanner by fastman92\n");

    auto seg, loc;

    Message("========================================\n");

    seg = FirstSeg();   // Get address pointed by a first segment

    auto PointerTable;
    auto XlatTable;
    auto DefaultLocation;
    auto NumberOfCases;
    auto LowestValue;
    auto NumberOfCasesSet;

    auto locPrev;
    auto cmd;
    auto j;
    auto destination;


    while(seg != BADADDR )
    {
        Message("----------------------------------------\n");

        loc = SegStart(seg);

        Message("Segment %s\n", SegName(seg));

        // loc = 0x7AA341;

        while(loc != BADADDR && loc < SegEnd(seg))
        {
            if(isCode(GetFlags(loc)))
            {
                auto mnem = GetMnem(loc);

                if(IsSwitchJmp(loc))
                {
                    if(Word(loc) == 0x24FF)
                    {
                        LowestValue = 0;
                        DefaultLocation = 0;
                        NumberOfCases = 0;
                        NumberOfCasesSet = false;

                        PointerTable = GetOperandValue(loc, 0);
                        auto switchReg = (Byte(loc + 2) >> 3) & 7;

                        XlatTable = NULL;
                        // Message("Instruction 0x%X Disasm: %s\n", loc, GetDisasm(loc));

                        auto state = 0;

                        locPrev = loc;

                        auto i = 0;

                        while(true)
                        {
                            i++;
                            // Message("Address 0x%X lowest value: %d\n", locPrev, LowestValue);

                            locPrev = PrevHead(locPrev, 0);
                            mnem = GetMnem(locPrev);

                            if(DefaultLocation == 0
                                && mnem == "movzx"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                )
                            {
                                XlatTable = GetOperandValue(locPrev, 1);
                                switchReg = Byte(locPrev + 2) & 7;
                                // Message("Xlat = 0x%X\n", XlatTable);

                                // Message("Instruction 0x%X Disasm: %s\n", locPrev, GetDisasm(locPrev));
                            }
                            else if(mnem == "ja")
                            {
                                cmd = cmd = DecodeInstruction(locPrev);
                                DefaultLocation = cmd.Op0.addr;
                            }
                            else if(
                                !NumberOfCasesSet &&
                                mnem == "cmp"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                && GetOpType(locPrev, 1) == o_imm
                                )
                            {
                                NumberOfCasesSet = true;
                                NumberOfCases = GetOperandValue(locPrev, 1) + 1;
                            }
                            else if(
                                !NumberOfCasesSet &&
                                mnem == "cmp"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                && GetOpType(locPrev, 1) == o_reg
                                )
                            {
                                auto registerForComparison = GetOperandValue(locPrev, 1);

                                auto locPrevPrev = locPrev;


                                while(true)
                                {
                                    locPrevPrev = PrevHead(locPrevPrev, 0);
                                    auto xMnem = GetMnem(locPrevPrev);


                                    if(xMnem == "mov"
                                        && GetOpType(locPrevPrev, 0) == o_reg
                                        && GetOperandValue(locPrevPrev, 0) == registerForComparison
                                        && GetOpType(locPrevPrev, 1) == o_imm
                                        )
                                    {
                                        NumberOfCases = GetOperandValue(locPrevPrev, 1) + 1;
                                        break;
                                    }
                                    else if(xMnem == "pop"
                                        && GetOpType(locPrevPrev, 0) == o_reg
                                        && GetOperandValue(locPrevPrev, 0) == registerForComparison
                                        )
                                    {


                                        while(true)
                                        {
                                            if(GetMnem(locPrevPrev) == "push")
                                            {
                                                if(GetOpType(locPrevPrev, 0) == o_imm)
                                                    NumberOfCases = GetOperandValue(locPrevPrev, 0) + 1;
                                                else
                                                    NumberOfCases = -1;

                                                break;
                                            }

                                            locPrevPrev = PrevHead(locPrevPrev, 0);
                                        }

                                        break;
                                    }


                                }
                            }

                            else if(mnem == "add"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                && GetOpType(locPrev, 1) == o_imm
                                )
                            {
                                LowestValue = -GetOperandValue(locPrev, 1);
                                break;
                            }
                            else if(mnem == "sub"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                && GetOpType(locPrev, 1) == o_imm
                                )
                            {
                                LowestValue = GetOperandValue(locPrev, 1);
                                break;
                            }
                            else if(mnem == "dec"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                )
                            {
                                LowestValue = 1;
                                break;
                            }
                            else if(mnem == "inc"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                )
                            {
                                LowestValue = -1;
                                break;
                            }
                            else if(
                                mnem == "lea"
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                            )
                            {
                                if(GetOpType(locPrev, 1) == o_displ)
                                    LowestValue = -GetOperandValue(locPrev, 1);

                                break;
                            }
                            else if(
                                (mnem == "movzx" || mnem == "movsx" || mnem == "mov" || mnem == "and")
                                && GetOpType(locPrev, 0) == o_reg
                                && GetOperandValue(locPrev, 0) == switchReg
                                )
                                break;
                            else if(mnem == "call")
                                break;

                            // if(i > 3)
                            //  break;

                        }

                        Message("\nSwitch 0x%X:\n", loc);
                        Message("\tXLAT table: 0x%X\n", XlatTable);
                        Message("\tPointer table: 0x%X\n", PointerTable);
                        Message("\tDefault case: 0x%X\n", DefaultLocation);

                        Message("\tNumber of cases: %d\n", NumberOfCases);
                        Message("\tLowest case: %d\n", LowestValue);

                        Message("\tCases:\n");

                        if(XlatTable)
                        {
                            for(j = 0; j < NumberOfCases; j++)
                            {
                                destination = Dword(PointerTable + Byte(XlatTable + j) * 4);


                                if(destination != DefaultLocation)
                                    Message("\t\t%d - 0x%X\n", LowestValue + j, destination);

                                /*
                                if(destination == DefaultLocation)
                                    Message(" (default)");
                                */

                            }
                        }
                        else
                        {
                            for(j = 0; j < NumberOfCases; j++)
                            {
                                destination = Dword(PointerTable + j * 4);

                                if(destination != DefaultLocation)
                                    Message("\t\t%d - 0x%X\n", LowestValue + j, destination);

                                /*
                                if(destination == DefaultLocation)
                                    Message(" (default)");
                                */

                            }
                        }

                    }
                    else
                    {
                        Message("Don't know how to process Instruction 0x%X Disasm: %s\n", loc, GetDisasm(loc));
                    }

                }
            }

            // break;

            loc = NextHead(loc, BADADDR);
        }

        // break;

        seg = NextSeg(seg);     // get address of the next segment
    }
}