make a basic python module

1 files changed, 219 insertions, 0 deletions

diff --git a/pokemontools/romstr.py b/pokemontools/romstr.py
new file mode 100644
index 0000000..69a4f2a
--- /dev/null
+++ b/pokemontools/romstr.py
@@ -0,0 +1,219 @@
+# -*- coding: utf-8 -*-
+
+import sys
+import os
+import time
+import datetime
+from ctypes import c_int8
+from copy import copy
+import json
+
+# New versions of json don't have read anymore.
+if not hasattr(json, "read"):
+    json.read = json.loads
+
+from labels import (
+    get_label_from_line,
+    get_address_from_line_comment,
+)
+
+relative_jumps = [0x38, 0x30, 0x20, 0x28, 0x18, 0xc3, 0xda, 0xc2, 0x32]
+relative_unconditional_jumps = [0xc3, 0x18]
+call_commands = [0xdc, 0xd4, 0xc4, 0xcc, 0xcd]
+end_08_scripts_with = [
+    0xe9, # jp hl
+    0xc9, # ret
+] # possibly also:
+    # 0xc3,  # jp
+    # 0xc18, # jr
+    # 0xda, 0xe9, 0xd2, 0xc2, 0xca, 0x38, 0x30, 0x20, 0x28, 0x18, 0xd8,
+    # 0xd0, 0xc0, 0xc8, 0xc9
+
+spacing = "\t"
+
+class RomStr(str):
+    """
+    Simple wrapper to prevent a giant rom from being shown on screen.
+    """
+
+    def __init__(self, *args, **kwargs):
+        if "labels" in kwargs.keys() and kwargs["labels"] == True:
+            self.load_labels()
+        str.__init__(self)
+
+    def __repr__(self):
+        """
+        Simplifies this object so that the output doesn't overflow stdout.
+        """
+        return "RomStr(too long)"
+
+    @classmethod
+    def load(cls, filename=None, crystal=True, red=False):
+        """
+        Load a ROM into a RomStr.
+        """
+        if   crystal and not red and not filename:
+            file_handler = open("../baserom.gbc", "r")
+        elif red and not crystal and not filename:
+            file_handler = open("../pokered-baserom.gbc", "r")
+        elif filename not in ["", None]:
+            file_handler = open(filename, "rb")
+        else:
+            raise Exception("not sure which rom to load?")
+        bytes = file_handler.read()
+        file_handler.close()
+        return RomStr(bytes)
+
+    def load_labels(self, filename="labels.json"):
+        """
+        Loads labels from labels.json.
+
+        (Or parses the source code file and
+        generates new labels.)
+        """
+        filename = os.path.join(os.path.dirname(__file__), filename)
+
+        # blank out the hash
+        self.labels = {}
+
+        # check if the labels file exists
+        file_existence = os.path.exists(filename)
+
+        generate_labels = False
+
+        # determine if the labels file needs to be regenerated
+        if file_existence:
+            modified = os.path.getmtime(filename)
+            modified = datetime.datetime.fromtimestamp(modified)
+            current  = datetime.datetime.fromtimestamp(time.time())
+
+            is_old = (current - modified) > datetime.timedelta(days=3)
+
+            if is_old:
+                generate_labels = True
+        else:
+            generate_labels = True
+
+        # scan the asm source code for labels
+        if generate_labels:
+            asm = open(os.path.join(os.path.dirname(__file__), "../main.asm"), "r").read().split("\n")
+
+            for line in asm:
+                label = get_label_from_line(line)
+
+                if label:
+                    address = get_address_from_line_comment(line)
+
+                    self.labels[address] = label
+
+            content = json.dumps(self.labels)
+            file_handler = open(filename, "w")
+            file_handler.write(content)
+            file_handler.close()
+
+        # load the labels from the file
+        self.labels = json.read(open(filename, "r").read())
+
+    def get_address_for(self, label):
+        """
+        Return the address of a label.
+
+        This is slow and could be improved dramatically.
+        """
+        label = str(label)
+        for address in self.labels.keys():
+            if self.labels[address] == label:
+                return address
+        return None
+
+    def length(self):
+        """
+        len(self)
+        """
+        return len(self)
+
+    def len(self):
+        """
+        len(self)
+        """
+        return self.length()
+
+    def interval(self, offset, length, strings=True, debug=True):
+        """
+        Return hex values for the rom starting at offset until offset+length.
+        """
+        returnable = []
+        for byte in self[offset:offset+length]:
+            if strings:
+                returnable.append(hex(ord(byte)))
+            else:
+                returnable.append(ord(byte))
+        return returnable
+
+    def until(self, offset, byte, strings=True, debug=False):
+        """
+        Return hex values from rom starting at offset until the given byte.
+        """
+        return self.interval(offset, self.find(chr(byte), offset) - offset, strings=strings)
+
+    def to_asm(self, address, end_address=None, size=None, max_size=0x4000, debug=None):
+        """
+        Disassemble ASM at some address.
+
+        This will stop disassembling when either the end_address or size is
+        met. Also, there's a maximum size that will be parsed, so that large
+        patches of data aren't parsed as code.
+        """
+        if type(address) in [str, unicode] and "0x" in address:
+            address = int(address, 16)
+
+        start_address = address
+
+        if start_address == None:
+            raise Exception, "address must be given"
+
+        if debug == None:
+            if not hasattr(self, "debug"):
+                debug = False
+            else:
+                debug = self.debug
+
+        # this is probably a terrible idea.. why am i doing this?
+        if size != None and max_size < size:
+            raise Exception, "max_size must be greater than or equal to size"
+        elif end_address != None and (end_address - start_address) > max_size:
+            raise Exception, "end_address is out of bounds"
+        elif end_address != None and size != None:
+            if (end_address - start_address) >= size:
+                size = end_address - start_address
+            else:
+                end_address = start_address + size
+        elif end_address == None and size != None:
+            end_address = start_address + size
+        elif end_address != None and size == None:
+            size = end_address - start_address
+
+        raise NotImplementedError("DisAsm was removed and never worked; hook up another disassembler please.")
+        #return DisAsm(start_address=start_address, end_address=end_address, size=size, max_size=max_size, debug=debug, rom=self)
+
+class AsmList(list):
+    """
+    Simple wrapper to prevent all asm lines from being shown on screen.
+    """
+
+    def length(self):
+        """
+        len(self)
+        """
+        return len(self)
+
+    def __repr__(self):
+        """
+        Simplifies this object so that the output doesn't overflow stdout.
+        """
+        return "AsmList(too long)"
+
+if __name__ == "__main__":
+    cryrom = RomStr(open("../pokecrystal.gbc", "r").read());
+    asm = cryrom.to_asm(sys.argv[1])
+    print asm