Add cyclic default args to context

docs/source/util/cyclic.rst

@@ -1,5 +1,6 @@
 .. testsetup:: *
 
+   from pwn import *
    from pwnlib.util.cyclic import *
 
 

pwnlib/context/__init__.py

@@ -337,6 +337,8 @@ class ContextType(object):
         'binary': None,
         'bits': 32,
         'buffer_size': 4096,
+        'cyclic_alphabet': string.ascii_lowercase,
+        'cyclic_size': 4,
         'delete_corefiles': False,
         'device': os.getenv('ANDROID_SERIAL', None) or None,
         'endian': 'little',
@@ -1249,6 +1251,32 @@ def gdbinit(self, value):
         """
         return str(value)
 
+    @_validator
+    def cyclic_alphabet(self, alphabet):
+        """Cyclic alphabet.
+
+        Default value is `string.ascii_lowercase`.
+        """
+
+        # Do not allow multiple occurrences
+        if len(set(alphabet)) != len(alphabet):
+            raise AttributeError("cyclic alphabet cannot contain duplicates")
+
+        return str(alphabet)
+
+    @_validator
+    def cyclic_size(self, size):
+        """Cyclic pattern size.
+
+        Default value is `4`.
+        """
+        size = int(size)
+
+        if size > self.bytes:
+            raise AttributeError("cyclic pattern size cannot be larger than word size")
+
+        return size
+
     #*************************************************************************
     #                               ALIASES
     #*************************************************************************

pwnlib/util/cyclic.py

@@ -2,15 +2,15 @@
 
 import string
 
-from pwnlib.context import context
+from pwnlib.context import context, LocalContext
 from pwnlib.log import getLogger
 from pwnlib.util import packing
 
 log = getLogger(__name__)
 
 # Taken from https://en.wikipedia.org/wiki/De_Bruijn_sequence but changed to a generator
-def de_bruijn(alphabet = string.ascii_lowercase, n = None):
-    """de_bruijn(alphabet = string.ascii_lowercase, n = 4) -> generator
+def de_bruijn(alphabet = None, n = None):
+    """de_bruijn(alphabet = None, n = None) -> generator
 
     Generator for a sequence of unique substrings of length `n`. This is implemented using a
     De Bruijn Sequence over the given `alphabet`.
@@ -21,8 +21,10 @@ def de_bruijn(alphabet = string.ascii_lowercase, n = None):
         alphabet: List or string to generate the sequence over.
         n(int): The length of subsequences that should be unique.
     """
+    if alphabet is None:
+        alphabet = context.cyclic_alphabet
     if n is None:
-        n = 4
+        n = context.cyclic_size
     k = len(alphabet)
     a = [0] * k * n
     def db(t, p):
@@ -42,8 +44,8 @@ def db(t, p):
 
     return db(1,1)
 
-def cyclic(length = None, alphabet = string.ascii_lowercase, n = None):
-    """cyclic(length = None, alphabet = string.ascii_lowercase, n = 4) -> list/str
+def cyclic(length = None, alphabet = None, n = None):
+    """cyclic(length = None, alphabet = None, n = None) -> list/str
 
     A simple wrapper over :func:`de_bruijn`. This function returns at most
     `length` elements.
@@ -56,17 +58,68 @@ def cyclic(length = None, alphabet = string.ascii_lowercase, n = None):
         alphabet: List or string to generate the sequence over.
         n(int): The length of subsequences that should be unique.
 
+    Notes:
+        The maximum length is `len(alphabet)**n`.
+
+        The default values for `alphabet` and `n` restrict the total space to ~446KB.
+
+        If you need to generate a longer cyclic pattern, provide a longer `alphabet`,
+        or if possible a larger `n`.
+
     Example:
-        >>> cyclic(alphabet = "ABC", n = 3)
-        'AAABAACABBABCACBACCBBBCBCCC'
+
+        Cyclic patterns are usually generated by providing a specific `length`.
+
         >>> cyclic(20)
         'aaaabaaacaaadaaaeaaa'
+
+        >>> cyclic(32)
+        'aaaabaaacaaadaaaeaaafaaagaaahaaa'
+
+        The `alphabet` and `n` arguments will control the actual output of the pattern
+
+        >>> cyclic(20, alphabet=string.ascii_uppercase)
+        'AAAABAAACAAADAAAEAAA'
+
+        >>> cyclic(20, n=8)
+        'aaaaaaaabaaaaaaacaaa'
+
+        >>> cyclic(20, n=2)
+        'aabacadaeafagahaiaja'
+
+        The size of `n` and `alphabet` limit the maximum length that can be generated.
+        Without providing `length`, the entire possible cyclic space is generated.
+
+        >>> cyclic(alphabet = "ABC", n = 3)
+        'AAABAACABBABCACBACCBBBCBCCC'
+
+        >>> cyclic(length=512, alphabet = "ABC", n = 3)
+        Traceback (most recent call last):
+        ...
+        PwnlibException: Can't create a pattern length=512 with len(alphabet)==3 and n==3
+
+        The `alphabet` can be set in `context`, which is useful for circumstances
+        when certain characters are not allowed.  See :obj:`.context.cyclic_alphabet`.
+
+        >>> context.cyclic_alphabet = "ABC"
+        >>> cyclic(10)
+        'AAAABAAACA'
+
+        The original values can always be restored with:
+
+        >>> context.clear()
+
+        The following just a test to make sure the length is correct.
+
         >>> alphabet, n = range(30), 3
         >>> len(alphabet)**n, len(cyclic(alphabet = alphabet, n = n))
         (27000, 27000)
     """
     if n is None:
-        n = 4
+        n = context.cyclic_size
+
+    if alphabet is None:
+        alphabet = context.cyclic_alphabet
 
     if len(alphabet) ** n < length:
         log.error("Can't create a pattern length=%i with len(alphabet)==%i and n==%i" \
@@ -84,8 +137,9 @@ def cyclic(length = None, alphabet = string.ascii_lowercase, n = None):
     else:
         return out
 
-def cyclic_find(subseq, alphabet = string.ascii_lowercase, n = None):
-    """cyclic_find(subseq, alphabet = string.ascii_lowercase, n = None) -> int
+@LocalContext
+def cyclic_find(subseq, alphabet = None, n = None):
+    """cyclic_find(subseq, alphabet = None, n = None) -> int
 
     Calculates the position of a substring into a De Bruijn sequence.
 
@@ -103,25 +157,71 @@ def cyclic_find(subseq, alphabet = string.ascii_lowercase, n = None):
                 integer. If an integer is provided it will be packed as a
                 little endian integer.
         alphabet: List or string to generate the sequence over.
+                  By default, uses :obj:`.context.cyclic_alphabet`.
         n(int): The length of subsequences that should be unique.
-
+                By default, uses :obj:`.context.cyclic_size`.
 
     Examples:
 
+        Let's generate an example cyclic pattern.
+
+        >>> cyclic(16)
+        'aaaabaaacaaadaaa'
+
+        Note that 'baaa' starts at offset 4.  The `cyclic_find` routine shows us this:
+
+        >>> cyclic_find('baaa')
+        4
+
+        The *default* length of a subsequence generated by `cyclic` is `4`.
+        If a longer value is submitted, it is automatically truncated to four bytes.
+
+        >>> cyclic_find('baaacaaa')
+        4
+
+        If you provided e.g. `n=8` to `cyclic` to generate larger subsequences,
+        you must explicitly provide that argument.
+
+        >>> cyclic_find('baaacaaa', n=8)
+        3515208
+
+        We can generate a large cyclic pattern, and grab a subset of it to
+        check a deeper offset.
+
         >>> cyclic_find(cyclic(1000)[514:518])
         514
+
+        Instead of passing in the byte representation of the pattern, you can
+        also pass in the integer value.  Note that this is sensitive to the
+        selected endianness via `context.endian`.
+
         >>> cyclic_find(0x61616162)
         4
+        >>> cyclic_find(0x61616162, endian='big')
+        1
+
+        You can use anything for the cyclic pattern, including non-printable
+        characters.
+
+        >>> cyclic_find(0x00000000, alphabet=unhex('DEADBEEF00'))
+        621
     """
+
+    if n is None:
+        n = context.cyclic_size
+
     if isinstance(subseq, (int, long)):
-        width = 'all' if n is None else n * 8
-        subseq = packing.pack(subseq, width, 'little', False)
-
-    if n is None and len(subseq) != 4:
-        log.warn_once("cyclic_find() expects 4-byte subsequences by default, you gave %r\n" % subseq \
-            + "Unless you specified cyclic(..., n=%i), you probably just want the first 4 bytes.\n" % len(subseq) \
-            + "Truncating the data at 4 bytes.  Specify cyclic_find(..., n=%i) to override this." % len(subseq))
-        subseq = subseq[:4]
+        subseq = packing.pack(subseq, bytes=n)
+
+    if len(subseq) != n:
+        log.warn_once("cyclic_find() expects %i-byte subsequences by default, you gave %r\n"\
+            + "Unless you specified cyclic(..., n=%i), you probably just want the first 4 bytes.\n"\
+            + "Truncating the data at 4 bytes.  Specify cyclic_find(..., n=%i) to override this.",
+            n, subseq, len(subseq), len(subseq))
+        subseq = subseq[:n]
+
+    if alphabet is None:
+        alphabet = context.cyclic_alphabet
 
     if any(c not in alphabet for c in subseq):
         return -1

pwnlib/util/packing.py

@@ -58,7 +58,7 @@ def pack(number, word_size = None, endianness = None, sign = None, **kwargs):
 
     Arguments:
         number (int): Number to convert
-        word_size (int): Word size of the converted integer or the string 'all'.
+        word_size (int): Word size of the converted integer or the string 'all' (in bits).
         endianness (str): Endianness of the converted integer ("little"/"big")
         sign (str): Signedness of the converted integer (False/True)
         kwargs: Anything that can be passed to context.local
@@ -171,7 +171,7 @@ def unpack(data, word_size = None):
 
     Arguments:
         number (int): String to convert
-        word_size (int): Word size of the converted integer or the string "all".
+        word_size (int): Word size of the converted integer or the string "all" (in bits).
         endianness (str): Endianness of the converted integer ("little"/"big")
         sign (str): Signedness of the converted integer (False/True)
         kwargs: Anything that can be passed to context.local
@@ -237,7 +237,7 @@ def unpack_many(data, word_size = None):
 
     Args
         number (int): String to convert
-        word_size (int): Word size of the converted integers or the string "all".
+        word_size (int): Word size of the converted integers or the string "all" (in bits).
         endianness (str): Endianness of the converted integer ("little"/"big")
         sign (str): Signedness of the converted integer (False/True)
         kwargs: Anything that can be passed to context.local
@@ -366,7 +366,7 @@ def make_packer(word_size = None, sign = None, **kwargs):
     faster to call this function, since it will then use a specialized version.
 
     Arguments:
-        word_size (int): The word size to be baked into the returned packer or the string all.
+        word_size (int): The word size to be baked into the returned packer or the string all (in bits).
         endianness (str): The endianness to be baked into the returned packer. ("little"/"big")
         sign (str): The signness to be baked into the returned packer. ("unsigned"/"signed")
         kwargs: Additional context flags, for setting by alias (e.g. ``endian=`` rather than index)
@@ -429,7 +429,7 @@ def make_unpacker(word_size = None, endianness = None, sign = None, **kwargs):
     faster to call this function, since it will then use a specialized version.
 
     Arguments:
-        word_size (int): The word size to be baked into the returned packer.
+        word_size (int): The word size to be baked into the returned packer (in bits).
         endianness (str): The endianness to be baked into the returned packer. ("little"/"big")
         sign (str): The signness to be baked into the returned packer. ("unsigned"/"signed")
         kwargs: Additional context flags, for setting by alias (e.g. ``endian=`` rather than index)
@@ -525,7 +525,7 @@ def flat(*args, **kwargs):
       preprocessor (function): Gets called on every element to optionally
          transform the element before flattening. If :const:`None` is
          returned, then the original value is uded.
-      word_size (int): Word size of the converted integer.
+      word_size (int): Word size of the converted integer (in bits).
       endianness (str): Endianness of the converted integer ("little"/"big").
       sign (str): Signedness of the converted integer (False/True)
 
@@ -577,7 +577,7 @@ def fit(pieces=None, **kwargs):
       preprocessor (function): Gets called on every element to optionally
          transform the element before flattening. If :const:`None` is
          returned, then the original value is used.
-      word_size (int): Word size of the converted integer.
+      word_size (int): Word size of the converted integer (in bits).
       endianness (str): Endianness of the converted integer ("little"/"big").
       sign (str): Signedness of the converted integer (False/True)