import os
import sys
import binascii
import struct
import re

def get_data(filename):
    totalbytes = os.path.getsize(filename)
    infile = open(filename, 'rb')
    totalfiledata = infile.read(totalbytes)
    return totalfiledata

opcodes = [
            [0x1,'retn',7],[0x2,'jmp',7],[0x3,'jmpcond',7],[0x4,'text',7],[0x5,'0x5',7],[0x6,'0x6',7],[0x7,'0x7',7],
            [0x8,'0x8',7],[0x9,'0x9',7],[0xa,'0xa',7],[0xb,'0xb',7],[0xc,'0xc',7],[0xd,'0xd',7],[0xe,'0xe',7],
            [0xf,'0xf',7],[0x10,'0x10',7],[0x11,'0x11',7],[0x12,'0x12',7],[0x13,'0x13',7],[0x14,'0x14',7],[0x15,'0x15',7],
            [0x16,'0x16',7],[0x17,'0x17',7],[0x18,'0x18',7],[0x19,'0x19',7],[0x1a,'0x1a',7],[0x1b,'0x1b',7],[0x1c,'0x1c',7],
            [0x1d,'0x1d',7],[0x1e,'0x1e',7],[0x1f,'0x1f',7],[0x20,'0x20',7],[0x21,'0x21',7],[0x22,'0x22',7],[0x23,'0x23',7],
            [0x24,'0x24',7],[0x25,'0x25',7],[0x26,'0x26',7],[0x27,'0x27',7],[0x28,'0x28',7],[0x29,'0x29',7],[0x2a,'0x2a',7],
            [0x2b,'0x2b',7],[0x2c,'0x2c',7],[0x2d,'0x2d',7],[0x2e,'0x2e',7],[0x2f,'0x2f',7],[0x30,'0x30',7],[0x31,'0x31',7],
            [0x32,'0x32',7],[0x33,'0x33',7],[0x34,'0x34',7],[0x35,'0x35',7],[0x36,'0x36',7],[0x37,'0x37',7],[0x38,'0x38',7],
            [0x39,'0x39',7],[0x3a,'0x3a',7],[0x3b,'0x3b',7],[0x3c,'shop',7],[0x3d,'0x3d',7],[0x3e,'0x3e',7],[0x3f,'0x3f',7],
            [0x40,'0x40',7],[0x41,'0x41',7],[0x42,'0x42',7],[0x43,'0x43',7],[0x44,'0x44',7],[0x45,'0x45',7],[0x46,'0x46',7],
            [0x47,'0x47',7],[0x48,'0x48',7],[0x49,'0x49',7],[0x4a,'0x4a',7],[0x4b,'0x4b',7],[0x4c,'0x4c',7],[0x4d,'0x4d',7],
            [0x4e,'0x4e',7],[0x4f,'char',7],[0x50,'0x50',7],[0x51,'0x51',7],[0x52,'0x52',7],[0x53,'0x53',7],[0x54,'0x54',7],
            [0x55,'0x55',7],[0x56,'0x56',7],[0x57,'0x57',7],[0x58,'0x58',7],[0x59,'0x59',7],[0x5a,'0x5a',7],[0x5b,'0x5b',7],
            [0x5c,'0x5c',7],[0x5d,'0x5d',7],[0x5e,'0x5e',7],[0x5f,'0x5f',7],[0x60,'0x60',7],[0x61,'0x61',7],[0x62,'0x62',7],
            [0x63,'0x63',7],[0x64,'0x64',7],[0x65,'0x65',7],[0x66,'0x66',7],[0x67,'0x67',7],[0x68,'0x68',7],[0x69,'0x69',7],
            [0x6a,'0x6a',7],[0x6b,'0x6b',7],[0x6c,'0x6c',7],[0x6d,'0x6d',7],[0x6e,'0x6e',7],[0x6f,'0x6f',7],[0x70,'0x70',7],
            [0x71,'0x71',7],[0x72,'0x72',7],[0x73,'0x73',7],[0x74,'0x74',7],[0x75,'0x75',7],[0x76,'0x76',7],[0x77,'0x77',7],
            [0x78,'video',7],[0x79,'0x79',7],[0x7a,'0x7a',7],[0x7b,'0x7b',7],[0x7c,'0x7c',7],[0x7d,'0x7d',7],[0x7e,'0x7e',7],
            [0x7f,'0x7f',7],[0x80,'0x80',7],[0x81,'0x81',7],[0x82,'0x82',7],[0x83,'0x83',7],[0x84,'0x84',7],[0x85,'0x85',7],
            [0x86,'0x86',7],[0x87,'0x87',7],
          ]

def extract(filename):
    filedata = get_data(filename)
    table1 = struct.unpack('<H',filedata[0x1c:0x1e])[0] * 0x10
    table2 = struct.unpack('<H',filedata[0x1e:0x20])[0] * 0x10
    table3 = struct.unpack('<H',filedata[0x20:0x22])[0] * 0x10
    table4 = struct.unpack('<H',filedata[0x22:0x24])[0] * 0x10
    table5 = struct.unpack('<H',filedata[0x2c:0x2e])[0] * 0x10
    table6 = 0x1000
    table7 = struct.unpack('<H',filedata[0x24:0x26])[0] * 0x20
    table8 = struct.unpack('<H',filedata[0x26:0x28])[0] * 0x2
    maxlength = struct.unpack('<H',filedata[0x28:0x2a])[0] * 0x20
    table10 = struct.unpack('<H',filedata[0x2a:0x2c])[0]
    outfiledata = ''
    opcodeblockstart = 0x40 + table1 + table2 + table3 + table4 + table5 + table6 + table7 + table8
    i = opcodeblockstart
    blocksmoved = 0
    outfile = open(os.path.splitext(filename)[0] + '.data','wb')

    outfiledata += filename + '\n'
    lastopcode = 0
    firstrun = True
    functioncount = 0
    functions = []
    functablelen = table8
    for a in range(i-functablelen,i,0x2):
        functions.append(struct.unpack('<H',filedata[a:a+0x2])[0])

    while i != opcodeblockstart+maxlength: # loop while not end of opcode block
        if functions[functioncount] == ((i-opcodeblockstart) / 0x20):
            if firstrun == True:
                firstrun = False
            else:
                outfiledata += '}\n\n'
            try:
                outfiledata += 'function_%x --- length: %x\n{\n' % (blocksmoved,(functions[functioncount+1]-blocksmoved))
                functioncount += 1
            except Exception:
                outfiledata += 'function_%x --- length: %x\n{\n' % (blocksmoved,(maxlength/0x20)-blocksmoved)

        c = 0
        opcode = struct.unpack('<I',filedata[i:i+0x4])[0] # first dword is always our opcode
        while c != len(opcodes):
            if opcodes[c][0] == opcode: # is this a match to one of our opcodes
                variables = opcodes[c][2]
                name = ('%x' % (blocksmoved)) + ' { ' + opcodes[c][1] # the text to replace the opcode with
                blocksmoved += 1
                break;
            if c+1 == len(opcodes):
                print 'Unknown opcode: %x at loc: %x' % (opcode,i)
                sys.exit()
            c += 1

        outfiledata += name + ' '
        for a in range(i+0x4,i+0x4+(variables*4),0x4): # run through each variable for the opcode
            var = struct.unpack('<I',filedata[a:a+0x4])[0]
            if opcode == 0x4 and a == i+0x4+0x4*4: # get string for 0x4 opcode
                outfiledata += '"' + filedata[opcodeblockstart+maxlength+var:filedata.find('\x00',opcodeblockstart+maxlength+var,len(filedata))] + '" '
            elif opcode == 0x3c and a == i+0x4+0x4*4: # get string for 0x3c opcode
                outfiledata += '"' + filedata[opcodeblockstart+maxlength+var:filedata.find('\x00',opcodeblockstart+maxlength+var,len(filedata))] + '" '
            elif opcode == 0x21 and a == i+0x4+0x4*4: # get string for 0x21 opcode
                outfiledata += '"' + filedata[opcodeblockstart+maxlength+var:filedata.find('\x00',opcodeblockstart+maxlength+var,len(filedata))] + '" '
            elif opcode == 0x78 and a == i+0x4+0x4*4: # get string for 0x78 opcode
                outfiledata += '"' + filedata[opcodeblockstart+maxlength+var:filedata.find('\x00',opcodeblockstart+maxlength+var,len(filedata))] + '" '
            elif opcode == 0x4f and a == i+0x4+0x4*4: # get string for 0x4f opcode
                outfiledata += '"' + filedata[opcodeblockstart+maxlength+var:filedata.find('\x00',opcodeblockstart+maxlength+var,len(filedata))] + '" '
            else:
                outfiledata += ('%0x' % var) + ' ' # not a string, so just add the variable

        outfiledata += ' }\n'
        lastopcode = opcode
        i += 0x20
        if i == opcodeblockstart+maxlength:
            outfiledata += '}'
        outfile.write(outfiledata) # write the opcode out to our data file
        outfiledata = ''

        blocksmoved2 = 0

def insert(filename):
    filedata = get_data(filename)
    newopcodes = filedata.split('\n') # split our data by lines for each opcode
    newstrings = ''
    newfiledata = ''
    oldfilename = newopcodes[0]
    newopcodes.remove(newopcodes[0])
    b = 0
    c = 0
    blocksmoved = 0
    variables = 0
    totallinesadded = 0
    linesadded = []
    blocksmoved = 0
    positionsofinsert = []
    positionsofremove = []
    lastorigline = 0
    firstrun = True

    for i in range(0,len(newopcodes)):
        thisline = newopcodes[i] # get our new opcode
        if (thisline != '') and (thisline.find('function_') == -1) and (thisline != '{') and (thisline != '}'): # make sure it's not empty or just a function label
            if thisline.find('"') != -1: # does this line have a string in it. Could match the opcodes against all known text ones, do that if " causes problems
                strings = thisline.find('"')
                stringe = thisline.find('"',strings+1)
                string = thisline[strings+1:stringe]
                string += '\x00'
                if newstrings.find(string) == -1:
                    newstrings += string
                    thisline = thisline[:strings] + 'TEXT ' + thisline[stringe+2:] # replace the string with a temp string, this is so it doesn't break the below split
                else:
                    newstringpos = newstrings.find(string)
                    thisline = thisline[:strings] + 'DUPE ' + thisline[stringe+2:] # replace the string with a temp string, this is so it doesn't break the below split

            newvariables = thisline.split(' ') # split the line into variables
            isorigline = ''
            if newvariables[1] == '{':
                isorigline = newvariables[0]
                newvariables.remove(isorigline)

            newvariables.remove('{') # remove excess
            newvariables.remove('}')
            newvariables.remove('')
            a = 0
            f = 0
            d = 0
            diffcount = 0
            if isorigline == '':
                positionsofinsert.append(lastorigline)
            else:
                isorigline = int(isorigline,16)
                if firstrun == True:
                    lastorigline = isorigline
                    firstrun = False
                elif lastorigline == isorigline:
                    print 'ERROR WITH LINE COUNT.\nOUTPUT WILL BE BROKEN\nReason could be that your edits are formatted badly\nor an unknown opcode.'
                elif (lastorigline != isorigline):
                    while lastorigline+1 < isorigline:
                        lastorigline += 1
                        positionsofremove.append(isorigline)
                lastorigline = isorigline

            while a != len(newvariables):
                if newvariables[f] != '{' and newvariables[f] != '}' and newvariables[f] != '':
                    if a == 0: # first is always our opcode
                        d = 0
                        while d != len(opcodes):
                            if opcodes[d][1] == newvariables[f]: #does this match on of our opcodes
                                newfiledata += struct.pack('<I',opcodes[d][0]) # pack the opcode into our new filedata
                                break;
                            if d+1 == len(opcodes):
                                print 'Unknown opcode: %x in line: %s' % (opcode,f)
                                sys.exit()
                            d += 1
                    else:
                        if newvariables[f] == 'TEXT': # if we had a string as this variable, replace with its pointer (current length of our new strings)
                            newfiledata += struct.pack('<I',len(newstrings)-(len(string)))
                            newvariables[f+1] = str(hex(len(string)-1))
                        elif newvariables[f] == 'DUPE': # if we had a string as this variable, we already have this string in the file, point to it
                            newfiledata += struct.pack('<I',newstringpos)
                            newvariables[f+1] = str(hex(len(string)-1))
                            newstringpos = 0
                        else: # not a string, so just pack the variable again
                            newvariables[f] = int(newvariables[f],16)
                            newfiledata += struct.pack('<I',newvariables[f])
                    a += 1
                f += 1
            string = ''
            f = 0
            d = 0
            c = 0
            a = 0
    blocksmoved = 0
    variables = 0
    s = 0
    tmpfix = 0
    tmpfix2 = 0

    if len(positionsofinsert) > 0 or len(positionsofremove) > 0: #fix all our 0x2 and 0x3 jumps to match the amount removed/added
        while s != len(newfiledata):
            thisval = struct.unpack('<I',newfiledata[s:s+0x4])[0]
            if thisval == 0x2:
                fixme = struct.unpack('<I',newfiledata[s+0x4:s+0x8])[0]
                tmpfix = fixme
                if len(positionsofinsert) > 0:
                    for g in range(0,len(positionsofinsert)):
                        if tmpfix >= positionsofinsert[g]:
                            fixme += 0x1
                if len(positionsofremove) > 0:
                    for g in range(0,len(positionsofremove)):
                        if tmpfix >= positionsofremove[g]:
                            fixme -= 0x1
                newfiledata = newfiledata[:s+0x4] + struct.pack('<I',fixme) + newfiledata[s+0x8:]

            elif thisval == 0x3:
                fixme = struct.unpack('<I',newfiledata[s+0xc:s+0x10])[0]
                tmpfix = fixme
                if len(positionsofinsert) > 0:
                    for g in range(0,len(positionsofinsert)):
                        if tmpfix >= positionsofinsert[g]:
                            fixme += 0x1
                if len(positionsofremove) > 0:
                    for g in range(0,len(positionsofremove)):
                        if tmpfix >= positionsofremove[g]:
                            fixme -= 0x1
                newfiledata = newfiledata[:s+0xc] + struct.pack('<I',fixme) + newfiledata[s+0x10:]

                fixme = struct.unpack('<I',newfiledata[s+0x10:s+0x14])[0]
                tmpfix = fixme
                if len(positionsofinsert) > 0:
                    for g in range(0,len(positionsofinsert)):
                        if tmpfix >= positionsofinsert[g]:
                            fixme += 0x1
                if len(positionsofremove) > 0:
                    for g in range(0,len(positionsofremove)):
                        if tmpfix >= positionsofremove[g]:
                            fixme -= 0x1
                newfiledata = newfiledata[:s+0x10] + struct.pack('<I',fixme) + newfiledata[s+0x14:]

            s += 0x20
            blocksmoved += 1

    filedata1 = get_data(oldfilename)
    table1 = struct.unpack('<H',filedata1[0x1c:0x1e])[0] * 0x10
    table2 = struct.unpack('<H',filedata1[0x1e:0x20])[0] * 0x10
    table3 = struct.unpack('<H',filedata1[0x20:0x22])[0] * 0x10
    table4 = struct.unpack('<H',filedata1[0x22:0x24])[0] * 0x10
    table5 = struct.unpack('<H',filedata1[0x2c:0x2e])[0] * 0x10
    table6 = 0x1000
    table7 = struct.unpack('<H',filedata1[0x24:0x26])[0] * 0x20
    table8 = struct.unpack('<H',filedata1[0x26:0x28])[0] * 0x2
    opcodeblockstart = 0x40 + table1 + table2 + table3 + table4 + table5 + table6 + table7 + table8

    newfiledata1 = filedata1[:opcodeblockstart-(struct.unpack('<H',filedata1[0x26:0x28])[0]*0x2)] # this gets original data up to the point we edit
    newfiledata2 = filedata1[opcodeblockstart-(struct.unpack('<H',filedata1[0x26:0x28])[0]*0x2):opcodeblockstart] # this is the table before the opcode block
    newfiledataloctable = ''
    tmpfix = 0
    tmpfix2 = 0
    first1run = True

    if len(positionsofinsert) > 0 or len(positionsofremove) > 0:
        for fixloc in range(0x0,len(newfiledata2),0x2): # we have to fix this table as well, as each word is a block number
            fixlocval = struct.unpack('<H',newfiledata2[fixloc:fixloc+0x2])[0]
            tmpfix = fixlocval
            if len(positionsofinsert) > 0:
                for g in range(0,len(positionsofinsert)):
                    if tmpfix >= positionsofinsert[g]:
                        fixlocval += 0x1
            if len(positionsofremove) > 0:
                for g in range(0,len(positionsofremove)):
                    if tmpfix >= positionsofremove[g]:
                        fixlocval -= 0x1
            newfiledataloctable += struct.pack('<H',fixlocval)
    elif len(positionsofinsert) == 0 and len(positionsofremove) == 0:
        newfiledataloctable = newfiledata2

    # put together our new file, and fix the 2 length variables for opcode block length, and text block length
    newfiledata = newfiledata1[:0x28] + struct.pack('<H',len(newfiledata)/0x20) + struct.pack('<H',len(newstrings)) + newfiledata1[0x2c:] + newfiledataloctable + newfiledata
    newfiledata += newstrings

    outfile = open(os.path.splitext(filename)[0],'wb')
    outfile.write(newfiledata)
    outfile.close()


if __name__ == '__main__':
    if sys.argv[2] == '-e':
        extract(sys.argv[1])
    elif sys.argv[2] == '-i':
        insert(sys.argv[1])