import os, sys, time, string, re

from idaapi   import *
from idautils import *
from idc      import *

def get_register(address, opnum):
    op = GetOpnd(address, opnum)
    
    register = op.split("[")[-1].split("+")[0].split("-")[0].split("]")[0]
                                   
    return register

def convert_size(register):
    dwords = ["eax", "ebx", "ecx", "edx", "esp", "ebp", "esi", "edi"]
    words  = ["ax", "bx", "cx", "dx", "bp", "si", "di"]
    bytes  = ["al", "ah", "bl", "bh", "cl", "ch", "dl", "dh"]
    
    if register in dwords:
        return 4
    elif register in words:
        return 2
    elif register in bytes:
        return 1
    
    return False

def convert_type(typeid):
    optypes = \
        {
            1   :   'r',
            2   :   'p',
            3   :   'p',
            4   :   'o',
        }

    for k,v in optypes.iteritems():
        if k == typeid:
            return v
    
    return False

def write_config(fh, entries):
    for entry in entries:
        for parsed_entry in entry:
            print "%x,%c,%d,%s" % (parsed_entry["address"], parsed_entry["type"], parsed_entry["size"], parsed_entry["data"])
            fh.write("%x,%c,%d,%s\n" % (parsed_entry["address"], parsed_entry["type"], parsed_entry["size"], parsed_entry["data"]))
    
    fh.close()
    
    return True

def parse_tag(current, entries):
    # R = Register
    # O = Operand
    # M = Memory
    
    parsed_entries = []
    
    for entry in entries:
        parsed_entry = {}
        if re.search('R', entry, re.IGNORECASE):
            dreftype, register = entry.split(':')
            
            parsed_entry["address"] = current
            parsed_entry["type"] = 'r'
            parsed_entry["size"] = convert_size(register)
            parsed_entry["data"] = register.upper()
        elif re.search('O', entry, re.IGNORECASE):
            dreftype, opnum = entry.split(':')
            optype = convert_type(GetOpType(current, int(opnum)))
            size = convert_size(get_register(current, int(opnum)))
        
            parsed_entry["address"] = current
            parsed_entry["type"] = optype
            parsed_entry["size"] = size
            parsed_entry["data"] = str(int(opnum) + 1)
        elif re.search('M', entry, re.IGNORECASE):
            dreftype, memory_address = entry.split(':')
            
            parsed_entry["address"] = current
            parsed_entry["type"] = 'p'
            parsed_entry["size"] = 4
            parsed_entry["data"] = memory_address
        else:
            return False
        
        parsed_entries.append(parsed_entry)

    return parsed_entries
        
def find_tags():
    start = SegByName(".text")
    end = SegEnd(start)
    
    entries = []
    
    current = start
    while current != BADADDR:
        comment = Comment(current)

        if comment != None and re.search('\*\*LA.*', comment, re.IGNORECASE):
            entry = []
            entry = re.split('\*\*LA\s', comment, re.IGNORECASE)[-1].strip().split(',')
            
            entries.append(parse_tag(current, entry))
            
        current = NextHead(current, end)
    
    return entries

#entries address, type, size, opnum

filename = AskFile(1, "LiveAnalysis.conf", "Config output file")

if filename == None:
    print "[!] Cancelled"
else:
    fh = open(filename, "w")
    
    write_config(fh, find_tags())