import io

import os

from typing import List # noqa: F401

class NonClosingBytesIO(io.BytesIO):

"""BytesIO that saves the underlying buffer on close().

This allows us to access written data after close().

"""

def __init__(self, *args, **kwargs):

super(NonClosingBytesIO, self).__init__(*args, **kwargs)

self._saved_buffer = None

def close(self):

self._saved_buffer = self.getvalue()

return super(NonClosingBytesIO, self).close()

def getvalue(self):

if self.closed:

return self._saved_buffer

else:

return super(NonClosingBytesIO, self).getvalue()

class CustomBytesIO(io.BytesIO):

def __init__(self, *args, **kwargs):

self._flush_count = 0

self._read_count = 0

self._write_count = 0

self.flush_exception = None

self.read_exception = None

self.write_exception = None

super(CustomBytesIO, self).__init__(*args, **kwargs)

def flush(self):

self._flush_count += 1

if self.flush_exception:

raise self.flush_exception

return super(CustomBytesIO, self).flush()

def read(self, *args):

self._read_count += 1

if self.read_exception:

raise self.read_exception

return super(CustomBytesIO, self).read(*args)

def write(self, data):

self._write_count += 1

if self.write_exception:

raise self.write_exception

return super(CustomBytesIO, self).write(data)

_source_files = [] # type: List[bytes]

def random_input_data():

"""Obtain the raw content of source files.

This is used for generating "random" data to feed into fuzzing, since it is

faster than random content generation.

"""

if _source_files:

return _source_files

for root, dirs, files in os.walk(os.path.dirname(__file__)):

# We filter out __pycache__ because there is a race between another

# process writing cache files and us reading them.

dirs[:] = list(sorted(d for d in dirs if d != "__pycache__"))

for f in sorted(files):

try:

with open(os.path.join(root, f), "rb") as fh:

data = fh.read()

# Exclude large files because it can cause us to easily exceed

# deadlines during fuzz testing.

if data and len(data) < 131072:

_source_files.append(data)

except OSError:

pass

# Also add some actual random data.

_source_files.append(os.urandom(100))

_source_files.append(os.urandom(1000))

_source_files.append(os.urandom(10000))

_source_files.append(os.urandom(100000))

_source_files.append(os.urandom(1000000))

return _source_files

def get_optimal_dict_size_heuristically(src):

return sum(len(ch) for ch in src) // 100

def generate_samples():

inputs = [

b"foo" * 32,

b"bar" * 16,

b"abcdef" * 64,

b"sometext" * 128,

b"baz" * 512,

]

samples = []

for i in range(128):

samples.append(inputs[i % 5])

samples.append(inputs[i % 5] * (i + 3))

samples.append(inputs[-(i % 5)] * (i + 2))

return samples