diff --git a/pwnlib/context/__init__.py b/pwnlib/context/__init__.py index 94a550d71..aad02f1b9 100644 --- a/pwnlib/context/__init__.py +++ b/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 #************************************************************************* diff --git a/pwnlib/util/cyclic.py b/pwnlib/util/cyclic.py index d3f934566..4cd984df3 100644 --- a/pwnlib/util/cyclic.py +++ b/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,64 @@ 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) + 'AAAABAAACAABBAABCAAC' + + 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 +133,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 +153,66 @@ 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_fid('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 """ + + 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] + width = n * 8 + subseq = packing.pack(subseq, width) + + 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 diff --git a/pwnlib/util/packing.py b/pwnlib/util/packing.py index 4c5188402..2e98eba75 100644 --- a/pwnlib/util/packing.py +++ b/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)