repeater.py

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import static
from Peach import generator, group
from Peach.generator import Generator


class Repeater(generator.Generator):
    """
    Will repeat a value (generated by a Generator) by round count. Can be used
    for basic buffer overflow testing.

    Example:

        gen = Repeater(None, String("A"), 3)
        gen.getValue()
        A
        gen.next()
        gen.getValue()
        AA
        gen.next()
        gen.getValue()
        AAA

    Example:

        gen = Repeater(None, Static("Peach "), 5, 3)
        gen.getValue()
        Peach
        gen.next()
        gen.getValue()
        Peach Peach Peach Peach Peach
        gen.next()
        gen.getValue()
        Peach Peach Peach Peach Peach Peach Peach Peach Peach Peach
    """

    def __init__(self, group, generator, incrementor=1, maxSteps=-1,
                 startStep=None):
        """
        @type	group: Group
        @param	group: Group this generator belongs to
        @type	generator: Generator
        @param	generator: Generator to repeat
        @type	incrementor: number
        @param	incrementor: Multiplier against round count
        @type	maxSteps: number
        @param	maxSteps: Maximum repeats
        @type	startSteps: number
        @param	startSteps: Start at this step
        """
        Generator.__init__(self)
        self._incrementor = None
        self._roundCount = 1
        self._generator = None
        self._maxSteps = -1
        self._generator = generator
        self._incrementor = incrementor
        self.setGroup(group)
        self._maxSteps = maxSteps
        self._startStep = startStep
        if self._startStep is not None:
            self._roundCount = self._startStep

    def next(self):
        self._roundCount += 1
        if self._maxSteps != -1 and self._roundCount > self._maxSteps:
            self._roundCount -= 1
            raise generator.GeneratorCompleted("Peach.repeater.Repeater")

    def getRawValue(self):
        return str(self._generator.getValue()) * \
            (self._roundCount * self._incrementor)

    def getGenerator(self):
        """
        Get Generator who's value we will repeat.

        @rtype: Generator
        @return: Generator we are repeating
        """
        return self._generator

    def setGenerator(self, generator):
        """
        Set Generator who's value we will repeat.

        @type	generator: Generator
        @param	generator: Generator to repeate
        """
        self._generator = generator

    def reset(self):
        self._roundCount = 1
        if self._startStep is not None:
            self._roundCount = self._startStep
        self._generator.reset()

    @staticmethod
    def unittest():
        g = group.Group()
        r = Repeater(g, static.Static('A'), 1, 10)
        try:
            while g.next():
                print(r.getValue())
        except group.GroupCompleted:
            pass


class RepeaterGI(generator.Generator):
    """
    Will repeat a value (generated by a Generator) by multiplier (generator).

    Example:

        Repeater(None, String("A"), BadUnsignedNumbers(None))

    Would produce a string of A's the length of each number returned by
    BadUnsignedNumbers.
    """

    def __init__(self, group, generator, incrementor):
        """
        @type	group: Group
        @param	group: Group this generator belongs to
        @type	generator: Generator
        @param	generator: Generator to repeat
        @type	incrementor: Generator
        @param	incrementor: Multiplier against round count
        """
        Generator.__init__(self)
        self._incrementor = None
        self._roundCount = 1
        self._generator = None
        self._generator = generator
        self._incrementor = incrementor
        self.setGroup(group)

    def next(self):
        self._roundCount += 1
        self._incrementor.next()

    def getRawValue(self):
        try:
            ret = str(self._generator.getValue()) * \
                int(self._incrementor.getValue())
        except OverflowError:
            # Integer overflow exception.  Oh well, we tried!
            ret = self._generator.getValue()
        except MemoryError:
            ret = self._generator.getValue()
        #print "RepeaterGI: MemoryError! Value is %d long multiplier is %d." % (
        #	len(str(ret)), int(self._incrementor.getValue()))
        return ret

    def reset(self):
        self._roundCount = 1
        self._incrementor.reset()
        self._generator.reset()