def save_bits(file_name: str, bool_array: nparray, packed=True) -> int: ''' Save bits to a file from a bool array. Parameters ---------- file_name: string The name of the file to save. bool_array: numpy.array The bool array. packed: bool Whether to pack the bits into bytes. Defaults to True. Returns the number of bytes saved. ''' with open(file_name, 'wb') as bit_file: writer = BufferedWriter(bit_file) count = 0 if packed: for byte in pack_bools_to_bytes(bool_array): writer.write(byte) count += 1 else: for byte in bools_to_bytes(bool_array): writer.write(byte) count += 1 writer.flush() return count
def consecutive_download(node_id: str, file: io.BufferedWriter, **kwargs): """Keyword args: write_callback""" r = BackOffRequest.get(get_content_url() + 'nodes/' + node_id + '/content', stream=True) if r.status_code not in OK_CODES: raise RequestError(r.status_code, r.text) write_callback = kwargs.get('write_callback', None) total_ln = int(r.headers.get('content-length')) length = kwargs.get('length', None) if length and total_ln != length: logging.info('Length mismatch: argument %d, content %d' % (length, total_ln)) pgo = progress.Progress() curr_ln = 0 try: for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() if write_callback: write_callback(chunk) curr_ln += len(chunk) pgo.print_progress(total_ln, curr_ln) except (ConnectionError, ReadTimeoutError) as e: raise RequestError(RequestError.CODE.READ_TIMEOUT, '[acd_cli] Timeout. ' + e.__str__()) print() # break progress line r.close() return
def consecutive_download(node_id: str, file: io.BufferedWriter, **kwargs): """Keyword args: write_callback""" r = BackOffRequest.get(get_content_url() + 'nodes/' + node_id + '/content', stream=True) if r.status_code not in OK_CODES: raise RequestError(r.status_code, r.text) write_callback = kwargs.get('write_callback', None) total_ln = int(r.headers.get('content-length')) length = kwargs.get('length', None) if length and total_ln != length: logging.info('Length mismatch: argument %d, content %d' % (length, total_ln)) pgo = progress.Progress() curr_ln = 0 try: for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() if write_callback: write_callback(chunk) curr_ln += len(chunk) pgo.print_progress(total_ln, curr_ln) except (ConnectionError, ReadTimeoutError) as e: raise RequestError(RequestError.CODE.READ_TIMEOUT, '[acd_cli] Timeout. ' + e.__str__()) print() # break progress line r.close() return
def chunked_download(self, node_id: str, file: io.BufferedWriter, **kwargs): """:param kwargs: offset (int): byte offset -- start byte for ranged request length (int): total file length[!], equal to end + 1 write_callbacks (list[function]) """ ok_codes = [http.PARTIAL_CONTENT] write_callbacks = kwargs.get('write_callbacks', []) chunk_start = kwargs.get('offset', 0) length = kwargs.get('length', 100 * 1024**4) retries = 0 while chunk_start < length: chunk_end = chunk_start + CHUNK_SIZE - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= CHUNK_MAX_RETRY: raise RequestError( RequestError.CODE.FAILED_SUBREQUEST, '[acd_api] Downloading chunk failed multiple times.') r = self.BOReq.get( self.content_url + 'nodes/' + node_id + '/content', stream=True, acc_codes=ok_codes, headers={'Range': 'bytes=%d-%d' % (chunk_start, chunk_end)}) logger.debug('Range %d-%d' % (chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: r.close() logger.debug('Invalid byte range requested %d-%d' % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug('Chunk [%d-%d], retry %d.' % (chunk_start, chunk_end, retries)) continue curr_ln = 0 try: for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() for wcb in write_callbacks: wcb(chunk) curr_ln += len(chunk) finally: r.close() chunk_start += CHUNK_SIZE retries = 0 return
def send_packet(writer: io.BufferedWriter, cmd: int, flags: int = 0, data: bytes = b'') -> None: packet = Packet(cmd, flags, len(data), data, protocol_utils.Formats.HEADER) writer.write(packet.to_bytes()) writer.flush()
def save_file(file: io.BufferedWriter): """Save file without closing it Args: file (:obj:`io.BufferedWriter`): A file-like object """ file.flush() os.fsync(file.fileno()) file.seek(0)
def chunked_download(node_id: str, file: io.BufferedWriter, **kwargs): """Keyword args: offset: byte offset length: total length, equal to end - 1 write_callback """ ok_codes = [http.PARTIAL_CONTENT] write_callback = kwargs.get('write_callback', None) length = kwargs.get('length', 100 * 1024 ** 4) pgo = progress.Progress() chunk_start = kwargs.get('offset', 0) retries = 0 while chunk_start < length: chunk_end = chunk_start + CHUNK_SIZE - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= CHUNK_MAX_RETRY: raise RequestError(RequestError.CODE.FAILED_SUBREQUEST, '[acd_cli] Downloading chunk failed multiple times.') r = BackOffRequest.get(get_content_url() + 'nodes/' + node_id + '/content', stream=True, acc_codes=ok_codes, headers={'Range': 'bytes=%d-%d' % (chunk_start, chunk_end)}) logger.debug('Range %d-%d' % (chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: logger.debug('Invalid byte range requested %d-%d' % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug('Chunk [%d-%d], retry %d.' % (retries, chunk_start, chunk_end)) continue try: curr_ln = 0 for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() if write_callback: write_callback(chunk) curr_ln += len(chunk) pgo.print_progress(length, curr_ln + chunk_start) chunk_start += CHUNK_SIZE retries = 0 r.close() except (ConnectionError, ReadTimeoutError) as e: file.close() raise RequestError(RequestError.CODE.READ_TIMEOUT, '[acd_cli] Timeout. ' + e.__str__()) print() # break progress line return
def chunked_download(self, node_id: str, file: io.BufferedWriter, **kwargs): """:param kwargs: offset (int): byte offset -- start byte for ranged request length (int): total file length[!], equal to end + 1 write_callbacks (list[function]) """ ok_codes = [http.PARTIAL_CONTENT] write_callbacks = kwargs.get("write_callbacks", []) chunk_start = kwargs.get("offset", 0) length = kwargs.get("length", 100 * 1024 ** 4) retries = 0 while chunk_start < length: chunk_end = chunk_start + CHUNK_SIZE - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= CHUNK_MAX_RETRY: raise RequestError( RequestError.CODE.FAILED_SUBREQUEST, "[acd_api] Downloading chunk failed multiple times." ) r = self.BOReq.get( self.content_url + "nodes/" + node_id + "/content", stream=True, acc_codes=ok_codes, headers={"Range": "bytes=%d-%d" % (chunk_start, chunk_end)}, ) logger.debug("Range %d-%d" % (chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: logger.debug("Invalid byte range requested %d-%d" % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug("Chunk [%d-%d], retry %d." % (chunk_start, chunk_end, retries)) continue curr_ln = 0 try: for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() for wcb in write_callbacks: wcb(chunk) curr_ln += len(chunk) finally: r.close() chunk_start += CHUNK_SIZE retries = 0 return
def chunked_download(self, node_id: str, file: io.BufferedWriter, **kwargs): """:param kwargs: offset (int): byte offset -- start byte for ranged request length (int): total file length[!], equal to end + 1 write_callbacks (list[function]) """ ok_codes = [http.PARTIAL_CONTENT] write_callbacks = kwargs.get('write_callbacks', []) chunk_start = kwargs.get('offset', 0) length = kwargs.get('length', 100 * 1024 ** 4) dl_chunk_sz = self._conf.getint('transfer', 'dl_chunk_size') retries = 0 while chunk_start < length: chunk_end = chunk_start + dl_chunk_sz - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= self._conf.getint('transfer', 'chunk_retries'): raise RequestError(RequestError.CODE.FAILED_SUBREQUEST, '[acd_api] Downloading chunk failed multiple times.') r = self.BOReq.get(self.content_url + 'nodes/' + node_id + '/content', stream=True, acc_codes=ok_codes, headers={'Range': 'bytes=%d-%d' % (chunk_start, chunk_end)}) logger.debug('Node "%s", range %d-%d' % (node_id, chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: r.close() logger.debug('Invalid byte range requested %d-%d' % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug('Chunk [%d-%d], retry %d.' % (chunk_start, chunk_end, retries)) continue curr_ln = 0 try: for chunk in r.iter_content(chunk_size=self._conf.getint('transfer', 'fs_chunk_size')): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() for wcb in write_callbacks: wcb(chunk) curr_ln += len(chunk) finally: r.close() chunk_start = file.tell() retries = 0 return
def chunked_download(node_id: str, file: io.BufferedWriter, **kwargs): """Keyword args: offset (int): byte offset -- start byte for ranged request length (int): total file length[!], equal to end + 1 write_callbacks: (list[function]) """ ok_codes = [http.PARTIAL_CONTENT] write_callbacks = kwargs.get('write_callbacks', []) chunk_start = kwargs.get('offset', 0) length = kwargs.get('length', 100 * 1024 ** 4) retries = 0 while chunk_start < length: chunk_end = chunk_start + CHUNK_SIZE - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= CHUNK_MAX_RETRY: raise RequestError(RequestError.CODE.FAILED_SUBREQUEST, '[acd_cli] Downloading chunk failed multiple times.') r = BackOffRequest.get(get_content_url() + 'nodes/' + node_id + '/content', stream=True, acc_codes=ok_codes, headers={'Range': 'bytes=%d-%d' % (chunk_start, chunk_end)}) logger.debug('Range %d-%d' % (chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: logger.debug('Invalid byte range requested %d-%d' % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug('Chunk [%d-%d], retry %d.' % (chunk_start, chunk_end, retries)) continue curr_ln = 0 # connection exceptions occur here for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() for wcb in write_callbacks: wcb(chunk) curr_ln += len(chunk) chunk_start += CHUNK_SIZE retries = 0 r.close() return
def serialise(self, datastructure: dict, filestream: BinaryIO) -> BinaryIO: buffer = BufferedWriter(filestream) kwargs = { "buffer": buffer, "encoding": self._encoding, "write_through": True } with TextIOWrapper(**kwargs) as textstream: self.dump(datastructure, textstream) buffer.flush() filestream.seek(0) return BytesIO(filestream.getvalue())
def send_generic(msg: KQMLPerformative, out: BufferedWriter): """Basic send mechanism copied (more or less) from pykqml. Writes the msg as a string to the output buffer then flushes it. Args: msg (KQMLPerformative): Message to be sent out (BufferedWriter): The output to write to, needed for sending to Companions and sending on our own port. """ LOGGER.debug('Sending: %s', msg) try: msg.write(out) except IOError: LOGGER.error('IOError during message sending') out.write(b'\n') out.flush()
def stream_noop(instr: io.BufferedReader, outstr: io.BufferedWriter, chunk_size=None): orig_data_size: int = 0 dest_data_size: int = 0 data = instr.read(chunk_size) while data: orig_data_size += len(data) # no transformation dest_data_size += len(data) written = outstr.write(data) outstr.flush() # detect early if we couldnt write everything assert written == len(data) # read more data in case it is now available data = instr.read(chunk_size) return orig_data_size, dest_data_size
def remove_duplicate_testcases(virgin_bits: list, queue: Queue, tar_file: BufferedWriter, csv_path: Path = None) -> None: """ Retrieve coverage information from the queue, and delete the testcase if it does *not* lead to new coverage. """ while True: testcase, cov = queue.get() new_bits, virgin_bits = has_new_bits(cov, virgin_bits) if new_bits: # Write testcase to GZIP with open(testcase, 'rb') as inf, \ TarFile.open(fileobj=tar_file, mode='w:gz') as tar: tar.addfile(TarInfo(testcase.name), inf) tar_file.flush() if csv_path: # If a CSV file has been provided, write coverage information to # the CSV file t_bytes = count_non_255_bytes(virgin_bits) t_byte_ratio = (t_bytes * 100.0) / MAP_SIZE execs = int(testcase.name.split('id:')[1]) csv_dict = dict(unix_time='%d' % testcase.stat().st_ctime, map_size='%.02f' % t_byte_ratio, execs=execs) with open(csv_path, 'a') as outf: writer = CsvDictWriter(outf, fieldnames=CSV_FIELDNAMES) writer.writerow(csv_dict) # Delete testcase Thread(target=lambda: os.unlink(testcase)).start() queue.task_done()
def GetData(self): writer = BufferedWriter(BytesIO()) writer.write(int.to_bytes(self.scores, length=4, byteorder='little', signed=True)) writer.write(self.name.encode('utf-8')) writer.flush() return writer.raw.getvalue()
def chunked_download(self, node_id: str, file: io.BufferedWriter, **kwargs): """:param kwargs: offset (int): byte offset -- start byte for ranged request length (int): total file length[!], equal to end + 1 write_callbacks (list[function]) """ ok_codes = [http.PARTIAL_CONTENT] write_callbacks = kwargs.get('write_callbacks', []) chunk_start = kwargs.get('offset', 0) length = kwargs.get('length', 100 * 1024**4) if self._conf.getboolean('prepend', 'enabled'): length += os.path.getsize(self._conf.get('prepend', 'path')) offset += os.path.getsize(self._conf.get('prepend', 'path')) dl_chunk_sz = self._conf.getint('transfer', 'dl_chunk_size') seekable = True try: file.tell() except OSError: seekable = False retries = 0 while chunk_start < length: chunk_end = chunk_start + dl_chunk_sz - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= self._conf.getint('transfer', 'chunk_retries'): raise RequestError( RequestError.CODE.FAILED_SUBREQUEST, '[acd_api] Downloading chunk failed multiple times.') r = self.BOReq.get( self.content_url + 'nodes/' + node_id + '/content', stream=True, acc_codes=ok_codes, headers={'Range': 'bytes=%d-%d' % (chunk_start, chunk_end)}) logger.debug('Node "%s", range %d-%d' % (node_id, chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: r.close() logger.debug('Invalid byte range requested %d-%d' % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug('Chunk [%d-%d], retry %d.' % (chunk_start, chunk_end, retries)) continue curr_ln = 0 try: for chunk in r.iter_content(chunk_size=self._conf.getint( 'transfer', 'fs_chunk_size')): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() for wcb in write_callbacks: wcb(chunk) curr_ln += len(chunk) finally: r.close() if seekable: chunk_start = file.tell() else: chunk_start = chunk_start + curr_ln retries = 0 return
def write(data, wfile: io.BufferedWriter) -> None: wfile.write(json.dumps(data).encode() + b"\n") wfile.flush()
class BottomBox: """ Utility class to display something at the bottom of the screen while the stdout can still print things. Notes ----- This instructs the terminal to create a scroll window above the text in the bottom. When an update happens, bottom lines are manually drawn. We also manage the cursor so that when stdout prints something it is at the bottom of the window it scrolls up and is not drawn downstairs. This happens on stderr, this way the output of stdout can be piped. The renderer will either get called on updates (aka when you call update()) either when the terminal is resized. Also, there is an inner buffer before writing out things to make sure that all characters are written at the same time and avoid messing things up. """ def __init__(self, renderer: Callable[[int], Sequence[Text]], output: Output): """ Constructor. Parameters ---------- renderer A callable that will be called when a drawing is requested """ self.last_height = 0 self.stdout: TextIO = getattr(sys, output.value) self.renderer = renderer # noinspection PyTypeChecker self.buffer = BufferedWriter(self.stdout.buffer, buffer_size=1_000_000) signal(SIGWINCH, lambda _, __: self.update()) def _write(self, s: bytes): """ Store bytes into the buffer Parameters ---------- s Bytes to be flushed later """ self.buffer.write(s) def _flush(self): """ Flushes the bytes to the output """ self.buffer.flush() self.stdout.buffer.flush() def _save_cursor(self) -> None: """ ANSI save of the cursor position """ self._write(b"\0337") def _restore_cursor(self) -> None: """ ANSI restore of the cursor position """ self._write(b"\0338") def _set_scroll_region(self, rows: int) -> None: """ ANSI definition of the scroll area Parameters ---------- rows Number of rows in the scroll area """ self._write(f"\033[0;{rows}r".encode()) def _move_cursor_to_row(self, row: int) -> None: """ Moves the cursor to that row Parameters ---------- row 1-indexed row number """ self._write(f"\033[{row};1f".encode()) def _move_cursor_up(self) -> None: """ Moves the cursor to the row up """ self._write(b"\033[1A") def _clear_row(self) -> None: """ Clears the current row """ self._write(b"\033[K") def _get_terminal_size(self) -> TerminalSize: """ Inquires the size of the terminal and returns it """ return TerminalSize.from_bytes( ioctl(self.stdout.fileno(), TIOCGWINSZ, "\0" * 8)) def _adjust_for_height(self, new_height) -> None: """ Makes sure that there is white space enough at the bottom of the output in order for the bottom box to be drawn. """ for _ in range(self.last_height, new_height): self._write(b"\n") for _ in range(self.last_height, new_height): self._move_cursor_up() self.last_height = new_height def _clear_screen_below_cursor(self) -> None: """ Clears everything below the cursor using dedicated ANSI code """ self._write(b"\033[J") def update(self) -> None: """ Updates the content of the bottom box """ size = self._get_terminal_size() lines = self.renderer(size.cols) self._adjust_for_height(len(lines)) self._save_cursor() self._set_scroll_region(size.rows - len(lines)) start = size.rows - len(lines) + 1 for i, line in enumerate(lines): self._move_cursor_to_row(start + i) self._clear_row() self._write(line.encode()) self._restore_cursor() self._flush() def cleanup(self): """ When done, restores the scrolling idea and the configuration of the terminal (as much as possible). """ size = self._get_terminal_size() self._save_cursor() self._set_scroll_region(size.rows) self._restore_cursor() self._clear_screen_below_cursor() self._flush()
def write(data, wfile: io.BufferedWriter) -> None: wfile.write(json.dumps(data).encode() + b"\n") wfile.flush()
def chunked_download(node_id: str, file: io.BufferedWriter, **kwargs): """Keyword args: offset: byte offset length: total length, equal to end - 1 write_callback """ ok_codes = [http.PARTIAL_CONTENT] write_callback = kwargs.get('write_callback', None) length = kwargs.get('length', 100 * 1024**4) pgo = progress.Progress() chunk_start = kwargs.get('offset', 0) retries = 0 while chunk_start < length: chunk_end = chunk_start + CHUNK_SIZE - 1 if chunk_end >= length: chunk_end = length - 1 if retries >= CHUNK_MAX_RETRY: raise RequestError( RequestError.CODE.FAILED_SUBREQUEST, '[acd_cli] Downloading chunk failed multiple times.') r = BackOffRequest.get( get_content_url() + 'nodes/' + node_id + '/content', stream=True, acc_codes=ok_codes, headers={'Range': 'bytes=%d-%d' % (chunk_start, chunk_end)}) logger.debug('Range %d-%d' % (chunk_start, chunk_end)) # this should only happen at the end of unknown-length downloads if r.status_code == http.REQUESTED_RANGE_NOT_SATISFIABLE: logger.debug('Invalid byte range requested %d-%d' % (chunk_start, chunk_end)) break if r.status_code not in ok_codes: r.close() retries += 1 logging.debug('Chunk [%d-%d], retry %d.' % (retries, chunk_start, chunk_end)) continue try: curr_ln = 0 for chunk in r.iter_content(chunk_size=FS_RW_CHUNK_SZ): if chunk: # filter out keep-alive new chunks file.write(chunk) file.flush() if write_callback: write_callback(chunk) curr_ln += len(chunk) pgo.print_progress(length, curr_ln + chunk_start) chunk_start += CHUNK_SIZE retries = 0 r.close() except (ConnectionError, ReadTimeoutError) as e: file.close() raise RequestError(RequestError.CODE.READ_TIMEOUT, '[acd_cli] Timeout. ' + e.__str__()) print() # break progress line return
class BinaryWriter: """ Small utility class to write binary data. Also creates a "Memory Stream" if necessary """ def __init__(self, stream=None): if not stream: stream = BytesIO() self.writer = BufferedWriter(stream) self.written_count = 0 # region Writing # "All numbers are written as little endian." |> Source: https://core.telegram.org/mtproto def write_byte(self, value): """Writes a single byte value""" self.writer.write(pack('B', value)) self.written_count += 1 def write_int(self, value, signed=True): """Writes an integer value (4 bytes), which can or cannot be signed""" self.writer.write( int.to_bytes( value, length=4, byteorder='little', signed=signed)) self.written_count += 4 def write_long(self, value, signed=True): """Writes a long integer value (8 bytes), which can or cannot be signed""" self.writer.write( int.to_bytes( value, length=8, byteorder='little', signed=signed)) self.written_count += 8 def write_float(self, value): """Writes a floating point value (4 bytes)""" self.writer.write(pack('<f', value)) self.written_count += 4 def write_double(self, value): """Writes a floating point value (8 bytes)""" self.writer.write(pack('<d', value)) self.written_count += 8 def write_large_int(self, value, bits, signed=True): """Writes a n-bits long integer value""" self.writer.write( int.to_bytes( value, length=bits // 8, byteorder='little', signed=signed)) self.written_count += bits // 8 def write(self, data): """Writes the given bytes array""" self.writer.write(data) self.written_count += len(data) # endregion # region Telegram custom writing def tgwrite_bytes(self, data): """Write bytes by using Telegram guidelines""" if len(data) < 254: padding = (len(data) + 1) % 4 if padding != 0: padding = 4 - padding self.write(bytes([len(data)])) self.write(data) else: padding = len(data) % 4 if padding != 0: padding = 4 - padding self.write(bytes([254])) self.write(bytes([len(data) % 256])) self.write(bytes([(len(data) >> 8) % 256])) self.write(bytes([(len(data) >> 16) % 256])) self.write(data) self.write(bytes(padding)) def tgwrite_string(self, string): """Write a string by using Telegram guidelines""" self.tgwrite_bytes(string.encode('utf-8')) def tgwrite_bool(self, boolean): """Write a boolean value by using Telegram guidelines""" # boolTrue boolFalse self.write_int(0x997275b5 if boolean else 0xbc799737, signed=False) def tgwrite_date(self, datetime): """Converts a Python datetime object into Unix time (used by Telegram) and writes it""" value = 0 if datetime is None else int(datetime.timestamp()) self.write_int(value) def tgwrite_object(self, tlobject): """Writes a Telegram object""" tlobject.on_send(self) def tgwrite_vector(self, vector): """Writes a vector of Telegram objects""" self.write_int(0x1cb5c415, signed=False) # Vector's constructor ID self.write_int(len(vector)) for item in vector: self.tgwrite_object(item) # endregion def flush(self): """Flush the current stream to "update" changes""" self.writer.flush() def close(self): """Close the current stream""" self.writer.close() def get_bytes(self, flush=True): """Get the current bytes array content from the buffer, optionally flushing first""" if flush: self.writer.flush() return self.writer.raw.getvalue() def get_written_bytes_count(self): """Gets the count of bytes written in the buffer. This may NOT be equal to the stream length if one was provided when initializing the writer""" return self.written_count # with block def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close()
class BinaryWriter: """ Small utility class to write binary data. Also creates a "Memory Stream" if necessary """ def __init__(self, stream=None, known_length=None): if not stream: stream = BytesIO() if known_length is None: # On some systems, DEFAULT_BUFFER_SIZE defaults to 8192 # That's over 16 times as big as necessary for most messages known_length = max(DEFAULT_BUFFER_SIZE, 1024) self.writer = BufferedWriter(stream, buffer_size=known_length) self.written_count = 0 # region Writing # "All numbers are written as little endian." # https://core.telegram.org/mtproto def write_byte(self, value): """Writes a single byte value""" self.writer.write(pack('B', value)) self.written_count += 1 def write_int(self, value, signed=True): """Writes an integer value (4 bytes), optionally signed""" self.writer.write( int.to_bytes(value, length=4, byteorder='little', signed=signed)) self.written_count += 4 def write_long(self, value, signed=True): """Writes a long integer value (8 bytes), optionally signed""" self.writer.write( int.to_bytes(value, length=8, byteorder='little', signed=signed)) self.written_count += 8 def write_float(self, value): """Writes a floating point value (4 bytes)""" self.writer.write(pack('<f', value)) self.written_count += 4 def write_double(self, value): """Writes a floating point value (8 bytes)""" self.writer.write(pack('<d', value)) self.written_count += 8 def write_large_int(self, value, bits, signed=True): """Writes a n-bits long integer value""" self.writer.write( int.to_bytes(value, length=bits // 8, byteorder='little', signed=signed)) self.written_count += bits // 8 def write(self, data): """Writes the given bytes array""" self.writer.write(data) self.written_count += len(data) # endregion # region Telegram custom writing def tgwrite_bytes(self, data): """Write bytes by using Telegram guidelines""" if len(data) < 254: padding = (len(data) + 1) % 4 if padding != 0: padding = 4 - padding self.write(bytes([len(data)])) self.write(data) else: padding = len(data) % 4 if padding != 0: padding = 4 - padding self.write(bytes([254])) self.write(bytes([len(data) % 256])) self.write(bytes([(len(data) >> 8) % 256])) self.write(bytes([(len(data) >> 16) % 256])) self.write(data) self.write(bytes(padding)) def tgwrite_string(self, string): """Write a string by using Telegram guidelines""" self.tgwrite_bytes(string.encode('utf-8')) def tgwrite_bool(self, boolean): """Write a boolean value by using Telegram guidelines""" # boolTrue boolFalse self.write_int(0x997275b5 if boolean else 0xbc799737, signed=False) def tgwrite_date(self, datetime): """Converts a Python datetime object into Unix time (used by Telegram) and writes it """ value = 0 if datetime is None else int(datetime.timestamp()) self.write_int(value) def tgwrite_object(self, tlobject): """Writes a Telegram object""" tlobject.on_send(self) def tgwrite_vector(self, vector): """Writes a vector of Telegram objects""" self.write_int(0x1cb5c415, signed=False) # Vector's constructor ID self.write_int(len(vector)) for item in vector: self.tgwrite_object(item) # endregion def flush(self): """Flush the current stream to "update" changes""" self.writer.flush() def close(self): """Close the current stream""" self.writer.close() def get_bytes(self, flush=True): """Get the current bytes array content from the buffer, optionally flushing first """ if flush: self.writer.flush() return self.writer.raw.getvalue() def get_written_bytes_count(self): """Gets the count of bytes written in the buffer. This may NOT be equal to the stream length if one was provided when initializing the writer """ return self.written_count # with block def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close()