
def _bin(x, width):
    return ''.join(str((x>>i)&1) for i in xrange(width-1,-1,-1))

from construct import *
import struct

d_opcode = BitStruct("double_arg",
     BitField("op_code", 4),
     BitField("r1", 14),
     BitField("r2", 14),
)

s_opcode = BitStruct("single_arg",
     BitField("blank", 4),
     BitField("op_code", 4),
     BitField("imm", 10),
     BitField("r1", 14),
)    

D = ['OOPS!', 
    'r[%(destination)i] = r[%(r1)i] + r[%(r2)i]', 
    'r[%(destination)i] = r[%(r1)i] - r[%(r2)i]', 
    'r[%(destination)i] = r[%(r1)i] * r[%(r2)i]', 
    'r[%(destination)i] = safe_div(r[%(r1)i], r[%(r2)i])', 
    'port[%(r1)i] = r[%(r2)i]', 
    'r[%(destination)i] = r[%(r1)i] if machine.status else r[%(r2)i]']
S = [
    '# offset 0x%(destination)x (%(destination)i): ', 
    'machine.status = %(imm)s(r[%(r1)i], 0.0)', 
    'r[%(destination)i] = sqrt(r[%(r1)i])',
    'r[%(destination)i] = r[%(r1)i]', 
    'r[%(destination)i] = input[%(r1)i]'
]

def little_endian(bytes):
    _1, _2, _3, _4 = bytes
    return _4+_3+_2+_1

def decode(little_endian_bytes, index):
    parsed = d_opcode.parse(little_endian_bytes)
    OP = parsed.op_code
    if OP:        
        result = D[OP] % dict( 
            r1 = parsed.r1,
            r2 = parsed.r2,
            destination = index, 
        )
        print result
        return result
    else:
        parsed = s_opcode.parse(little_endian_bytes)
        OP = parsed.op_code
        result = S[OP] % dict( 
            r1 = parsed.r1,
            imm = ['lt','le','eq','ge','gt'][parsed.imm >> 7],
            destination = index, 
        )
        print result
        return result
    return ''
    
def translate(f, output, *a, **kw):
    f_bytes = f.read(8)
    i_bytes = f.read(4)

    instructions = []
    float_data = []
    
    def decode_instructions_and_data(f_bytes, i_bytes):
        float_data.append(struct.unpack('d', f_bytes)[0])
        instructions.append(little_endian(i_bytes))        
    
    while f_bytes and i_bytes:
        decode_instructions_and_data(f_bytes, i_bytes)
        i_bytes = f.read(4)
        f_bytes = f.read(8)
        if f_bytes and i_bytes:
            decode_instructions_and_data(f_bytes, i_bytes)
            f_bytes = f.read(8)
            i_bytes = f.read(4)
    
    output(instructions, float_data, *a, **kw)

def write_file(python, file_name):
    output_file = open(file_name, "w")
    output_file.write(python)
    output_file.close()
    
def generate_python_class(instructions, float_data, do, *a, **kw):
    decoded_instructions = []
    index = 0
    for instruction in instructions:
        decoded_instructions.append(decode(instruction, index))
        index += 1
    do("""

from math import sqrt
from operator import *

def safe_div(a,b):
    try:
        return a/b
    except ZeroDivisionError:
        return 0.0

class OrbitalVirtualMachine(object):
    def __init__(machine):
        machine.r = %(float_data)s
        machine.port = list(0.0 for i in xrange(20))
        machine.input = list(0.0 for i in xrange(16001))
            
    def __iter__(machine):
        r = machine.r
        port = machine.port
        input = machine.input
        while True:
            %(instructions)s
            yield port

    """ % dict(
        instructions = '\n            '.join(decoded_instructions),
        float_data = float_data
    ), *a, **kw)
    

def write_python_class(instructions, float_data, filename):
    generate_python_class(instructions, float_data, write_file, filename)
    
from functools import partial
def generate_bytes(instructions, float_data, do, *a, **kw):
    do("""
{
    "data":
        %(float_data)s
    ,
    "code":
        [%(instructions)s]
}
    """ % dict(
        instructions = ', '.join(map(lambda x:str(x[0]).replace('L',''), map(partial(struct.unpack, 'I'), instructions))),
        float_data = float_data
    ), *a, **kw)

    
def write_bytes(instructions, float_data, filename):
    generate_bytes(instructions, float_data, write_file, filename)    

translate(open("../binaries/bin1.obf", "rb"), write_bytes, "../Python/translated_binaries/bin1.json")
translate(open("../binaries/bin2.obf", "rb"), write_bytes, "../Python/translated_binaries/bin2.json")
translate(open("../binaries/bin3.obf", "rb"), write_bytes, "../Python/translated_binaries/bin3.json")


translate(open("../binaries/bin1.obf", "rb"), write_python_class, "../Python/translated_binaries/bin1.py")
translate(open("../binaries/bin2.obf", "rb"), write_python_class, "../Python/translated_binaries/bin2.py")
translate(open("../binaries/bin3.obf", "rb"), write_python_class, "../Python/translated_binaries/bin3.py")
