compose_reaction

#!/usr/bin/env python3
import os
import sys
import stat
import traceback
import collections
from json import loads, JSONDecoder
from struct import pack, calcsize
from argparse import ArgumentParser

SOCKET_TYPE_TCP =  (1 << 0)
SOCKET_TYPE_UDP = (1 << 1)
SOCKET_TYPE_IPV4 = (1 << 2)
SOCKET_TYPE_IPV6 = (1 << 3)

import os

supported_quarks = ['execve', 'execveat', 'fork-and-rename',
    'connect', 'listen', 'copy', 'remove']

class AtomDecoder(JSONDecoder):
    def __init__(self, *args, **kwargs):
        self.quark_id = 1
        JSONDecoder.__init__(self, object_pairs_hook=self.object_pairs_hook, *args, **kwargs)

    def object_pairs_hook(self, obj):
        new_obj = collections.OrderedDict()
        for k,v in obj:
            if k.lower() in supported_quarks:
                new_obj[k + ':' + str(self.quark_id)] = v
                self.quark_id += 1
            else:
                new_obj[k] = v

        return new_obj

script_directory = os.path.dirname(os.path.realpath(__file__))

def parse_argv(json_list):
    argv = b''
    for s in json_list:
        argv += pack("<{}s".format(len(s)+1), s.encode('utf8'))
    argv += b'\x00'
    return argv

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("atom_file", help="The path to the atoms definition file.")
    parser.add_argument("reaction_file", help="The path to the reaction chain file.")
    parser.add_argument("output", help="The path to where the completed reaction should be written")

    options = parser.parse_args()

    assert os.path.exists(os.path.join(script_directory, 'build/_X86_64/chain_reactor')), "You must `make` chain_reactor first"

    try:
        reaction_json = loads(open(options.reaction_file, 'rb').read().decode('utf8'))
        atom_json = loads(open(options.atom_file, 'rb').read().decode('utf8'), cls=AtomDecoder)

        reaction_section = pack("<64sL",
            reaction_json['name'].encode('utf8'),
            len(reaction_json['atoms']))

        for atom in reaction_json['atoms']:
            reaction_section += pack("<{}s".format(len(atom)+1), atom.encode('utf8'))

        atoms = b''
        for atom in atom_json:
            name = atom.get('name')
            del atom['name']
            atom_cb = 0
            atom_blob = b''
            for k, v in atom.items():
                quark = b''
                if ':' in k:
                    k = k.lower().split(':')[0]
                type = k

                if k == 'execve' or k == 'execveat':
                    assert isinstance(v, list), '{} argument must be a list'.format(k)
                    assert len(v) > 0, '{} must have arguments'.format(k)
                    method = { 'execve' : 1, 'execveat' : 2  }.get(k)
                    type = 'exec'
                    quark += pack("<L", method) + parse_argv(v)
                elif k == 'fork-and-rename' or k == 'fork-and-rename-x86':
                    type = 'fork-and-rename'
                    assert isinstance(v, list), '{} argument must be a list'.format(k)
                    assert len(v) > 0, '{} must have arguments'.format(k)
                    method = { 'fork-and-rename': 2, 'fork-and-rename-x86' : 1 }.get(k)
                    quark += pack("<L", method) + parse_argv(v)
                elif k == 'connect' or k == 'listen':
                    method = { 'syscall' : 1, 'socketcall' : 2  }.get(v.get('method'))
                    socket_type = { 'tcp4' : SOCKET_TYPE_TCP | SOCKET_TYPE_IPV4,
                        'tcp6' : SOCKET_TYPE_TCP | SOCKET_TYPE_IPV6,
                        'udp4' : SOCKET_TYPE_UDP | SOCKET_TYPE_IPV4,
                        'udp6' : SOCKET_TYPE_UDP | SOCKET_TYPE_IPV6 }.get(v.get('protocol'))
                    quark += pack("<LLH{}s".format(len(v.get('address'))+1),
                        method, socket_type, v.get('port'), v.get('address').encode('utf8'))
                elif k == 'copy':
                    assert isinstance(v, list), '{} argument must be a list'.format(k)
                    assert len(v) == 2, 'only two arguments are supported for copy'
                    quark += pack("<L", 0) + parse_argv(v)
                elif k == 'remove':
                    assert isinstance(v, list), '{} argument must be a list'.format(k)
                    assert len(v) > 0, '{} must have arguments'.format(k)
                    quark += pack("<L", 0) + parse_argv(v)
                else:
                    raise Exception("unknown quark {}".format(k))

                quark_cb = calcsize("<L64s") + len(quark)
                quark = pack("<L64s", quark_cb, type.encode('utf8')) + quark
                atom_cb += quark_cb
                atom_blob += quark

            atoms += pack("<L64sL", atom_cb + calcsize("<L64sL"),
                name.encode('utf8'), len(atom.items())) + atom_blob

        atoms_section = pack("<LL", len(atoms) + calcsize("<LL"), len(atom_json)) + atoms

        with open(os.path.join(script_directory, 'build/_X86_64/chain_reactor'), 'rb') as rf:
            with open(options.output, 'wb') as wf:
                wf.write(rf.read() + atoms_section + reaction_section)

        st = os.stat(options.output)
        os.chmod(options.output, st.st_mode | stat.S_IEXEC)
    except:
        traceback.print_exc()
        sys.exit(1)