/sulley/blocks.py

import pgraph

import primitives

import sex



import zlib

import hashlib

import struct



REQUESTS = {}

CURRENT  = None



########################################################################################################################

class request (pgraph.node):

    def __init__ (self, name):

        '''

        Top level container instantiated by s_initialize(). Can hold any block structure or primitive. This can

        essentially be thought of as a super-block, root-block, daddy-block or whatever other alias you prefer.



        @type  name: String

        @param name: Name of this request

        '''



        self.name          = name



        self.label         = name    # node label for graph rendering.

        self.stack         = []      # the request stack.

        self.block_stack   = []      # list of open blocks, -1 is last open block.

        self.closed_blocks = {}      # dictionary of closed blocks.

        self.callbacks     = {}      # dictionary of list of sizers / checksums that were unable to complete rendering.

        self.names         = {}      # dictionary of directly accessible primitives.

        self.rendered      = ""      # rendered block structure.

        self.mutant_index  = 0       # current mutation index.

        self.mutant        = None    # current primitive being mutated.





    def mutate (self):

        mutated = False



        for item in self.stack:

            if item.fuzzable and item.mutate():

                mutated = True



                if not isinstance(item, block):

                    self.mutant = item



                break



        if mutated:

            self.mutant_index += 1



        return mutated





    def num_mutations (self):

        '''

        Determine the number of repetitions we will be making.



        @rtype:  Integer

        @return: Number of mutated forms this primitive can take.

        '''



        num_mutations = 0



        for item in self.stack:

            if item.fuzzable:

                num_mutations += item.num_mutations()



        return num_mutations





    def pop (self):

        '''

        The last open block was closed, so pop it off of the block stack.

        '''



        if not self.block_stack:

            raise sex.error("BLOCK STACK OUT OF SYNC")



        self.block_stack.pop()





    def push (self, item):

        '''

        Push an item into the block structure. If no block is open, the item goes onto the request stack. otherwise,

        the item goes onto the last open blocks stack.

        '''



        # if the item has a name, add it to the internal dictionary of names.

        if hasattr(item, "name") and item.name:

            # ensure the name doesn't already exist.

            if item.name in self.names.keys():

                raise sex.error("BLOCK NAME ALREADY EXISTS: %s" % item.name)



            self.names[item.name] = item



        # if there are no open blocks, the item gets pushed onto the request stack.

        # otherwise, the pushed item goes onto the stack of the last opened block.

        if not self.block_stack:

            self.stack.append(item)

        else:

            self.block_stack[-1].push(item)



        # add the opened block to the block stack.

        if isinstance(item, block):

            self.block_stack.append(item)





    def render (self):

        # ensure there are no open blocks lingering.

        if self.block_stack:

            raise sex.error("UNCLOSED BLOCK: %s" % self.block_stack[-1].name)



        # render every item in the stack.

        for item in self.stack:

            item.render()



        # process remaining callbacks.

        for key in self.callbacks.keys():

            for item in self.callbacks[key]:

                item.render()



        def update_size(stack, name):

            # walk recursively through each block to update its size

            blocks = []



            for item in stack:

                if isinstance(item, size):

                    item.render()

                elif isinstance(item, block):

                    blocks += [item]



            for b in blocks:

                update_size(b.stack, b.name)

                b.render()



        # call update_size on each block of the request

        for item in self.stack:

            if isinstance(item, block):

                update_size(item.stack, item.name)

                item.render()



        # now collect, merge and return the rendered items.

        self.rendered = ""



        for item in self.stack:

            self.rendered += item.rendered



        return self.rendered





    def reset (self):

        '''

        Reset every block and primitives mutant state under this request.

        '''



        self.mutant_index  = 1

        self.closed_blocks = {}



        for item in self.stack:

            if item.fuzzable:

                item.reset()





    def walk (self, stack=None):

        '''

        Recursively walk through and yield every primitive and block on the request stack.



        @rtype:  Sulley Primitives

        @return: Sulley Primitives

        '''



        if not stack:

            stack = self.stack



        for item in stack:

            # if the item is a block, step into it and continue looping.

            if isinstance(item, block):

                for item in self.walk(item.stack):

                    yield item

            else:

                yield item





########################################################################################################################

class block:

    def __init__ (self, name, request, group=None, encoder=None, dep=None, dep_value=None, dep_values=[], dep_compare="=="):

        '''

        The basic building block. Can contain primitives, sizers, checksums or other blocks.



        @type  name:        String

        @param name:        Name of the new block

        @type  request:     s_request

        @param request:     Request this block belongs to

        @type  group:       String

        @param group:       (Optional, def=None) Name of group to associate this block with

        @type  encoder:     Function Pointer

        @param encoder:     (Optional, def=None) Optional pointer to a function to pass rendered data to prior to return

        @type  dep:         String

        @param dep:         (Optional, def=None) Optional primitive whose specific value this block is dependant on

        @type  dep_value:   Mixed

        @param dep_value:   (Optional, def=None) Value that field "dep" must contain for block to be rendered

        @type  dep_values:  List of Mixed Types

        @param dep_values:  (Optional, def=[]) Values that field "dep" may contain for block to be rendered

        @type  dep_compare: String

        @param dep_compare: (Optional, def="==") Comparison method to apply to dependency (==, !=, >, >=, <, <=)

        '''



        self.name          = name

        self.request       = request

        self.group         = group

        self.encoder       = encoder

        self.dep           = dep

        self.dep_value     = dep_value

        self.dep_values    = dep_values

        self.dep_compare   = dep_compare



        self.stack         = []     # block item stack.

        self.rendered      = ""     # rendered block contents.

        self.fuzzable      = True   # blocks are always fuzzable because they may contain fuzzable items.

        self.group_idx     = 0      # if this block is tied to a group, the index within that group.

        self.fuzz_complete = False  # whether or not we are done fuzzing this block.

        self.mutant_index  = 0      # current mutation index.





    def mutate (self):

        mutated = False



        # are we done with this block?

        if self.fuzz_complete:

            return False



        #

        # mutate every item on the stack for every possible group value.

        #



        if self.group:

            group_count = self.request.names[self.group].num_mutations()



            # update the group value to that at the current index.

            self.request.names[self.group].value = self.request.names[self.group].values[self.group_idx]



            # mutate every item on the stack at the current group value.

            for item in self.stack:

                if item.fuzzable and item.mutate():

                    mutated = True



                    if not isinstance(item, block):

                        self.request.mutant = item



                    break



            # if the possible mutations for the stack are exhausted.

            if not mutated:

                # increment the group value index.

                self.group_idx += 1



                # if the group values are exhausted, we are done with this block.

                if self.group_idx == group_count:

                    # restore the original group value.

                    self.request.names[self.group].value = self.request.names[self.group].original_value



                # otherwise continue mutating this group/block.

                else:

                    # update the group value to that at the current index.

                    self.request.names[self.group].value = self.request.names[self.group].values[self.group_idx]



                    # this the mutate state for every item in this blocks stack.

                    # NOT THE BLOCK ITSELF THOUGH! (hence why we didn't call self.reset())

                    for item in self.stack:

                        if item.fuzzable:

                            item.reset()



                    # now mutate the first field in this block before continuing.

                    # (we repeat a test case if we don't mutate something)

                    for item in self.stack:

                        if item.fuzzable and item.mutate():

                            mutated = True



                            if not isinstance(item, block):

                                self.request.mutant = item



                            break



        #

        # no grouping, mutate every item on the stack once.

        #



        else:

            for item in self.stack:

                if item.fuzzable and item.mutate():

                    mutated = True



                    if not isinstance(item, block):

                        self.request.mutant = item



                    break



        # if this block is dependant on another field, then manually update that fields value appropriately while we

        # mutate this block. we'll restore the original value of the field prior to continuing.

        if mutated and self.dep:

            # if a list of values was specified, use the first item in the list.

            if self.dep_values:

                self.request.names[self.dep].value = self.dep_values[0]



            # if a list of values was not specified, assume a single value is present.

            else:

                self.request.names[self.dep].value = self.dep_value





        # we are done mutating this block.

        if not mutated:

            self.fuzz_complete = True



            # if we had a dependancy, make sure we restore the original value.

            if self.dep:

                self.request.names[self.dep].value = self.request.names[self.dep].original_value



        if mutated:

            if not isinstance(item, block):

                self.request.mutant = item



        return mutated





    def num_mutations (self):

        '''

        Determine the number of repetitions we will be making.



        @rtype:  Integer

        @return: Number of mutated forms this primitive can take.

        '''



        num_mutations = 0



        for item in self.stack:

            if item.fuzzable:

                num_mutations += item.num_mutations()



        # if this block is associated with a group, then multiply out the number of possible mutations.

        if self.group:

            num_mutations *= len(self.request.names[self.group].values)



        return num_mutations





    def push (self, item):

        '''

        Push an arbitrary item onto this blocks stack.

        '''



        self.stack.append(item)





    def render (self):

        '''

        Step through every item on this blocks stack and render it. Subsequent blocks recursively render their stacks.

        '''



        # add the completed block to the request dictionary.

        self.request.closed_blocks[self.name] = self



        #

        # if this block is dependant on another field and the value is not met, render nothing.

        #



        if self.dep:

            if self.dep_compare == "==":

                if self.dep_values and self.request.names[self.dep].value not in self.dep_values:

                    self.rendered = ""

                    return



                elif not self.dep_values and self.request.names[self.dep].value != self.dep_value:

                    self.rendered = ""

                    return



            if self.dep_compare == "!=":

                if self.dep_values and self.request.names[self.dep].value in self.dep_values:

                    self.rendered = ""

                    return



                elif self.request.names[self.dep].value == self.dep_value:

                    self.rendered = ""

                    return



            if self.dep_compare == ">" and self.dep_value <= self.request.names[self.dep].value:

                self.rendered = ""

                return



            if self.dep_compare == ">=" and self.dep_value < self.request.names[self.dep].value:

                self.rendered = ""

                return



            if self.dep_compare == "<" and self.dep_value >= self.request.names[self.dep].value:

                self.rendered = ""

                return



            if self.dep_compare == "<=" and self.dep_value > self.request.names[self.dep].value:

                self.rendered = ""

                return



        #

        # otherwise, render and encode as usual.

        #



        # recursively render the items on the stack.

        for item in self.stack:

            item.render()



        # now collect and merge the rendered items.

        self.rendered = ""



        for item in self.stack:

            self.rendered += item.rendered



        # if an encoder was attached to this block, call it.

        if self.encoder:

            self.rendered = self.encoder(self.rendered)



        # the block is now closed, clear out all the entries from the request back splice dictionary.

        if self.request.callbacks.has_key(self.name):

            for item in self.request.callbacks[self.name]:

                item.render()





    def reset (self):

        '''

        Reset the primitives on this blocks stack to the starting mutation state.

        '''



        self.fuzz_complete = False

        self.group_idx     = 0



        for item in self.stack:

            if item.fuzzable:

                item.reset()





########################################################################################################################

class checksum:

    checksum_lengths = {"crc32":4, "adler32":4, "md5":16, "sha1":20}



    def __init__(self, block_name, request, algorithm="crc32", length=0, endian="<", name=None):

        '''

        Create a checksum block bound to the block with the specified name. You *can not* create a checksm for any

        currently open blocks.



        @type  block_name: String

        @param block_name: Name of block to apply sizer to

        @type  request:    s_request

        @param request:    Request this block belongs to

        @type  algorithm:  String

        @param algorithm:  (Optional, def=crc32) Checksum algorithm to use. (crc32, adler32, md5, sha1)

        @type  length:     Integer

        @param length:     (Optional, def=0) Length of checksum, specify 0 to auto-calculate

        @type  endian:     Character

        @param endian:     (Optional, def=LITTLE_ENDIAN) Endianess of the bit field (LITTLE_ENDIAN: <, BIG_ENDIAN: >)

        @type  name:       String

        @param name:       Name of this checksum field

        '''



        self.block_name = block_name

        self.request    = request

        self.algorithm  = algorithm

        self.length     = length

        self.endian     = endian

        self.name       = name



        self.rendered   = ""

        self.fuzzable   = False



        if not self.length and self.checksum_lengths.has_key(self.algorithm):

            self.length = self.checksum_lengths[self.algorithm]





    def checksum (self, data):

        '''

        Calculate and return the checksum (in raw bytes) over the supplied data.



        @type  data: Raw

        @param data: Rendered block data to calculate checksum over.



        @rtype:  Raw

        @return: Checksum.

        '''



        if type(self.algorithm) is str:

            if self.algorithm == "crc32":

                return struct.pack(self.endian+"L", zlib.crc32(data))



            elif self.algorithm == "adler32":

                return struct.pack(self.endian+"L", zlib.adler32(data))



            elif self.algorithm == "md5":

                digest = hashlib.md5(data).digest()



                # XXX - is this right?

                if self.endian == ">":

                    (a, b, c, d) = struct.unpack("<LLLL", digest)

                    digest       = struct.pack(">LLLL", a, b, c, d)



                return digest



            elif self.algorithm == "sha1":

                digest = hashlib.sha1(data).digest()



                # XXX - is this right?

                if self.endian == ">":

                    (a, b, c, d, e) = struct.unpack("<LLLLL", digest)

                    digest          = struct.pack(">LLLLL", a, b, c, d, e)



                return digest



            else:

                raise sex.error("INVALID CHECKSUM ALGORITHM SPECIFIED: %s" % self.algorithm)

        else:

            return self.algorithm(data)





    def render (self):

        '''

        Calculate the checksum of the specified block using the specified algorithm.

        '''



        self.rendered = ""



        # if the target block for this sizer is already closed, render the checksum.

        if self.block_name in self.request.closed_blocks:

            block_data    = self.request.closed_blocks[self.block_name].rendered

            self.rendered = self.checksum(block_data)



        # otherwise, add this checksum block to the requests callback list.

        else:

            if not self.request.callbacks.has_key(self.block_name):

                self.request.callbacks[self.block_name] = []



            self.request.callbacks[self.block_name].append(self)





########################################################################################################################

class repeat:

    '''

    This block type is kind of special in that it is a hybrid between a block and a primitive (it can be fuzzed). The

    user does not need to be wary of this fact.

    '''



    def __init__ (self, block_name, request, min_reps=0, max_reps=None, step=1, variable=None, fuzzable=True, name=None):

        '''

        Repeat the rendered contents of the specified block cycling from min_reps to max_reps counting by step. By

        default renders to nothing. This block modifier is useful for fuzzing overflows in table entries. This block

        modifier MUST come after the block it is being applied to.



        @type  block_name: String

        @param block_name: Name of block to apply sizer to

        @type  request:    s_request

        @param request:    Request this block belongs to

        @type  min_reps:   Integer

        @param min_reps:   (Optional, def=0) Minimum number of block repetitions

        @type  max_reps:   Integer

        @param max_reps:   (Optional, def=None) Maximum number of block repetitions

        @type  step:       Integer

        @param step:       (Optional, def=1) Step count between min and max reps

        @type  variable:   Sulley Integer Primitive

        @param variable:   (Optional, def=None) Repititions will be derived from this variable, disables fuzzing

        @type  fuzzable:   Boolean

        @param fuzzable:   (Optional, def=True) Enable/disable fuzzing of this primitive

        @type  name:       String

        @param name:       (Optional, def=None) Specifying a name gives you direct access to a primitive

        '''



        self.block_name    = block_name

        self.request       = request

        self.variable      = variable

        self.min_reps      = min_reps

        self.max_reps      = max_reps

        self.step          = step

        self.fuzzable      = fuzzable

        self.name          = name



        self.value         = self.original_value = ""   # default to nothing!

        self.rendered      = ""                         # rendered value

        self.fuzz_complete = False                      # flag if this primitive has been completely fuzzed

        self.fuzz_library  = []                         # library of static fuzz heuristics to cycle through.

        self.mutant_index  = 0                          # current mutation number

        self.current_reps  = min_reps                   # current number of repetitions



        # ensure the target block exists.

        if self.block_name not in self.request.names:

            raise sex.error("CAN NOT ADD REPEATER FOR NON-EXISTANT BLOCK: %s" % self.block_name)



        # ensure the user specified either a variable to tie this repeater to or a min/max val.

        if self.variable == None and self.max_reps == None:

            raise sex.error("REPEATER FOR BLOCK %s DOES NOT HAVE A MIN/MAX OR VARIABLE BINDING" % self.block_name)



        # if a variable is specified, ensure it is an integer type.

        if self.variable and not isinstance(self.variable, primitives.bit_field):

            print self.variable

            raise sex.error("ATTEMPT TO BIND THE REPEATER FOR BLOCK %s TO A NON INTEGER PRIMITIVE" % self.block_name)



        # if not binding variable was specified, propogate the fuzz library with the repetition counts.

        if not self.variable:

            self.fuzz_library = range(self.min_reps, self.max_reps + 1, self.step)

        # otherwise, disable fuzzing as the repitition count is determined by the variable.

        else:

            self.fuzzable = False





    def mutate (self):

        '''

        Mutate the primitive by stepping through the fuzz library, return False on completion. If variable-bounding is

        specified then fuzzing is implicitly disabled. Instead, the render() routine will properly calculate the

        correct repitition and return the appropriate data.



        @rtype:  Boolean

        @return: True on success, False otherwise.

        '''



        # render the contents of the block we are repeating.

        self.request.names[self.block_name].render()



        # if the target block for this sizer is not closed, raise an exception.

        if self.block_name not in self.request.closed_blocks:

            raise sex.error("CAN NOT APPLY REPEATER TO UNCLOSED BLOCK: %s" % self.block_name)



        # if we've run out of mutations, raise the completion flag.

        if self.mutant_index == self.num_mutations():

            self.fuzz_complete = True



        # if fuzzing was disabled or complete, and mutate() is called, ensure the original value is restored.

        if not self.fuzzable or self.fuzz_complete:

            self.value        = self.original_value

            self.current_reps = self.min_reps

            return False



        if self.variable:

            self.current_reps = self.variable.value

        else:

            self.current_reps = self.fuzz_library[self.mutant_index]



        # set the current value as a multiple of the rendered block based on the current fuzz library count.

        block      = self.request.closed_blocks[self.block_name]

        self.value = block.rendered * self.fuzz_library[self.mutant_index]



        # increment the mutation count.

        self.mutant_index += 1



        return True





    def num_mutations (self):

        '''

        Determine the number of repetitions we will be making.



        @rtype:  Integer

        @return: Number of mutated forms this primitive can take.

        '''



        return len(self.fuzz_library)





    def render (self):

        '''

        Nothing fancy on render, simply return the value.

        '''



        # if the target block for this sizer is not closed, raise an exception.

        if self.block_name not in self.request.closed_blocks:

            raise sex.error("CAN NOT APPLY REPEATER TO UNCLOSED BLOCK: %s" % self.block_name)



        # if a variable-bounding was specified then set the value appropriately.

        if self.variable:

            block      = self.request.closed_blocks[self.block_name]

            self.value = block.rendered * self.variable.value



        self.rendered = self.value

        return self.rendered





    def reset (self):

        '''

        Reset the fuzz state of this primitive.

        '''



        self.fuzz_complete  = False

        self.mutant_index   = 0

        self.value          = self.original_value





########################################################################################################################

class size:

    '''

    This block type is kind of special in that it is a hybrid between a block and a primitive (it can be fuzzed). The

    user does not need to be wary of this fact.

    '''



    def __init__ (self, block_name, request, length=4, endian="<", format="binary", inclusive=False, signed=False, math=None, fuzzable=False, name=None):

        '''

        Create a sizer block bound to the block with the specified name. You *can not* create a sizer for any

        currently open blocks.



        @type  block_name: String

        @param block_name: Name of block to apply sizer to

        @type  request:    s_request

        @param request:    Request this block belongs to

        @type  length:     Integer

        @param length:     (Optional, def=4) Length of sizer

        @type  endian:     Character

        @param endian:     (Optional, def=LITTLE_ENDIAN) Endianess of the bit field (LITTLE_ENDIAN: <, BIG_ENDIAN: >)

        @type  format:     String

        @param format:     (Optional, def=binary) Output format, "binary" or "ascii"

        @type  inclusive:  Boolean

        @param inclusive:  (Optional, def=False) Should the sizer count its own length?

        @type  signed:     Boolean

        @param signed:     (Optional, def=False) Make size signed vs. unsigned (applicable only with format="ascii")

        @type  math:       Function

        @param math:       (Optional, def=None) Apply the mathematical operations defined in this function to the size

        @type  fuzzable:   Boolean

        @param fuzzable:   (Optional, def=False) Enable/disable fuzzing of this sizer

        @type  name:       String

        @param name:       Name of this sizer field

        '''



        self.block_name    = block_name

        self.request       = request

        self.length        = length

        self.endian        = endian

        self.format        = format

        self.inclusive     = inclusive

        self.signed        = signed

        self.math          = math

        self.fuzzable      = fuzzable

        self.name          = name



        self.original_value = "N/A"    # for get_primitive

        self.s_type         = "size"   # for ease of object identification

        self.bit_field      = primitives.bit_field(0, self.length*8, endian=self.endian, format=self.format, signed=self.signed)

        self.rendered       = ""

        self.fuzz_complete  = self.bit_field.fuzz_complete

        self.fuzz_library   = self.bit_field.fuzz_library

        self.mutant_index   = self.bit_field.mutant_index

        self.value          = self.bit_field.value



        if self.math == None:

            self.math = lambda (x): x





    def exhaust (self):

        '''

        Exhaust the possible mutations for this primitive.



        @rtype:  Integer

        @return: The number of mutations to reach exhaustion

        '''



        num = self.num_mutations() - self.mutant_index



        self.fuzz_complete          = True

        self.mutant_index           = self.num_mutations()

        self.bit_field.mutant_index = self.num_mutations()

        self.value                  = self.original_value



        return num





    def mutate (self):

        '''

        Wrap the mutation routine of the internal bit_field primitive.



        @rtype:  Boolean

        @return: True on success, False otherwise.

        '''



        if self.mutant_index == self.num_mutations():

            self.fuzz_complete = True



        self.mutant_index += 1



        return self.bit_field.mutate()





    def num_mutations (self):

        '''

        Wrap the num_mutations routine of the internal bit_field primitive.



        @rtype:  Integer

        @return: Number of mutated forms this primitive can take.

        '''



        return self.bit_field.num_mutations()





    def render (self):

        '''

        Render the sizer.

        '''



        self.rendered = ""



        # if the sizer is fuzzable and we have not yet exhausted the the possible bit field values, use the fuzz value.

        if self.fuzzable and self.bit_field.mutant_index and not self.bit_field.fuzz_complete:

            self.rendered = self.bit_field.render()



        # if the target block for this sizer is already closed, render the size.

        elif self.block_name in self.request.closed_blocks:

            if self.inclusive: self_size = self.length

            else:              self_size = 0



            block                = self.request.closed_blocks[self.block_name]

            self.bit_field.value = self.math(len(block.rendered) + self_size)

            self.rendered        = self.bit_field.render()



        # otherwise, add this sizer block to the requests callback list.

        else:

            if not self.request.callbacks.has_key(self.block_name):

                self.request.callbacks[self.block_name] = []



            self.request.callbacks[self.block_name].append(self)





    def reset (self):

        '''

        Wrap the reset routine of the internal bit_field primitive.

        '''



        self.bit_field.reset()