Add Type Annotations for parts of pwnlib/util

CHANGELOG.md

@@ -116,6 +116,7 @@ The table below shows which release corresponds to each branch, and what date th
 - [#2647][2647] packing: Add `overlap` to overlap structures easily
 - [#2669][2669] asm: try native binutils before fallback architectures
 - [#2673][2673] Add libc module for libc-related functions
+- [#2679][2679] Add type hints to parts of pwnlib.utils
 - [#2680][2680] Cleanup Python 2 legacy
 - [#2687][2687] Add (un)pack shorthands for 40-56 bit numbers `u48()`/`p48()`
 - [#2699][2699] Fix `tty` and `raw` arguments in `ssh.process()`
@@ -169,6 +170,7 @@ The table below shows which release corresponds to each branch, and what date th
 [2647]: https://github.com/Gallopsled/pwntools/pull/2647
 [2669]: https://github.com/Gallopsled/pwntools/pull/2669
 [2673]: https://github.com/Gallopsled/pwntools/pull/2673
+[2679]: https://github.com/Gallopsled/pwntools/pull/2679
 [2680]: https://github.com/Gallopsled/pwntools/pull/2680
 [2687]: https://github.com/Gallopsled/pwntools/pull/2687
 [2699]: https://github.com/Gallopsled/pwntools/pull/2699

pwnlib/util/crc/known.py

@@ -1,8 +1,9 @@
 import os
 import re
+from typing import Any
 
 
-def generate():
+def generate() -> dict[str, dict]:
     """Generates a dictionary of all the known CRC formats from:
     https://reveng.sourceforge.io/crc-catalogue/all.htm
 
@@ -15,7 +16,7 @@ def generate():
         data = fd.read()
     out = {}
 
-    def fixup(s):
+    def fixup(s: str) -> Any | int | bool:
         if s == 'true':
             return True
         elif s == 'false':

pwnlib/util/fiddling.py

@@ -6,6 +6,8 @@
 import string
 
 from io import BytesIO
+from typing import Any, Generator, Iterable, Optional, BinaryIO
+from collections.abc import Sequence
 
 from pwnlib.context import LocalNoarchContext
 from pwnlib.context import context
@@ -20,7 +22,7 @@
 
 log = getLogger(__name__)
 
-def unhex(s):
+def unhex(s: Sequence) -> bytes:
     r"""unhex(s) -> str
 
     Hex-decodes a string.
@@ -42,7 +44,7 @@ def unhex(s):
             s = '0' + s
     return binascii.unhexlify(s)
 
-def enhex(x):
+def enhex(x: Sequence) -> str:
     """enhex(x) -> str
 
     Hex-encodes a string.
@@ -58,7 +60,7 @@ def enhex(x):
     return x
 
 
-def hexstr(s, force=False):
+def hexstr(s: bytes, force: bool = False) -> str:
     r"""
     hexstr(x, force=False) -> str
 
@@ -82,7 +84,7 @@ def hexstr(s, force=False):
     return out.decode()
 
 
-def urlencode(s):
+def urlencode(s: str) -> str:
     """urlencode(s) -> str
 
     URL-encodes a string.
@@ -94,7 +96,7 @@ def urlencode(s):
     """
     return ''.join(['%%%02x' % ord(c) for c in s])
 
-def urldecode(s, ignore_invalid = False):
+def urldecode(s: str, ignore_invalid: bool = False) -> str:
     """urldecode(s, ignore_invalid = False) -> str
 
     URL-decodes a string.
@@ -128,13 +130,13 @@ def urldecode(s, ignore_invalid = False):
                 raise ValueError("Invalid input to urldecode")
     return res
 
-def bits(s, endian = 'big', zero = 0, one = 1):
+def bits(s: int | bytes, endian: str = 'big', zero: str = 0, one: str = 1) -> list[int | str]:
     """bits(s, endian = 'big', zero = 0, one = 1) -> list
 
     Converts the argument into a list of bits.
 
     Arguments:
-        s: A string or number to be converted into bits.
+        s: A bytestring or number to be converted into bits.
         endian (str): The binary endian, default 'big'.
         zero: The representing a 0-bit.
         one: The representing a 1-bit.
@@ -185,7 +187,7 @@ def bits(s, endian = 'big', zero = 0, one = 1):
 
     return out
 
-def bits_str(s, endian = 'big', zero = '0', one = '1'):
+def bits_str(s: Sequence, endian: str = 'big', zero: str = '0', one: str = '1') -> str:
     """bits_str(s, endian = 'big', zero = '0', one = '1') -> str
 
     A wrapper around :func:`bits`, which converts the output into a string.
@@ -199,7 +201,7 @@ def bits_str(s, endian = 'big', zero = '0', one = '1'):
     """
     return ''.join(bits(s, endian, zero, one))
 
-def unbits(s, endian = 'big'):
+def unbits(s: Iterable, endian: str = 'big') -> str:
     r"""unbits(s, endian = 'big') -> str
 
     Converts an iterable of bits into a string.
@@ -247,7 +249,7 @@ def unbits(s, endian = 'big'):
     return out
 
 
-def bitswap(s):
+def bitswap(s: bytes) -> bytes:
     r"""bitswap(s) -> str
 
     Reverses the bits in every byte of a given string.
@@ -265,7 +267,7 @@ def bitswap(s):
 
     return b''.join(out)
 
-def bitswap_int(n, width):
+def bitswap_int(n: int, width: int) -> str:
     """bitswap_int(n) -> int
 
     Reverses the bits of a numbers and returns the result as a new number.
@@ -295,7 +297,7 @@ def bitswap_int(n, width):
     return int(s, 2)
 
 
-def b64e(s):
+def b64e(s: bytes) -> str:
     """b64e(s) -> str
 
     Base64 encodes a string
@@ -310,7 +312,7 @@ def b64e(s):
         x = x.decode('ascii')
     return x
 
-def b64d(s):
+def b64d(s: str) -> bytes:
     """b64d(s) -> str
 
     Base64 decodes a string
@@ -323,7 +325,7 @@ def b64d(s):
     return base64.b64decode(s)
 
 # misc binary functions
-def xor(*args, **kwargs):
+def xor(*args: tuple, **kwargs: dict[str, Any]) -> bytes:
     """xor(*args, cut = 'max') -> str
 
     Flattens its arguments using :func:`pwnlib.util.packing.flat` and
@@ -375,15 +377,15 @@ def xor(*args, **kwargs):
     else:
         raise ValueError("Not a valid argument for 'cut'")
 
-    def get(n):
+    def get(n: int) -> bytes:
         rv = 0
         for s in strs: rv ^= s[n%len(s)]
         return packing._p8lu(rv)
 
     return b''.join(map(get, range(cut)))
 
-def xor_pair(data, avoid = b'\x00\n'):
-    r"""xor_pair(data, avoid = '\x00\n') -> None or (str, str)
+def xor_pair(data: int | Sequence, avoid: bytes = b'\x00\n') -> Optional[tuple[str, str]]:
+    r"""xor_pair(data, avoid = '\\x00\\n') -> None or (str, str)
 
     Finds two strings that will xor into a given string, while only
     using a given alphabet.
@@ -427,7 +429,7 @@ def xor_pair(data, avoid = b'\x00\n'):
 
     return res1, res2
 
-def xor_key(data, avoid=b'\x00\n', size=None):
+def xor_key(data: str, avoid: bytes = b'\x00\n', size: int = None) -> Optional[tuple[str, str]]:
     r"""xor_key(data, size=None, avoid='\x00\n') -> None or (int, str)
 
     Finds a ``size``-width value that can be XORed with a string
@@ -476,7 +478,7 @@ def xor_key(data, avoid=b'\x00\n', size=None):
 
     return result, xor(data, result)
 
-def randoms(count, alphabet = string.ascii_lowercase):
+def randoms(count: int, alphabet: str = string.ascii_lowercase) -> str:
     """randoms(count, alphabet = string.ascii_lowercase) -> str
 
     Returns a random string of a given length using only the specified alphabet.
@@ -497,7 +499,7 @@ def randoms(count, alphabet = string.ascii_lowercase):
     return ''.join(random.choice(alphabet) for _ in range(count))
 
 
-def rol(n, k, word_size = None):
+def rol(n: Sequence | int, k: int, word_size: int = None) -> str:
     """Returns a rotation by `k` of `n`.
 
     When `n` is a number, then means ``((n << k) | (n >> (word_size - k)))`` truncated to `word_size` bits.
@@ -540,12 +542,12 @@ def rol(n, k, word_size = None):
     else:
         raise ValueError("rol(): 'n' must be an integer, string, list or tuple")
 
-def ror(n, k, word_size = None):
+def ror(n: Sequence | int, k: int, word_size: int = None) -> str:
     """A simple wrapper around :func:`rol`, which negates the values of `k`."""
 
     return rol(n, -k, word_size)
 
-def naf(n):
+def naf(n: int) -> Generator[int, None, None]:
     """naf(int) -> int generator
 
     Returns a generator for the non-adjacent form (NAF[1]) of a number, `n`.  If
@@ -571,7 +573,7 @@ def naf(n):
         n = (n - z) // 2
         yield z
 
-def isprint(c):
+def isprint(c: str | int) -> bool:
     """isprint(c) -> bool
 
     Return True if a character is printable"""
@@ -581,7 +583,7 @@ def isprint(c):
     return c in t
 
 
-def hexii(s, width = 16, skip = True):
+def hexii(s: str, width: int = 16, skip: bool = True) -> str:
     """hexii(s, width = 16, skip = True) -> str
 
     Return a HEXII-dump of a string.
@@ -597,7 +599,7 @@ def hexii(s, width = 16, skip = True):
 
     return hexdump(s, width, skip, True)
 
-def _hexiichar(c):
+def _hexiichar(c: int) -> str:
     HEXII = bytearray((string.punctuation + string.digits + string.ascii_letters).encode())
     if c in HEXII:
         return ".%c " % c
@@ -619,16 +621,16 @@ def _hexiichar(c):
 cyclic_pregen = b''
 de_bruijn_gen = de_bruijn()
 
-def sequential_lines(a,b):
+def sequential_lines(a: int, b: int) -> bool:
     return (a+b) in cyclic_pregen
 
-def update_cyclic_pregenerated(size):
+def update_cyclic_pregenerated(size: int) -> None:
     global cyclic_pregen
     while size > len(cyclic_pregen):
         cyclic_pregen += packing._p8lu(next(de_bruijn_gen))
 
-def hexdump_iter(fd, width=16, skip=True, hexii=False, begin=0, style=None,
-                 highlight=None, cyclic=False, groupsize=4, total=True):
+def hexdump_iter(fd: BinaryIO, width: int = 16, skip: bool = True, hexii: bool = False, begin: int = 0, style: dict = None,
+                 highlight: Iterable = None, cyclic: bool = False, groupsize: int = 4, total: bool = True) -> Generator[str, None, None]:
     r"""hexdump_iter(s, width = 16, skip = True, hexii = False, begin = 0, style = None,
                     highlight = None, cyclic = False, groupsize=4, total = True) -> str generator
 
@@ -688,7 +690,7 @@ def hexdump_iter(fd, width=16, skip=True, hexii=False, begin=0, style=None,
     marker      = (style.get('marker') or (lambda s:s))('│')
 
     if not hexii:
-        def style_byte(by):
+        def style_byte(by: int) -> tuple[str, str]:
             hbyte = '%02x' % by
             b = packing._p8lu(by)
             abyte = chr(by) if isprint(b) else '·'
@@ -791,8 +793,8 @@ def style_byte(by):
         line = "%08x" % (begin + numb)
         yield line
 
-def hexdump(s, width=16, skip=True, hexii=False, begin=0, style=None,
-            highlight=None, cyclic=False, groupsize=4, total=True):
+def hexdump(s: bytes, width: int = 16, skip: bool = True, hexii: bool = False, begin: int = 0, style: dict = None,
+            highlight: Iterable = None, cyclic: bool = False, groupsize: int = 4, total: bool = True) -> str:
     r"""hexdump(s, width = 16, skip = True, hexii = False, begin = 0, style = None,
                 highlight = None, cyclic = False, groupsize=4, total = True) -> str
 
@@ -988,7 +990,7 @@ def hexdump(s, width=16, skip=True, hexii=False, begin=0, style=None,
                                   groupsize,
                                   total))
 
-def negate(value, width = None):
+def negate(value: int, width: int = None) -> int:
     """
     Returns the two's complement of 'value'.
     """
@@ -997,7 +999,7 @@ def negate(value, width = None):
     mask = ((1<<width)-1)
     return ((mask+1) - value) & mask
 
-def bnot(value, width=None):
+def bnot(value: int, width: int = None) -> int:
     """
     Returns the binary inverse of 'value'.
     """
@@ -1007,7 +1009,7 @@ def bnot(value, width=None):
     return mask ^ value
 
 @LocalNoarchContext
-def js_escape(data, padding=context.cyclic_alphabet[0:1], **kwargs):
+def js_escape(data: bytes, padding: bytes = context.cyclic_alphabet[0:1], **kwargs: dict[str, Any]) -> bytes:
     r"""js_escape(data, padding=context.cyclic_alphabet[0:1], endian = None, **kwargs) -> str
 
     Pack data as an escaped Unicode string for use in JavaScript's `unescape()` function
@@ -1049,7 +1051,7 @@ def js_escape(data, padding=context.cyclic_alphabet[0:1], **kwargs):
         return ''.join(f'%u{a:02x}{b:02x}' for a, b in iters.group(2, data))
 
 @LocalNoarchContext
-def js_unescape(s, **kwargs):
+def js_unescape(s: str, **kwargs: dict[str, Any]) -> bytes:
     r"""js_unescape(s, endian = None, **kwargs) -> bytes
 
     Unpack an escaped Unicode string from JavaScript's `escape()` function
@@ -1117,7 +1119,7 @@ def js_unescape(s, **kwargs):
 
     return b''.join(res)
 
-def tty_escape(s, lnext=b'\x16', dangerous=bytes(bytearray(range(0x20)))):
+def tty_escape(s: bytes, lnext: bytes = b'\x16', dangerous: bytes = bytes(bytearray(range(0x20)))) -> bytes:
     r"""tty_escape(s, lnext=b'\x16', dangerous=bytes(bytearray(range(0x20)))) -> bytes
 
     Escape data for terminal output. This is useful when sending data to a

pwnlib/util/getdents.py

@@ -65,7 +65,7 @@ class linux_dirent:
     d_type: Dtype
     d_name: str
 
-    def __init__(self, buf: bytes, is_dirent64: bool):
+    def __init__(self, buf: bytes, is_dirent64: bool) -> None:
         size_t = 8 if is_dirent64 else context.bytes
 
         self.d_ino = unpack(buf[0:size_t], size_t * 8)
@@ -82,13 +82,13 @@ def __init__(self, buf: bytes, is_dirent64: bool):
         self.d_name = d_name.split(b'\x00', 1)[0].decode('utf-8')
         self.d_type = Dtype(d_type)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.d_reclen
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.d_name
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f'{self.d_type.name:<8}{self.d_name}'