class BytesIOWrapper(BufferedIOBase):
__slots__ = ("wrapped", "encoded", "encoding")
def __init__(self, wrapped: str, encoding="utf-8"):
super(BytesIOWrapper, self).__init__()
self.wrapped = wrapped
self.encoding = encoding
self.encoded = None
def read(self, *args, **kwargs):
if self.encoded is None:
b, blen = codecs.getencoder(self.encoding)(self.wrapped)
self.encoded = BytesIO(b)
return self.encoded.read(*args, **kwargs)
def read1(self, *args, **kwargs):
if self.encoded is None:
b = codecs.getencoder(self.encoding)(self.wrapped)
self.encoded = BytesIO(b)
return self.encoded.read1(*args, **kwargs)
def readinto(self, *args, **kwargs):
raise NotImplementedError()
def readinto1(self, *args, **kwargs):
raise NotImplementedError()
def write(self, *args, **kwargs):
raise NotImplementedError()
def initialize_vm(bytecode: io.BytesIO) -> VmState:
"""Init vm state with given bytecode.
:param bytecode: Bytecode
:type bytecode: io.BytesIO
:return: Initialized VmState
:rtype: VmState
"""
code_size = len(bytecode.read1())
vm_state = VmState(vm_code_buffer=bytecode)
# Prefetch all labels
while vm_state.vm_code_pointer < code_size:
vm_state.vm_code_buffer.seek(vm_state.vm_code_pointer)
bcode = vm_state.vm_code_buffer.read1(2)
opcode = struct.unpack('=h', bcode)[0]
vm_state.vm_code_buffer.seek(vm_state.vm_code_pointer)
vm_state = VM_LABEL_FUNC[opcode](vm_state)
vm_state.vm_code_pointer = 0
vm_state.vm_code_buffer.seek(0)
return vm_state
def write_gtc(data, handle):
logger.debug('In sub-method of write_gtc.py, write_gtc(data, handle)')
handle.write(b'g')
handle.write(b't')
handle.write(b'c')
handle.write(handle_byte(5))
num_entries = len(data)
handle.write(handle_int(num_entries))
offset = 8 + num_entries * 6
buffer = BytesIO()
for toc_id in data:
# write the toc ID
handle.write(handle_short(toc_id))
# write the data into the buffer
if toc_id in list_types:
handle.write(handle_int(offset + buffer.tell()))
buffer.write(handle_int(len(data[toc_id])))
for element in data[toc_id]:
buffer.write(toc2handler[toc_id](element))
else:
if toc2handler[toc_id] == handle_int:
handle.write(handle_int(data[toc_id]))
else:
handle.write(handle_int(offset + buffer.tell()))
buffer.write(toc2handler[toc_id](data[toc_id]))
buffer.seek(0)
handle.write(buffer.read1(-1))
def test_nameRoundTrip(self):
"""
Encoding and then parsing the object.
"""
for n in self.names:
f = BytesIO()
dns.Name(n).encode(f)
f.seek(0, 0)
self.parser.updateData(f.read1(-1))
self.assertEqual(self.parser.name().name, n)
def _recordRoundtripTest(self, record, name, ttl=None):
"""
Assert that encoding C{record} and then parsing the resulting bytes
creates a record which compares equal to C{record}.
"""
stream = BytesIO()
record.encode(stream)
length = stream.tell()
stream.seek(0, 0)
data = stream.read1(-1)
self.parser.updateData(data)
self.assertEqual(record, getattr(self.parser, 'payload')(name, ttl, length))
def test_queryRoundTrip(self):
"""
Encoding and then parsing the object.
"""
for n in self.names:
for dnstype in range(1, 17):
for dnscls in range(1, 5):
f = BytesIO()
dns.Query(n, dnstype, dnscls).encode(f)
f.seek(0, 0)
self.parser.updateData(f.read1(-1))
query = self.parser.query()
self.assertEqual(query.name.name, n)
self.assertEqual(query.type, dnstype)
self.assertEqual(query.cls, dnscls)
def test_NULLRecordMessage(self):
"""
A I{NULL} record with an arbitrary payload can be encoded and parsed as
part of a L{dns.Message}.
"""
bytes = b''.join([dns._ord2bytes(i) for i in range(256)])
rec = dns.Record_NULL(bytes)
rr = dns.RRHeader(b'testname', dns.NULL, payload=rec)
msg1 = dns.Message()
msg1.answers.append(rr)
s = BytesIO()
msg1.encode(s)
s.seek(0, 0)
self.parser.updateData(s.read1(-1))
msg2 = self.parser.message()
self.failUnless(isinstance(msg2.answers[0].payload, dns.Record_NULL))
self.assertEqual(msg2.answers[0].payload.payload, bytes)
def is_stl(cls, buffer):
# This is pretty much impossible to know for sure,
# but I guess we'll give it a
try:
stream = BytesIO(buffer)
start = stream.read(5)
# The only consistent thing in text STL files is the string
# 'solid' at the beginning, so if that's there, assume we have
# a valid ASCII STL.
# Also, any files served from Wikimedia's Swift loader will have
# this string at the beginning for sure. Hacky, but it works.
if start == 'solid':
return True
# Now we either have a non-STL, or we have a binary STL file.
# There is literally no consistent way to tell whether a binary
# file is actually an STL file, but the header is 80 bytes...
stream.read(75)
# ...and then there's 4 bytes that indicate how many triangles are
# present in the file...
triangles = stream.read(4)
if len(triangles) < 4:
return False
triangles = struct.unpack("<L", triangles)[0]
# ...so there should be exactly that number, times 386 (the number
# of bytes in each triangle), left in the file...
remainder = stream.read(triangles * 386)
# ...so if we got at east that many bytes...
if len(remainder) == triangles * 386:
# ...and there aren't any extras...
extra = stream.read1(1)
if len(extra) == 0:
# ...this is as close to a golden STL file check as we're
# going to get.
return True
except IndexError:
pass
return False
def test_NAPTR(self):
naptrs = [
(100, 10, b"u", b"sip+E2U",
b"!^.*$!sip:[email protected]!", b""),
(100, 50, b"s", b"http+I2L+I2C+I2R",
b"", b"_http._tcp.gatech.edu")]
for (order, preference, flags, service, regexp, replacement) in naptrs:
rin = dns.Record_NAPTR(order, preference, flags, service, regexp,
replacement)
e = BytesIO()
rin.encode(e)
length = e.tell()
e.seek(0, 0)
self.parser.updateData(e.read1(-1))
rout = getattr(self.parser, 'payload')('NAPTR', None, length)
self.assertEqual(rin.order, rout.order)
self.assertEqual(rin.preference, rout.preference)
self.assertEqual(rin.flags, rout.flags)
self.assertEqual(rin.service, rout.service)
self.assertEqual(rin.regexp, rout.regexp)
self.assertEqual(rin.replacement.name, rout.replacement.name)
self.assertEqual(rin.ttl, rout.ttl)
class VerifiableStream(BinaryIO):
"""A binary stream whose contents can be verified to not have changed.
The stream does not accept a HMAC key, but generates it randomly as a nonce. While unusual,
this is intentional -- these streams are meant to be used as part of model serialization,
where their nonces and HMAC codes are stored in a cryptographically signed metadata file.
In other words, the HMAC simply ensures that stream's data has not changed, and does not
guarantee the data's origin -- that's the metadata signature's job.
The stream is meant to be used in the following sequence:
- instantiate the stream
- write all data to the stream (the stream is not readable yet!)
- call "finalize()" on the stream, saving the returned nonce and HMAC code
- read data from the stream (the stream is not writable any more!)
"""
def __init__(self):
"""Create a new VerifiableStream with a random nonce."""
self._finalized = False
self._random_nonce = os.urandom(
16) # this is bytes, be careful trying to add strings to it
self._underlying_stream = BytesIO()
self._hmac_state = hmac.new(self._random_nonce, digestmod=HASHER)
def _ensure_finalized(self):
"""Raise an error if the stream has not already been finalized."""
if not self._finalized:
raise AssertionError(
"Expected the stream to be finalized, but it was not!")
def _ensure_not_finalized(self):
"""Raise an error if the stream has already been finalized."""
if self._finalized:
raise AssertionError(
"Expected the stream to not be finalized, but it was!")
def finalize(self):
"""Calculate the HMAC code for the stream, disable writing and enable reading.
Returns:
tuple (nonce, HMAC code) (both of type string)
"""
self._ensure_not_finalized()
self._finalized = True
nonce_string = _convert_base64_bytes_to_string(self._random_nonce)
hmac_string = _convert_base64_bytes_to_string(
self._hmac_state.digest())
return nonce_string, hmac_string
# methods for writing require that the stream not be finalized
def writable(self) -> bool:
"""Return True if the stream is writable, and False otherwise."""
if self._finalized:
return False
else:
return self._underlying_stream.writable()
@validate(b=bytes)
def write(self, b: bytes) -> int:
"""Write the given binary data to the stream, and include it in the HMAC calculation."""
self._ensure_not_finalized()
num_bytes = self._underlying_stream.write(b)
self._hmac_state.update(b)
return num_bytes
def writelines(self, lines: Iterable[bytes]) -> None:
"""Write lines to a stream"""
self._ensure_not_finalized(
) # technically done by `write` but doesn't hurt to be safe
for line in lines:
self.write(line)
return None
# methods for reading require that the stream is finalized
def readable(self) -> bool:
"""Return True if the stream is readable, and False otherwise."""
if self._finalized:
return self._underlying_stream.readable()
else:
return False
def read(self, size=None) -> bytes:
"""Read bytes from stream"""
self._ensure_finalized()
return self._underlying_stream.read(size)
def readall(self) -> bytes:
"""Read lines from stream"""
raise NotImplementedError(
"`VerifiablStream` does not implement `readall` since the underlying BtytesIO does not "
"implement it.")
def readline(self, size=None) -> bytes:
"""Read a line from stream"""
self._ensure_finalized()
return self._underlying_stream.readline(size)
def readlines(self, size=None) -> List[bytes]:
"""Read lines from stream"""
self._ensure_finalized()
return self._underlying_stream.readlines(size)
def read1(self, size) -> bytes:
"""Read bytes from stream"""
self._ensure_finalized()
return self._underlying_stream.read1(size)
def readinto(self, b) -> Optional[int]:
"""Read bytes into another buffer"""
self._ensure_finalized()
return self._underlying_stream.readinto(b)
def readinto1(self, b) -> Optional[int]:
"""Read bytes into another buffer"""
self._ensure_finalized()
return self._underlying_stream.readinto1(b)
# seeking requires a finalized stream
def seekable(self):
"""Return True if the read pointer in the stream can be moved, and False otherwise."""
if self._finalized:
return self._underlying_stream.seekable()
else:
return False
def seek(self, *args, **kwargs) -> int:
"""Seek to a new position. Return the new position"""
self._ensure_finalized()
return self._underlying_stream.seek(*args, **kwargs)
def truncate(self, size: Optional[int] = ...) -> None:
"""Truncate the stream"""
raise NotImplementedError(
"`VerifiableStream` does not support truncation. It is too "
"complicated to keep track of the hmac digests")
def close(self):
"""Close the stream, discarding its data. Will raise an error if not finalized yet."""
if self._finalized:
return self._underlying_stream.close()
else:
raise AssertionError(
"Attempting to close an unfinalized VerifiableStream. This is "
"almost certainly a bug.")
# a bunch of attributes/methods that are always accessible
def isatty(self) -> bool:
"""Determine whether this is a terminal"""
return self._underlying_stream.isatty()
@property
def closed(self) -> bool:
"""Determine whether the stream is closed"""
return self._underlying_stream.closed
def fileno(self) -> int:
"""Return the underlying file descriptor"""
# this will technically raise UnsuportedOperation, but better to let BytesIO do that
return self._underlying_stream.fileno()
def mode(self) -> str:
"""Return the underlying file descriptor"""
# this doesn't exist for the underlying stream
raise AssertionError(
"`VerifiableStream` does not have a mode. This is probably a bug in "
"something assuming that the stream is a backed by a file")
def name(self) -> str:
"""Return the underlying file descriptor"""
# this doesn't exist for the underlying stream
raise AssertionError(
"`VerifiableStream` does not have a name. This is probably a bug in "
"something assuming the stream is a file descriptor")
def flush(self) -> None:
"""Flush the underlying stream"""
# this technically does nothing in BytesIO
return self._underlying_stream.flush()
def tell(self) -> int:
"""Tell the current position"""
return self._underlying_stream.tell()
# context manager methods
def __enter__(self) -> "VerifiableStream":
"""Enter"""
return self
def __exit__(
self,
exc_type: Optional[Type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional[TracebackType],
) -> bool:
"""Exit"""
return self._underlying_stream.__exit__(exc_type, exc_val, exc_tb)
def _read_bytes(bio: io.BytesIO, n: int) -> bytes:
x = bio.read1(n)
if len(x) != n:
raise EOFError('Unexpected number of bytes read')
return x
