#!/usr/bin/python
""" geckoasm
WIP tool for interconverting between PowerPC assembly and Gecko codesets,
powered by Keystone and Capstone.

NOTES
=====
* Keystone parses assembly somewhat differently than the GNU assembler.
  Be careful making assumptions about how Keystone will handle input.

* Disassembly is experimental. Recall that Capstone does not [*can not*]
  necessarily make a perfect distinction between code and plain-ol' data.
  Right now, it's up to the user to make this distinction.
"""

import argparse
import os
import re

from sys import argv
#from keystone import *
from binascii import hexlify, unhexlify

def check_files(file_list):
    """ Return True if all files in the given list are valid """
    for fn in file_list:
        if not (os.path.isfile(fn)):
            print("[!] Couldn't find infile '{}'".format(fn))
            return False
    return True

def get_metadata(filename):
    """ Return the geckoasm metadata associated with the given file.
    Expects comments (ideally in the header) of the form:

        '# key: value'

    Where some valid keys and options are:

        - name: <string>
        - author: <string>
        - code: [c2|06]
        - addr: <hex number>
        - dol-stitch: [append|prepend]
    """

    res = {'author': None,
            'name': None,
            'addr': None,
            'code': None,
            'dol_stitch': None,
            'repl': [],
            }
    with open(filename, "r") as f:
        for line in f:
            m = re.search('(?:^# ?@author: )(.*)', line)
            if (m is not None):
                res['author'] = m.group(1)

            m = re.search('(?:^# ?@name: )(.*)', line)
            if (m is not None):
                res['name'] = m.group(1)

            m = re.search('(?:^# ?@code: )(0?x?[0-9a-fA-F]{2})', line)
            if (m is not None):
                res['code'] = int(m.group(1), 16)

            m = re.search('(?:^# ?@addr: )(0?x?[0-9a-fA-f]{8})', line)
            if (m is not None):
                res['addr'] = int(m.group(1), 16)

            m = re.search('(?:^# ?@dol-stitch: )([aA]ppend|[pP]repend)', line)
            if (m is not None):
                res['dol_stitch'] = m.group(1)

            m = re.search('(?:^# ?@([0-9a-fA-F]{8}): )(.*)', line)
            if (m is not None):
                entry = {'addr': None, 'asm': None}
                entry['addr'] = int(m.group(1), 16)
                entry['asm'] = m.group(2)
                res['repl'].append(entry)
    return res

def fix_regs(data):
    """ We have to parse up assembly somewhat in order to get Keystone
    to play nicely, which sucks. This will probably break things. """
    res = data
    for reg in range(0, 31):
        # Replace "rN" with "N"
        res = res.replace("r"+str(reg), str(reg))
        res = res.replace("(r"+str(reg)+")", "("+str(reg)+")")
        res = res.replace(" r"+str(reg), " "+str(reg))
        res = res.replace("\tr"+str(reg), "\t"+str(reg))
        res = res.replace(",r"+str(reg), ","+str(reg))

        # Replace "fN" with "N"
        res = res.replace(" f"+str(reg), " "+str(reg))
        res = res.replace("\tf"+str(reg), "\t"+str(reg))

    #res = res.replace("sp", "1")
    res = res.replace("(sp)", "(1)")
    res = res.replace(" sp", " 1")
    res = res.replace("\tsp", "\t1")
    return res

def get_target_branch(filename):
    res = None
    with open(filename, "r") as f:
        for line in f:
            # For the actual codehandler command
            if ( ("#To be injected @" in line) or ("#To be inserted at" in line)):
                s = line.split(" ")
                res = int(s[4], 16)
    return res

def read_file(filename):
    with open(filename, "r") as f:
        code = f.read()
    return code

def asdf(symbol, value):
    """ overload default symbol resolving function in assembler"""
    print(symbol)
    print(value)

def assemble(text, addr):
    """ Given some set of assembly, use Keystone to assemble it into
    a bytearray of 32-bit big-endian PPC instructions """
    ks = Ks(KS_ARCH_PPC, KS_MODE_PPC32 + KS_MODE_BIG_ENDIAN)
    ks.sym_resolver = asdf
    encoding, count = ks.asm(text, addr)
    data = bytearray()
    for instr in encoding:
        data.append(instr)
    return data

def bin_to_gecko(data):
    """ Given some bytearray of binary PPC instructions, turn them into
    hex strings and pad accordingly for C2 codes """
    sdata = hexlify(data)
    instrs = [sdata[i:i+8] for i in range(0, len(sdata), 8)]
    if ((len(instrs) % 2) == 0):
        instrs += [b"60000000", b"00000000"]
        added_bytes = 8
    else:
        instrs.append(b"00000000")
        added_bytes = 4
    return (instrs, added_bytes)

def get_c2_command(branch_addr):
    """ Given some hex int, return the appropriate C2 command """
    cmd = (0xC2000000 | (branch_addr & 0x00ffffff)).to_bytes(4, byteorder='big')
    return hexlify(cmd)

def get_04_command(addr):
    """ Given some hex int, return the appropriate C2 command """
    cmd = (0x04000000 | (addr & 0x00ffffff)).to_bytes(4, byteorder='big')
    return hexlify(cmd)


def get_06_command(branch_addr):
    """ Given some hex int, return the appropriate C2 command """
    cmd = (0x06000000 | (branch_addr & 0x00ffffff)).to_bytes(4, byteorder='big')
    return hexlify(cmd)

def to_unixy_filename(x):
    """ Convert some string to lowercase, spaces to dashes """

def get_gecko_codes(filename):
    """ Pull gecko codes out of the given file, returning a list """
    gecko_codes = []
    with open(filename, "r") as f:
        block = 0
        # At the start of some C2 code
        for line in f:
            m = re.search('^[0-9a-fA-F]{8} [0-9a-fA-F]{8}', line)
            if (m is not None):
                gecko_codes += m.group(0).split(" ")
    return gecko_codes

def gecko_to_binary(codelist):
    """ Given some list of PPC instructions as hex strings,
    return the corresponding bytearray with binary data """


def disassemble(data):
    """ Given some bytearray of binary PPC instructions, produce some
    disassembly and return it """
    md = Cs(CS_ARCH_PPC, CS_MODE_32 + CS_MODE_BIG_ENDIAN)

def assemble_repl(entry):
    """ Given some dict {'asm' : <assembly string>, 'addr': <base addr> },
    representing a 04 code, emit a single instruction in hex string format """

    asm = entry['asm']
    addr = entry['addr']

    # Deal with injecting 'bl' instructions. Note that Keystone will transform
    # the immediate into an absolute address. Note that this is *not* the way
    # you expect GNU `as` to parse assembly. For instance, for PC=0x8016e74c,
    # the string 'bl 0x32ed94' will emit a branching instruction to 0x8032ed94,
    # *not* a branching instruction to (PC+0x32ed94).

    instr = None
    m = re.search('(bl).*(0?x?[0-9a-fA-F]{8})', asm)
    if (m is not None):
        bl_inst = m.group(1)
        bl_addr = int(m.group(2),16)
        if ((bl_addr > 0x80000000) and (bl_addr < 0x817fffff)):
            bl_imm = bl_addr & 0x007fffff
            bl_asm = "bl {}".format(hex(bl_imm))
            instr = assemble(bl_asm, addr)
            return hexlify(instr)
        else:
            print("[!] Branching to {} is invalid".format(hex(bl_addr)))
            return None

    # Deal with all other instructions (might have to add more particular
    # cases in the future where you need to do extra parsing, or deal with
    # invalid user input here)

    instr = assemble(entry['asm'], entry['addr'])
    return hexlify(instr)




""" ---------------------------------------------------------------------------
Use argparse to deal with user input.
"""

desc = """
DESCRIPTION
===========
Compiles sets of Gecko codes from one or more assembly files. Also supports
converting Gecko codes back into assembly (currently experimental). You will
need Python bindings for Keystone/Capstone for this script to work, see:

    https://www.keystone-engine.org/
    https://www.capstone-engine.org/

HEADER FORMAT
=============
In order to emit Gecko codes, geckoasm expects assembly files to start with a
header including some metadata about the file.

name    - The name of the code (optional)
author  - The author's name (optional)
addr    - Address (required)
code    - The Gecko code type (required)
            C2 - inject
            06 - replace block

For example:

    # @name: Do a thing
    # @author: meta
    # @addr: 0x8005c4b8
    # @code : c2
    <assembly follows>
    ...

REPLACING AN INSTRUCTION
========================
In your assembly file, you can include comments with the following form:

    # @8037b348: nop
    # @8037b348: nop
    # @8006184c: bl 0x8038fe44

This will instruction geckoasm to assemble the instruction and produce an '04'
code to patch the word at the given address.
"""

from argparse import RawTextHelpFormatter
parser = argparse.ArgumentParser(description=desc, formatter_class=RawTextHelpFormatter)

parser.add_argument('mode', help="'asm' OR 'disas'")
parser.add_argument('infile', nargs='+',
        help='assembly file[s] to-be-compiled')

parser.add_argument('-w', '--write',
        help="write output to files based on the code name",
        action="store_true")

args = parser.parse_args()

if not (check_files(args.infile)):
    exit(-1)

metadata = []
for fn in args.infile:
    metadata.append(get_metadata(fn))


if (args.mode == 'asm'):
    try:
        from keystone import *
    except ModuleNotFoundError as e:
        print("[!] It seems like you don't have Keystone installed.")
        print("[!] You can't use the 'asm' mode without Keystone.")
        print(e)
        exit(0)

    print("[*] Assembling files ...")
    for idx in range(0, len(args.infile)):
        print("============================================")
        print("[*] Assembling {}".format(args.infile[idx]))
        #print("[*]   Codetype: {:02X}".format(metadata[idx]['code']))
        #print("[*]    Address: {}".format(hex(metadata[idx]['addr'])))


        addr = metadata[idx]['addr']

        asm = read_file(args.infile[idx])
        asm = fix_regs(asm)

        data    = assemble(asm, addr)
        instrs, padding = bin_to_gecko(data)

        if (metadata[idx]['code'] == 0xc2):
            command = get_c2_command(addr)
            num_lines = hexlify((len(instrs) // 2).to_bytes(4, byteorder='big'))
            output = [command, num_lines] + instrs

        if (metadata[idx]['code'] == 0x06):
            command = get_06_command(addr)
            len_bytes = len(data) + padding
            num_lines = hexlify(len_bytes.to_bytes(4, byteorder='big'))
            output = [command, num_lines] + instrs

        # Render any single word replacements (Gecko 0x04 commands)
        repl = []
        for entry in metadata[idx]['repl']:
            repl_asm = entry['asm']
            repl_addr = entry['addr']
            repl_instr = assemble_repl(entry)
            repl_command = get_04_command(entry['addr'])
            print((repl_command + b" " + repl_instr).decode('utf8').upper())


        # Main Output
        for i in range(0, len(output), 2):
            print((output[i] + b" " + output[i+1]).decode('utf8').upper())

    exit(0)

if (args.mode == 'disas'):
    try:
        from capstone import *
    except ModuleNotFoundError as e:
        print("[!] It seems like you don't have Capstone installed.")
        print("[!] You can't use the 'disas' mode without Capstone.")
        print(e)
        exit(0)

    print("[*] Disassembling files ...")
    for idx in range(0, len(args.infile)):
        print("============================================")
        print("[*] Disassembling {}".format(args.infile[idx]))

        gecko_codes = get_gecko_codes(args.infile[idx])
        addr = metadata[idx]['addr']

        # Remove C2 code header and ending zeros
        if (metadata[idx]['code'] == 0xc2):
            gecko_codes = gecko_codes[2:]
            gecko_codes.pop()
        if (metadata[idx]['code'] == 0x06):
            gecko_codes = gecko_codes[2:]
            gecko_codes.pop()

        # Convert to binary
        data = bytearray()
        for item in gecko_codes:
            data += unhexlify(item)

        # Produce some disassembly
        md = Cs(CS_ARCH_PPC, CS_MODE_32 + CS_MODE_BIG_ENDIAN)
        disasm = md.disasm(data, addr)
        for instr in disasm:
            print("0x{:08X}:\t{}\t{}\t{}".format(instr.address,
                hexlify(instr.bytes).decode('utf8'),
                instr.mnemonic, instr.op_str))
