diff options
author | Shadab Zafar <dufferzafar0@gmail.com> | 2016-06-02 12:59:14 +0530 |
---|---|---|
committer | Shadab Zafar <dufferzafar0@gmail.com> | 2016-06-02 13:00:02 +0530 |
commit | a04d7fd166498b54316571a4e71a58837cf11df8 (patch) | |
tree | 55a2f60c98e3c4b2e88e0054465724151919a4b9 /pathod/language | |
parent | 92cdca50c7a43a5ab59b435e73e3b17dbe01cd3b (diff) | |
download | mitmproxy-a04d7fd166498b54316571a4e71a58837cf11df8.tar.gz mitmproxy-a04d7fd166498b54316571a4e71a58837cf11df8.tar.bz2 mitmproxy-a04d7fd166498b54316571a4e71a58837cf11df8.zip |
Py3: Return bytes from RandomGenerator
and use __getitem__ rather than __getslice__
Diffstat (limited to 'pathod/language')
-rw-r--r-- | pathod/language/generators.py | 37 |
1 files changed, 22 insertions, 15 deletions
diff --git a/pathod/language/generators.py b/pathod/language/generators.py index a17e7052..fdba1ab2 100644 --- a/pathod/language/generators.py +++ b/pathod/language/generators.py @@ -3,16 +3,16 @@ import random import mmap DATATYPES = dict( - ascii_letters=string.ascii_letters, - ascii_lowercase=string.ascii_lowercase, - ascii_uppercase=string.ascii_uppercase, - digits=string.digits, - hexdigits=string.hexdigits, - octdigits=string.octdigits, - punctuation=string.punctuation, - whitespace=string.whitespace, - ascii=string.printable, - bytes="".join(chr(i) for i in range(256)) + ascii_letters=string.ascii_letters.encode(), + ascii_lowercase=string.ascii_lowercase.encode(), + ascii_uppercase=string.ascii_uppercase.encode(), + digits=string.digits.encode(), + hexdigits=string.hexdigits.encode(), + octdigits=string.octdigits.encode(), + punctuation=string.punctuation.encode(), + whitespace=string.whitespace.encode(), + ascii=string.printable.encode(), + bytes=bytes(bytearray(range(256))) ) @@ -45,6 +45,15 @@ class TransformGenerator(object): return "'transform(%s)'" % self.gen +def rand_byte(chars): + """ + Return a random character as byte from a charset. + """ + # bytearray has consistent behaviour on both Python 2 and 3 + # while bytes does not + return bytes(bytearray([random.choice(chars)])) + + class RandomGenerator(object): def __init__(self, dtype, length): @@ -55,12 +64,10 @@ class RandomGenerator(object): return self.length def __getitem__(self, x): - return random.choice(DATATYPES[self.dtype]) - - def __getslice__(self, a, b): - b = min(b, self.length) chars = DATATYPES[self.dtype] - return "".join(random.choice(chars) for x in range(a, b)) + if isinstance(x, slice): + return b"".join(rand_byte(chars) for _ in range(*x.indices(self.length))) + return rand_byte(chars) def __repr__(self): return "%s random from %s" % (self.length, self.dtype) |