pokecrystal/extras/graph.py

#!/usr/bin/python
# author: Bryan Bishop <kanzure@gmail.com>
# date: 2012-06-20

import networkx as nx

from romstr import RomStr, DisAsm, \
    relative_jumps, call_commands, \
    relative_unconditional_jumps

class RomGraph(nx.DiGraph):
    """ Graphs various functions pointing to each other.

        TODO: Bank switches are nasty. They should be detected. Otherwise,
        functions will point to non-functions within the same bank. Another way
        to detect bankswitches is retroactively. By disassembling one function
        after another within the function banks, it can be roughly assumed that
        anything pointing to something else (within the same bank) is really
        actually a bankswitch. An even better method to handle bankswitches
        would be to just detect those situations in the asm (but I presently
        forget how bankswitches are performed in pokecrystal).
    """

    # some areas shouldn't be parsed as asm
    exclusions = []

    # where is the first function located?
    start_address = 0x150

    # and where is a good place to stop?
    end_address = 0x4000 * 0x01 # only do the first bank? sure..

    # where is the rom stored?
    rompath = "../baserom.gbc"

    def __init__(self, rom=None, **kwargs):
        """ Loads and parses the ROM into a function graph.
        """
        # continue the initialization
        nx.DiGraph.__init__(self, **kwargs)

        # load the graph
        if rom == None:
            self.load_rom()
        else:
            self.rom = rom

        # start parsing the ROM
        self.parse()

    def load_rom(self):
        """ Creates a RomStr from rompath.
        """
        file_handler = open(self.rompath, "r")
        self.rom = RomStr(file_handler.read())
        file_handler.close()

    def parse(self):
        """ Parses the ROM starting with the first function address. Each
            function is disassembled and parsed to find where else it leads to.
        """
        functions = {}

        address = self.start_address

        other_addresses = set()

        count = 0

        while True:
            if count > 100:
                break

            if address < self.end_address and address not in functions.keys():
                # address is okay to parse at, keep going
                pass
            elif len(other_addresses) > 0:
                # parse some other address possibly in a remote bank
                address = other_addresses.pop()
            else:
                # no more addresses detected- exit loop
                break

            # parse the asm
            func = self.rom.to_asm(address)

            # store this parsed function
            functions[address] = func

            # where does this function jump to?
            used_addresses = set(func.used_addresses())

            # add this information to the graph
            for used_address in used_addresses:
                # only add this remote address if it's not yet parsed
                if used_address not in functions.keys():
                    other_addresses.update([used_address])

                # add this other address to the graph
                self.add_node(used_address)

                # add this as an edge between the two nodes
                self.add_edge(address, used_address)

            # setup the next function to be parsed
            address = func.last_address

            count += 1

        self.functions = functions

    def pretty_printer(self):
        """ Shows some text output describing which nodes point to which other
            nodes.
        """
        print self.edges()

    def to_d3(self):
        """ Exports to d3.js because we're gangster like that.
        """
        import networkx.readwrite.json_graph as json_graph
        content = json_graph.dumps(self)
        fh = open("graphs.json", "w")
        fh.write(content)
        fh.close()

class RedGraph(RomGraph):
    """ Not implemented. Go away.
    """

    rompath = "../pokered-baserom.gbc"

class CryGraph(RomGraph):
    exclusions = [
        [0x000, 0x149],
    ]

    rompath = "../baserom.gbc"

if __name__ == "__main__":
    crygraph = CryGraph()
    crygraph.pretty_printer()
    crygraph.to_d3()