diff --git a/utils/packeter.py b/utils/packeter.py index 14608ba..f150ea8 100644 --- a/utils/packeter.py +++ b/utils/packeter.py @@ -1,6 +1,46 @@ 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 -class StructWrapper: @staticmethod def _endian(e): return ">" if e else "<" @@ -17,18 +57,11 @@ class StructWrapper: assert f in format_sizes return format_sizes[f] -class Parser(StructWrapper): - def __init__(self, data, big_endian=False): - self.data = data - self.offset = 0 +class Parser(Buffer): + def __init__(self, buffer, big_endian=False): + super().__init__(buffer) self.big_endian = big_endian - def eof(self): - return self.offset >= len(self.data) - - def remaining_size(self): - return len(self.data[self.offset:]) - def _struct_unpack(self, big_endian, f): size = self._format_size(f) assert self.remaining_size() >= size @@ -37,8 +70,7 @@ class Parser(StructWrapper): if big_endian is None: big_endian = self.big_endian - value = struct.unpack(self._endian(big_endian) + f, self.data[self.offset:self.offset+size])[0] - self.offset += size + value = struct.unpack(self._endian(big_endian) + f, self._read(size))[0] return value def read_byte(self): @@ -66,10 +98,7 @@ class Parser(StructWrapper): return self._struct_unpack(big_endian, "Q") def read_bin(self, length): - d = self.data[self.offset:self.offset+length] - assert len(self.data[self.offset:]) >= length - self.offset += length - return d + return self._read(length) def read_until(self, byte): data = b"" @@ -82,24 +111,16 @@ class Parser(StructWrapper): class Packer(StructWrapper): def __init__(self, big_endian=False): - self.buffer = b"" - self.offset = 0 + super().__init__(b"") self.big_endian = big_endian - def get(self): - return self.buffer - - def length(self): - return len(self.buffer) - 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.buffer += struct.pack(self._endian(big_endian) + f, value) - self.offset += size + self.write(struct.pack(self._endian(big_endian) + f, value)) def write_byte(self, value): self._struct_pack(big_endian, "c", value) @@ -126,8 +147,25 @@ class Packer(StructWrapper): self._struct_unpack(big_endian, "Q", value) def write_bin(self, value): - self.buffer += value - self.offset += len(value) + self.write(value) def write_string(self, value, encoding="UTF-8"): - self.write_bin(value.encode(encoding)) \ No newline at end of file + 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()