Add type hints to timeout, term.text, hashes, and more (#2723)

* Add typehints to timeout module

Co-authored-by: Copilot <copilot@github.com>

* Add typehints to device class

* Add typehints to PwnlibException

* Add general typehints for term.text module

Since the available module attributes are handled dynamically like `text.green()` or `text.green_on_white()` or `text.underline_red()`, it's not feasible to type out all possible combinations.

Just tell the type checker that all module attibutes have the same type and thus avoid the `error: Module has no attribute "bold_blue"  [attr-defined]` errors.

* Add typehints for config module

* Add typehints for hashes module

The guaranteed algorithms in hashlib are pretty stable and we can generate the list out proper type checking and auto complete.

Still add a check in CI that notifies us if there are changes in the future.

* Add atexception and atexit types

atexception.unregister was not documented correctly since the function expects an identifier returned by `atexception.register` instead of the handler function.

* Fix toplevel type hints in shellcraft.internal

* Fix toplevel type hints in term.key

---------

Co-authored-by: Copilot <copilot@github.com>

Author: peace-maker

.github/workflows/ci.yml

@@ -133,7 +133,9 @@ jobs:
         pip install 'unicorn==2.1.2'
 
     - name: Sanity checks
-      run:  PWNLIB_NOTERM=1 python -bb -c 'from pwn import *; print(pwnlib.term.term_mode)'
+      run: |
+        PWNLIB_NOTERM=1 python -bb -c 'from pwn import *; print(pwnlib.term.term_mode)'
+        python -c "import hashlib; assert hashlib.algorithms_guaranteed == {'sha256', 'blake2s', 'sha3_384', 'shake_256', 'sha1', 'md5', 'shake_128', 'sha224', 'sha512', 'blake2b', 'sha3_512', 'sha3_224', 'sha3_256', 'sha384'}"
 
     - name: Install documentation dependencies
       run: pip install -r docs/requirements.txt

mypy-baseline.txt

@@ -4,9 +4,7 @@ pwn/toplevel.py:0: error: Incompatible import of "size" (imported name has type
 pwn/toplevel.py:0: error: Incompatible import of "size" (imported name has type "Callable[[Any, Any, Any], Any]", local name has type "str")  [assignment]
 pwn/toplevel.py:0: error: Incompatible import of "size" (imported name has type "int", local name has type "str")  [assignment]
 pwnlib/abi.py:0: error: Need type annotation for "register_arguments" (hint: "register_arguments: list[<type>] = ...")  [var-annotated]
-pwnlib/asm.py:0: error: Module "pwnlib.util.hashes" has no attribute "sha1sumhex"  [attr-defined]
 pwnlib/asm.py:0: error: Need type annotation for "util_versions"  [var-annotated]
-pwnlib/atexception.py:0: error: Function "_run_handlers" could always be true in boolean context  [truthy-function]
 pwnlib/commandline/cyclic.py:0: error: "Callable[[Any, Any, Any, Any], Any]" has no attribute "add_argument"  [attr-defined]
 pwnlib/commandline/cyclic.py:0: error: "Callable[[Any, Any, Any, Any], Any]" has no attribute "add_argument"  [attr-defined]
 pwnlib/commandline/cyclic.py:0: error: Incompatible types in assignment (expression has type "_MutuallyExclusiveGroup", variable has type "Callable[[Any, Any, Any, Any], Any]")  [assignment]
@@ -37,21 +35,6 @@ pwnlib/filesystem/path.py:0: error: "classmethod" used with a non-method  [misc]
 pwnlib/filesystem/path.py:0: error: "type[Path]" has no attribute "mkdtemp"  [attr-defined]
 pwnlib/filesystem/path.py:0: error: "type[Path]" has no attribute "mktemp"  [attr-defined]
 pwnlib/gdb_api_bridge.py:0: error: Name "socket_path" is not defined  [name-defined]
-pwnlib/libcdb.py:0: error: Module "pwnlib.util.hashes" has no attribute "md5filehex"  [attr-defined]
-pwnlib/libcdb.py:0: error: Module "pwnlib.util.hashes" has no attribute "sha1filehex"  [attr-defined]
-pwnlib/libcdb.py:0: error: Module "pwnlib.util.hashes" has no attribute "sha256filehex"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_blue"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_blue"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_blue"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_green"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_red"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_red"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_yellow"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "bold_yellow"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "magenta"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "on_red"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "on_red"  [attr-defined]
-pwnlib/log.py:0: error: Module has no attribute "on_red"  [attr-defined]
 pwnlib/log.py:0: error: Need type annotation for "_one_time_infos" (hint: "_one_time_infos: set[<type>] = ...")  [var-annotated]
 pwnlib/log.py:0: error: Need type annotation for "_one_time_warnings" (hint: "_one_time_warnings: set[<type>] = ...")  [var-annotated]
 pwnlib/memleak.py:0: error: Module "pwnlib.util.packing" has no attribute "_p8lu"  [attr-defined]
@@ -86,22 +69,11 @@ pwnlib/shellcraft/__init__.py:0: error: Argument 1 to "append" of "list" has inc
 pwnlib/shellcraft/__init__.py:0: error: Need type annotation for "_templates" (hint: "_templates: list[<type>] = ...")  [var-annotated]
 pwnlib/term/__init__.py:0: error: Module has no attribute "height"  [attr-defined]
 pwnlib/term/__init__.py:0: error: Module has no attribute "width"  [attr-defined]
-pwnlib/term/key.py:0: error: Need type annotation for "_kbuf" (hint: "_kbuf: list[<type>] = ...")  [var-annotated]
-pwnlib/term/readline.py:0: error: Module has no attribute "reverse"  [attr-defined]
 pwnlib/term/readline.py:0: error: Need type annotation for "history" (hint: "history: list[<type>] = ...")  [var-annotated]
 pwnlib/term/readline.py:0: error: Need type annotation for "search_results" (hint: "search_results: list[<type>] = ...")  [var-annotated]
 pwnlib/term/term.py:0: error: Need type annotation for "on_winch" (hint: "on_winch: list[<type>] = ...")  [var-annotated]
+pwnlib/timeout.py:0: error: Incompatible return value type (got "float | TimeoutDefault", expected "float")  [return-value]
 pwnlib/tubes/process.py:0: error: Cannot assign to a type  [misc]
 pwnlib/tubes/process.py:0: error: Incompatible types in assignment (expression has type "PTY", variable has type "type[PTY]")  [assignment]
-pwnlib/tubes/process.py:0: error: Module "pwnlib.util.hashes" has no attribute "sha256file"  [attr-defined]
 pwnlib/tubes/tube.py:0: error: Argument 1 to "make_wrapper" has incompatible type "function"; expected "tube"  [arg-type]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "blue"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "blue"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "gray"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "gray"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "green"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "has_gray"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "has_gray"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "red"  [attr-defined]
-pwnlib/util/fiddling.py:0: error: Module has no attribute "red"  [attr-defined]
 pwnlib/util/iters.py:0: error: Name "pairwise" already defined (possibly by an import)  [no-redef]

pwnlib/atexception.py

@@ -5,16 +5,21 @@
 import sys
 import threading
 import traceback
+from types import TracebackType
+from typing import Any, Callable, ParamSpec, TypeVar
 
 from pwnlib.context import context
 
 __all__ = ['register', 'unregister']
 
 _lock = threading.Lock()
 _ident = 0
-_handlers = {}
+_handlers: dict[int, tuple[Callable[..., Any], Any, Any, dict[str, Any]]] = {}
 
-def register(func, *args, **kwargs):
+_P = ParamSpec('_P')
+_R = TypeVar('_R')
+
+def register(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> int:
     """register(func, *args, **kwargs)
 
     Registers a function to be called when an unhandled exception occurs.  The
@@ -40,7 +45,8 @@ def handler():
     actual exception-handler.  The exception-handler can then be unregistered
     with::
 
-      atexception.unregister(handler)
+      ident = atexception.register(handler)
+      atexception.unregister(ident)
 
     This function is thread safe.
 
@@ -52,16 +58,16 @@ def handler():
     _handlers[ident] = (func, args, kwargs, vars(context))
     return ident
 
-def unregister(func):
-    """unregister(func)
+def unregister(ident: int) -> None:
+    """unregister(ident)
 
-    Remove `func` from the collection of registered functions.  If `func` isn't
-    registered this is a no-op.
+    Remove the exception-handler identified by `ident` from the list of registered
+    handlers.  If `ident` isn't registered this is a no-op.
     """
-    if func in _handlers:
-        del _handlers[func]
+    if ident in _handlers:
+        del _handlers[ident]
 
-def _run_handlers():
+def _run_handlers() -> None:
     """_run_handlers()
 
     Run registered handlers.  They run in the reverse order of which they were
@@ -83,20 +89,20 @@ def _run_handlers():
             # extract the current exception and rewind the traceback to where it
             # originated
             typ, val, tb = sys.exc_info()
-            traceback.print_exception(typ, val, tb.tb_next)
+            traceback.print_exception(typ, value=val, tb=tb.tb_next if tb else None)
 
 # we rely on the existing excepthook to print exceptions
 _oldhook = getattr(sys, 'excepthook', None)
 
-def _newhook(typ, val, tb):
+def _newhook(typ: type[BaseException], val: BaseException, tb: TracebackType | None) -> None:
     """_newhook(typ, val, tb)
 
     Our excepthook replacement.  First the original hook is called to print the
     exception, then each handler is called.
     """
     if _oldhook:
         _oldhook(typ, val, tb)
-    if _run_handlers:
+    if _handlers:
         _run_handlers()
 
 sys.excepthook = _newhook

pwnlib/atexit.py

@@ -11,16 +11,20 @@
 import threading
 import traceback
 import atexit as std_atexit
+from typing import Any, Callable, ParamSpec, TypeVar
 
 from pwnlib.context import context
 
 __all__ = ['register', 'unregister']
 
 _lock = threading.Lock()
 _ident = 0
-_handlers = {}
+_handlers: dict[int, tuple[Callable[..., Any], Any, Any, dict[str, Any]]] = {}
 
-def register(func, *args, **kwargs):
+_P = ParamSpec('_P')
+_R = TypeVar('_R')
+
+def register(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> int:
     """register(func, *args, **kwargs)
 
     Registers a function to be called on program termination.  The function will
@@ -44,7 +48,8 @@ def handler():
     Notice however that this will bind ``handler`` to the identifier and not the
     actual exit-handler.  The exit-handler can then be unregistered with::
 
-      atexit.unregister(handler)
+      ident = atexit.register(handler)
+      atexit.unregister(ident)
 
     This function is thread safe.
 
@@ -56,7 +61,7 @@ def handler():
     _handlers[ident] = (func, args, kwargs, vars(context))
     return ident
 
-def unregister(ident):
+def unregister(ident: int) -> None:
     """unregister(ident)
 
     Remove the exit-handler identified by `ident` from the list of registered
@@ -65,7 +70,7 @@ def unregister(ident):
     if ident in _handlers:
         del _handlers[ident]
 
-def _run_handlers():
+def _run_handlers() -> None:
     """_run_handlers()
 
     Run registered exit-handlers.  They run in the reverse order of which they
@@ -87,13 +92,6 @@ def _run_handlers():
             # extract the current exception and rewind the traceback to where it
             # originated
             typ, val, tb = sys.exc_info()
-            traceback.print_exception(typ, val, tb.tb_next)
-
-# if there's already an exitfunc registered be sure to run that too
-if hasattr(sys, "exitfunc"):
-    register(sys.exitfunc)
+            traceback.print_exception(typ, value=val, tb=tb.tb_next if tb else None)
 
-if sys.version_info[0] < 3:
-    sys.exitfunc = _run_handlers
-else:
-    std_atexit.register(_run_handlers)
+std_atexit.register(_run_handlers)

pwnlib/config.py

@@ -13,7 +13,7 @@
 file in library mode, invoke :func:`.config.initialize`.
 
 The ``context`` section supports complex types, at least as far as is
-supported by ``pwnlib.util.safeeval.expr``.
+supported by :func:`pwnlib.util.safeeval.expr`.
 
 ::
 
@@ -33,10 +33,13 @@
 """
 import configparser
 import os
+from typing import Any, Callable
 
-registered_configs = {}
+ConfigHandler = Callable[[dict[str, Any]], None]
 
-def register_config(section, function):
+registered_configs: dict[str, ConfigHandler] = {}
+
+def register_config(section: str, function: ConfigHandler) -> None:
     """Registers a configuration section.
 
     Arguments:
@@ -46,7 +49,7 @@ def register_config(section, function):
     """
     registered_configs[section] = function
 
-def initialize():
+def initialize() -> None:
     """Read the configuration files."""
     from pwnlib.log import getLogger
     log = getLogger(__name__)

pwnlib/device.py

@@ -1,15 +1,15 @@
 class Device:
-    arch = None
-    bits = None
-    endian = None
-    serial = None
-    os = None
+    arch: str | None = None
+    bits: int | None = None
+    endian: str | None = None
+    serial: str | None = None
+    os: str | None = None
 
-    def __init__(self, serial=None):
+    def __init__(self, serial: str | None = None):
         self.serial = serial
 
-    def __str__(self):
-        return self.serial
+    def __str__(self) -> str:
+        return self.serial or "<no serial>"
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         return self.serial == other or self.serial == str(other)

pwnlib/exception.py

@@ -1,21 +1,27 @@
 __all__ = ['PwnlibException']
 import sys
 import traceback
+from types import TracebackType
 
+OptExcInfo = tuple[
+    type[BaseException],
+    BaseException,
+    TracebackType,
+] | tuple[None, None, None]
 
 class PwnlibException(Exception):
     '''Exception thrown by :func:`pwnlib.log.error`.
 
     Pwnlib functions that encounters unrecoverable errors should call the
     :func:`pwnlib.log.error` function instead of throwing this exception directly.'''
-    def __init__(self, msg, reason = None, exit_code = None):
+    def __init__(self, msg: str, reason: OptExcInfo | None = None, exit_code: int | None = None):
         '''bar'''
         Exception.__init__(self, msg)
         self.reason = reason
         self.exit_code = exit_code
         self.message = msg
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         s = 'PwnlibException: %s' % self.message
         if self.reason:
             s += '\nReason:\n'

pwnlib/shellcraft/internal.py

@@ -1,11 +1,16 @@
+from __future__ import annotations
 import os
 import sys
+from typing import TYPE_CHECKING
 
 from pwnlib.context import context
 
 __all__ = ['make_function']
 
-loaded = {}
+if TYPE_CHECKING:
+    from mako.template import Template
+
+loaded: dict[str, Template] = {}
 lookup = None
 def init_mako():
     global lookup, render_global