Python MCP Obf mappings generator

import re, os, shutil, sys, fnmatch, csv
from pprint import pprint
from zipfile import ZipFile
from classfile import ClassFile

base = 'C:\\Users\\Lex\\Dropbox\\Public\\MinecraftForge\\'

def main():
    genMap('client.srg', 'minecraft.jar', '0')
    genMap('server.srg', 'minecraft_server.jar', '1')

def genMap(srg, jar, side):
    print 'Generating reflection files for %s' % ('Client' if side == '0' else 'Server')
    with open(srg) as f:
        srg = f.readlines()
    fCsv = readCSV('fields.csv', side)
    mCsv = readCSV('methods.csv', side)
    data = {}
    classes = {}

    for line in srg:
        pts = line.split()

        if pts[0] == 'CL:':
            #print ('Class %s -> %s' % (pts[1], pts[2]))
            data[pts[2] + '|' + pts[1]] = {}
            classes[pts[1]] = pts[2]

        elif pts[0] == 'FD:':
            ps1 = pts[1].rsplit('/', 1)
            ps2 = pts[2].rsplit('/', 1)
            cls = ps2[0] + '|' + ps1[0]
            if ps2[1] in fCsv:
                ps2[1] = fCsv[ps2[1]]
            #print 'Field: %s|%s -> %s -> %s' % (ps1[0], ps2[0], ps1[1], ps2[1])
            if not 'Fields' in data[cls]:
                data[cls]['Fields'] = {}
            data[cls]['Fields'][ps1[1]] = ps2[1]

        elif pts[0] == 'MD:':
            ps1 = pts[1].rsplit('/', 1)
            ps2 = pts[3].rsplit('/', 1)
            cls = ps2[0] + '|' + ps1[0]
            if '/' in ps1[0]:
                pts[2] = pts[2].replace(ps1[0].rsplit('/', 1)[0] + '/', '')
            if '/' in ps2[0]:
                pts[4] = pts[4].replace(ps2[0].rsplit('/', 1)[0] + '/', '')
            if ps2[1] in mCsv:
                ps2[1] = mCsv[ps2[1]]
            #print 'Method: %s|%s -> %s%s -> %s%s' % (ps1[0], ps2[0], ps1[1], pts[2], ps2[1], pts[4])
            if not 'Methods' in data[cls]:
                data[cls]['Methods'] = {}
            data[cls]['Methods'][ps1[1] + pts[2]] = ps2[1] + pts[4]

    parseJar(jar, classes, data)

    with open('%s%s_Hierarchy.txt' % (base, ('Client' if side == '0' else 'Server')), 'w') as f:
        writeHierarchy(f, makeClassHierarchy(data))

    with open('%s%s.map' % (base, ('Client' if side == '0' else 'Server')), 'w') as f:
        f.write('{\n');
        writeData(f, data, '    ')
        f.write('}')

def writeData(file, data, indent='', hasMore=True):
    keylist = data.keys()
    keylist.sort()
    for x in range(0, len(keylist)):
        key = keylist[x]
        value = data[key]
        ending = (',' if x < len(keylist) - 1 else '')
        if type(value) == dict:
            if len(value) == 0:
                file.write('%s"%s": {}%s\n' % (indent, key, ending))
            else:
                file.write('%s"%s": {\n' % (indent, key))
                writeData(file, value, indent + '    ', x < len(keylist))
                file.write('%s}%s\n' % (indent, ending))

        elif type(value) == str or type(value) == unicode:
            file.write('%s"%s": "%s"%s\n' % (indent, key, value, ending))

        elif type(value) == list:
            if len(value) == 0:
                file.write('%s"%s": []%s\n' % (indent, key, ending))
            elif len(value) == 1:
                file.write('%s"%s": [ "%s" ]%s\n' % (indent, key, value[0], ending))
            else:
                file.write('%s"%s" : [\n' % (indent, key))
                value.sort()
                for v in range(0, len(value)):
                    file.write('%s    "%s"%s\n' % (indent, value[v], (',' if v < len(value) - 1 else '')))
                file.write('%s]%s\n' % (indent, ending))
        else:
            print type(value)

def readCSV(filename, side):
    tCsv = csv.DictReader(open(filename, 'rb'))
    data = {}
    for row in tCsv:
        if row['side'] == side:
            data[row['searge']] = row['name']
    return data

def parseJar(jar, classes, data):
    zip = ZipFile(jar, 'r')
    for className in zip.namelist():
        if className.endswith('.class'):
            className = className[0:-6]
            if not className in classes:
                continue

            key = classes[className] + '|' + className

            classData = zip.read(className + '.class')

            c = ClassFile(classData)
            super = str(c.super_class.get_name())

            package1 = className.rsplit('/', 1)[0] + '/'
            package2 = classes[className].rsplit('/', 1)[0] + '/'

            if super != 'java/lang/Object':
                if super in classes:
                    data[key]['Super'] = classes[super].replace(package2, '') + '|' + super.replace(package1, '')
                else:
                    data[key]['Super'] = super

            for interface in c.interfaces:
                if not 'Interfaces' in data[key]:
                    data[key]['Interfaces'] = []
                intName = str(interface.get_name())
                if intName in classes:
                    data[key]['Interfaces'].append(classes[intName].replace(package2, '') + '|' + intName.replace(package1, ''))
                else:
                    data[key]['Interfaces'].append(intName)

    zip.close()

def makeClassHierarchy(data):
    classes = {}
    classes['java/lang/Enum'] = JavaNode('java/lang/Enum')
    classes['java/lang/Thread'] = JavaNode('java/lang/Thread')

    for key, value in data.items():
        classes[key.split('|')[0]] = JavaNode(key.split('|')[0])

    for key, value in data.items():
        cls = key.split('|')[0]
        package = cls.rsplit('/', 1)[0] + '/'

        super = 'java/lang/Object'
        if 'Super' in value:
            super = value['Super'].split('|')[0]
            if not '/' in super:
                super = package + super

        if super in classes:
            classes[super].children.append(classes[cls])
            classes[cls].parent = classes[super]

    for key, value in classes.items():
        if not value.parent == None:
            del classes[key]
    return classes

def writeHierarchy(file, classes, indent='', hasMore=True):
    if type(classes) == dict:
        keylist = classes.keys()
        keylist.sort()

        file.write('{\r\n')
        for x in range(0, len(keylist)):
            writeHierarchy(file, classes[keylist[x]], '    ', x < len(keylist) - 1)
        file.write('}')
    else:
        if len(classes.children) == 0:
            file.write('%s"%s": {}%s\r\n' % (indent, classes.name.replace('net/minecraft/src/', ''), (',' if hasMore else '')))
        else:
            file.write('%s"%s": {\r\n' % (indent, classes.name.replace('net/minecraft/src/', '')))
            classes.children.sort()
            for x in range(0, len(classes.children)):
                writeHierarchy(file, classes.children[x], '    ' + indent, x < len(classes.children) - 1)
            file.write('%s}%s\r\n' % (indent, ',' if hasMore else ''))

class JavaNode:
    def __init__(self, name):
        self.name = name
        self.parent = None
        self.children = []

if __name__ == '__main__':
    main()