diff --git a/tools/misc/tcp_template.py b/tools/misc/tcp_template.py old mode 100644 new mode 100755 index e7f6163..b3b5f39 --- a/tools/misc/tcp_template.py +++ b/tools/misc/tcp_template.py @@ -101,6 +101,17 @@ if __name__ == "__main__": start_proxy({repr(listen_address)}, {listen_port}, {repr(remote_host)}, {remote_port}) """ +def port(value): + try: + value = int(value) + except ValueError: + raise argparse.ArgumentTypeError(f"{value!r} is not a valid integer") + + if not (1 <= value <= 65535): + raise argparse.ArgumentTypeError("port must be between 1 and 65535") + + return value + if __name__ == "__main__": parser = argparse.ArgumentParser( @@ -109,9 +120,9 @@ if __name__ == "__main__": ) parser.add_argument("la", type=str, help="Listen Address") - parser.add_argument("lp", type=int, help="Listen Port", choices=range(1,65535+1)) + parser.add_argument("lp", type=port, help="Listen Port") parser.add_argument("rh", type=str, help="Remote Host") - parser.add_argument("rp", type=int, help="Remote Port", choices=range(1,65535+1)) + parser.add_argument("rp", type=port, help="Remote Port") args = parser.parse_args() template = generate_template(args.la, args.lp, args.rh, args.rp) diff --git a/utils/packeter.py b/utils/packeter.py index f150ea8..580cdaa 100644 --- a/utils/packeter.py +++ b/utils/packeter.py @@ -73,11 +73,11 @@ class Parser(Buffer): value = struct.unpack(self._endian(big_endian) + f, self._read(size))[0] return value - def read_byte(self): + def read_byte(self, big_endian=None): return self._struct_unpack(big_endian, "c") - def read_char(self): - return chr(self.read_byte()) + def read_char(self, big_endian=None): + return chr(self.read_byte(big_endian)) def read_signed_short(self, big_endian=None): return self._struct_unpack(big_endian, "h") @@ -109,7 +109,7 @@ class Parser(Buffer): data += c return data -class Packer(StructWrapper): +class Packer(Buffer): def __init__(self, big_endian=False): super().__init__(b"") self.big_endian = big_endian @@ -120,12 +120,15 @@ class Packer(StructWrapper): big_endian = self.big_endian size = self._format_size(f) - self.write(struct.pack(self._endian(big_endian) + f, value)) + self._write(struct.pack(self._endian(big_endian) + f, value)) + + def write_byte(self, value, big_endian=None): + if isinstance(value, int): + value = bytes([value]) - def write_byte(self, value): self._struct_pack(big_endian, "c", value) - def write_char(self, value): + def write_char(self, value, big_endian=None): self._struct_pack(big_endian, "c", value.encode()) def write_signed_short(self, value, big_endian=None): @@ -135,25 +138,28 @@ class Packer(StructWrapper): self._struct_pack(big_endian, "H", value) def write_signed_int(self, value, big_endian=None): - self._struct_unpack(big_endian, "i", value) + self._struct_pack(big_endian, "i", value) def write_unsigned_int(self, value, big_endian=None): - self._struct_unpack(big_endian, "I", value) + self._struct_pack(big_endian, "I", value) def write_signed_long(self, value, big_endian=None): - self._struct_unpack(big_endian, "q", value) + self._struct_pack(big_endian, "q", value) def rwrite_unsigned_long(self, value, big_endian=None): - self._struct_unpack(big_endian, "Q", value) + self._struct_pack(big_endian, "Q", value) def write_bin(self, value): - self.write(value) + self._write(value) + + def write_bytes(self, value): + self._write(value) def write_string(self, value, encoding="UTF-8"): - self.write_bin(value.encode(encoding)) + self._write_bin(value.encode(encoding)) def fill(self, b: bytes, n: int): - self.write_bin(b * n) + self._write_bin(b * n) def crc32(self) -> int: return zlib.crc32(self.get()) & 0xFFFFFFFF