import struct import zlib import hashlib import hmac class Buffer: def __init__(self, buffer): self._offset = 0 self._buffer = buffer def _write(self, data): size = len(data) if self._offset == len(self._buffer): self._buffer += data else: self._buffer = self._buffer[:self._offset] + data + self._buffer[self._offset+size:] self._offset += size def _read(self, n): data = self._buffer[self._offset:self._offset+n] self._offset = min(len(self._buffer), self._offset+n) def eof(self): return self._offset >= len(self._buffer) def remaining_size(self): return len(self._buffer[self._offset:]) def get(self): return self._buffer def length(self): return len(self._buffer) def seek(self, pos): if pos < 0: self._offset = max(0, self._offset - pos) else: self._offset = min(len(self._buffer, pos)) def tell(self): return self._offset @staticmethod def _endian(e): return ">" if e else "<" @staticmethod def _format_size(f): format_sizes = { "c": 1, "h": 2, "H": 2, "i": 4, "I": 4, "q": 8, "Q": 8, } assert f in format_sizes return format_sizes[f] class Parser(Buffer): def __init__(self, buffer, big_endian=False): super().__init__(buffer) self.big_endian = big_endian def _struct_unpack(self, big_endian, f): size = self._format_size(f) assert self.remaining_size() >= size # grab default endianess, when none is given if big_endian is None: big_endian = self.big_endian value = struct.unpack(self._endian(big_endian) + f, self._read(size))[0] return value def read_byte(self): return self._struct_unpack(big_endian, "c") def read_char(self): return chr(self.read_byte()) def read_signed_short(self, big_endian=None): return self._struct_unpack(big_endian, "h") def read_unsigned_short(self, big_endian=None): return self._struct_unpack(big_endian, "H") def read_signed_int(self, big_endian=None): return self._struct_unpack(big_endian, "i") def read_unsigned_int(self, big_endian=None): return self._struct_unpack(big_endian, "I") def read_signed_long(self, big_endian=None): return self._struct_unpack(big_endian, "q") def read_unsigned_long(self, big_endian=None): return self._struct_unpack(big_endian, "Q") def read_bin(self, length): return self._read(length) def read_until(self, byte): data = b"" while not self.eof(): c = self.read_byte() if c == byte: break data += c return data class Packer(StructWrapper): def __init__(self, big_endian=False): super().__init__(b"") self.big_endian = big_endian def _struct_pack(self, big_endian, f, value): # grab default endianess, when none is given if big_endian is None: big_endian = self.big_endian size = self._format_size(f) self.write(struct.pack(self._endian(big_endian) + f, value)) def write_byte(self, value): self._struct_pack(big_endian, "c", value) def write_char(self, value): self._struct_pack(big_endian, "c", value.encode()) def write_signed_short(self, value, big_endian=None): self._struct_pack(big_endian, "c", value) def write_unsigned_short(self, value, big_endian=None): self._struct_pack(big_endian, "H", value) def write_signed_int(self, value, big_endian=None): self._struct_unpack(big_endian, "i", value) def write_unsigned_int(self, value, big_endian=None): self._struct_unpack(big_endian, "I", value) def write_signed_long(self, value, big_endian=None): self._struct_unpack(big_endian, "q", value) def rwrite_unsigned_long(self, value, big_endian=None): self._struct_unpack(big_endian, "Q", value) def write_bin(self, value): self.write(value) def write_string(self, value, encoding="UTF-8"): self.write_bin(value.encode(encoding)) def fill(self, b: bytes, n: int): self.write_bin(b * n) def crc32(self) -> int: return zlib.crc32(self.get()) & 0xFFFFFFFF def md5(self) -> bytes: return hashlib.md5(self.get()).digest() def sha1(self) -> bytes: return hashlib.sha1(self.get()).digest() def sha256(self) -> bytes: return hashlib.sha256(self.get()).digest() def hmac(self, algorithm: str = "sha256"): return hmac.new(key, self._data, getattr(hashlib, algorithm)).digest()