/sulley/blocks.py
Python | 824 lines | 811 code | 5 blank | 8 comment | 0 complexity | f36a62b1e0a652a3ce24065c6982942c MD5 | raw file
Possible License(s): AGPL-1.0
- 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()