HackingScripts/utils/packeter.py

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()