diff --git a/malduck/string/bin.py b/malduck/string/bin.py index 289d416..3e37f8a 100644 --- a/malduck/string/bin.py +++ b/malduck/string/bin.py @@ -89,10 +89,9 @@ def pack(self, other: int, size: Optional[int] = None) -> bytes: :type size: bytes, optional :rtype: bytes """ - packed = unhex(f"{other:x}")[::-1] - if size: - packed = packed[:size].ljust(size, b"\x00") - return packed + if size is None: + size = (other.bit_length() + 7) // 8 + return other.to_bytes(size, byteorder="little") def unpack_be(self, other: bytes, size: Optional[int] = None) -> int: """ @@ -123,10 +122,9 @@ def pack_be(self, other: int, size: Optional[int] = None) -> bytes: :type size: bytes, optional :rtype: bytes """ - packed = unhex(f"{other:x}") - if size: - packed = packed[:size].rjust(size, b"\x00") - return packed + if size is None: + size = (other.bit_length() + 7) // 8 + return other.to_bytes(size, byteorder="big") def __call__(self, s, bitsize): warnings.warn( diff --git a/tests/test_string.py b/tests/test_string.py index 977876c..5ddfac9 100644 --- a/tests/test_string.py +++ b/tests/test_string.py @@ -120,6 +120,13 @@ def test_bigint(): assert bigint.unpack_be(b"ABCDE", 4) == 0x41424344 assert bigint.pack_be(0x41424344, 8) == b"\x00\x00\x00\x00ABCD" + assert bigint.pack(1) == b"\x01" + assert bigint.pack_be(1) == b"\x01" + assert bigint.pack(1234) == b"\xd2\x04" + assert bigint.pack_be(1234) == b"\x04\xd2" + assert bigint.unpack(b"\xd2\x04") == 1234 + assert bigint.unpack_be(b"\x04\xd2") == 1234 + def test_pack(): assert pack(