##########################################################################
# An IDAPython plugin that outputs function information for ScratchABit .py files
##########################################################################
import idc
import idaapi
import idautils

import os
import sys
import time
import collections


def PYBuild(sigfile=None):
    print "Building PY file, this may take a few minutes..."
    start = time.time()
    print ("Saving signatures to %s..." % sigfile),
    fp = open(sigfile, "wb")
    fp.write("from idc import *\n")

    for i in idautils.Strings():
        fp.write("MakeComm(0x%x,\"%s\")\n" % (i.ea, str(i).replace('\r','\\r').replace('\n', '\\n').replace('"','\\"')))
    for ea in idautils.Functions():
        func = idaapi.get_func(ea)
        fp.write("SetRegEx(0x%lx,\"T\",%d,2)\n" % (func.startEA,GetReg(func.startEA,"T")))
        fp.write("MakeFunction(0x%lx,4294967295)\nMakeNameEx(0x%lx,\"%s\",1)\n" % (func.startEA,func.startEA,GetFunctionName(ea)))
        for block in idaapi.FlowChart(func):
            fp.write("#Block 0x%x;0x%x;  block;type:%d\n" % (block.startEA, block.endEA,block.type))
    fp.close()
    end = time.time()
    print "Built PY file in %.2f seconds" % (end-start)

	
class PYPlugin(idaapi.plugin_t):
    flags = 0
    comment = "Function signature"
    help = ""
    wanted_name = "PY"
    wanted_hotkey = ""

    NAME = "exportScratchABit.py"

    def init(self):
        self.menu_context_produce = idaapi.add_menu_item("File/Produce file/", "ScratchABit Python file...", "", 0, self.py_produce, (True,))
        return idaapi.PLUGIN_KEEP

    def term(self):
        idaapi.del_menu_item(self.menu_context_produce)
        return None

    def run(self, arg):
        return None

    def py_script(self):
        idaapi.IDAPython_ExecScript(self.script, globals())

    def py_produce(self, arg):
        defau=idc.GetInputFile()
        fname = idc.AskFile(1, defau+".py", "Save signature file as")
        if fname:
            if '.' not in fname:
                fname += ".py"
            PYBuild(fname)
        return None

def PLUGIN_ENTRY():
    return PYPlugin()