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