mitmproxy/pathod/language/generators.py
Maximilian Hils 3262b6e705 Squashed commit of the following:
commit edfbd41200a854f0bb7bb99f8bb70af9dbb9b8e0
Author: Matt Weidner <matt.weidner@gmail.com>
Date:   Tue Jul 25 01:19:53 2017 -0500

    Extended view.load test to check for unhandled IOError exception.

commit a523b534bc59ea97ed1fd5a3e6f78112fee19b6f
Author: requires.io <support@requires.io>
Date:   Mon Jul 24 21:25:04 2017 +0200

    [requires.io] dependency update

commit c725540c6eb92c003616b649ba43bee1f14e56ac
Author: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date:   Mon Jul 24 21:01:25 2017 +0200

    update travis

commit eeb6cfb4c76e60ac1813b839f589cd489c041c6c
Author: Thomas Kriechbaumer <Kriechi@users.noreply.github.com>
Date:   Mon Jul 24 21:03:14 2017 +0200

    [requires.io] dependency update on master branch (#2435)

commit 51a2672c782ee8ba8c7d5c7116073feccb4d8430
Author: Maximilian Hils <git@maximilianhils.com>
Date:   Mon Jul 24 19:03:01 2017 +0200

    require latest mypy version (refs #2452)

commit 5685a4850af6edda7100cae900487955c8b7a3ab
Author: Maximilian Hils <git@maximilianhils.com>
Date:   Fri Jul 21 11:24:42 2017 +0200

    fix addon tracebacks

    `.tb_next` discards the first interesting frame,
    this shouldn't happen.

commit a2da9b6c02030293f3a412d16df819868c581a29
Author: Matt Weidner <matt.weidner@gmail.com>
Date:   Sat Jul 22 12:30:15 2017 -0500

    Added os.path.expanduser() before open() calls with user supplied paths

commit 05db6e32c7957f267e97c34aa8f5a3cd6cb7dbb2
Author: Matt Weidner <matt.weidner@gmail.com>
Date:   Fri Jul 21 16:25:16 2017 -0500

    Added support for the ~ path shortcut when loading flows from disk.
    for consistency. Saving flows supports using paths with the ~ shortcut.

commit b7f864b6bbd221093aeb1c384dd16038f490441a
Author: Matt Weidner <matt.weidner@gmail.com>
Date:   Fri Jul 21 16:06:38 2017 -0500

    Fixed crash when loading flows from a file that does not exist
2017-07-25 21:55:26 +02:00

94 lines
2.4 KiB
Python

import os
import string
import random
import mmap
import sys
DATATYPES = dict(
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(range(256))
)
class TransformGenerator:
"""
Perform a byte-by-byte transform another generator - that is, for each
input byte, the transformation must produce one output byte.
gen: A generator to wrap
transform: A function (offset, data) -> transformed
"""
def __init__(self, gen, transform):
self.gen = gen
self.transform = transform
def __len__(self):
return len(self.gen)
def __getitem__(self, x):
d = self.gen.__getitem__(x)
if isinstance(x, slice):
return self.transform(x.start, d)
return self.transform(x, d)
def __repr__(self):
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([random.choice(chars)])
class RandomGenerator:
def __init__(self, dtype, length):
self.dtype = dtype
self.length = length
def __len__(self):
return self.length
def __getitem__(self, x):
chars = DATATYPES[self.dtype]
if isinstance(x, slice):
return b"".join(rand_byte(chars) for _ in range(*x.indices(min(self.length, sys.maxsize))))
return rand_byte(chars)
def __repr__(self):
return "%s random from %s" % (self.length, self.dtype)
class FileGenerator:
def __init__(self, path):
self.path = os.path.expanduser(path)
def __len__(self):
return os.path.getsize(self.path)
def __getitem__(self, x):
with open(self.path, mode="rb") as f:
if isinstance(x, slice):
with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mapped:
return mapped.__getitem__(x)
else:
f.seek(x)
return f.read(1)
def __repr__(self):
return "<%s" % self.path