More type fixes (#1506)

* More type fixes

* Properly handle missing imports during typechecking

* Fix endian type

pwndbg/commands/context.py

@@ -806,7 +806,7 @@ def context_backtrace(with_banner=True, target=sys.stdout, width=None):
         prefix = bt_prefix if frame == this_frame else " " * len(bt_prefix)
         prefix = " %s" % c.prefix(prefix)
         addrsz = c.address(pwndbg.ui.addrsz(frame.pc()))
-        symbol = c.symbol(pwndbg.gdblib.symbol.get(frame.pc()))
+        symbol = c.symbol(pwndbg.gdblib.symbol.get(int(frame.pc())))
         if symbol:
             addrsz = addrsz + " " + symbol
         line = map(str, (prefix, c.frame_label("%s%i" % (backtrace_frame_label, i)), addrsz))

pwndbg/commands/slab.py

@@ -1,6 +1,8 @@
 import argparse
 import sys
 from typing import Iterator
+from typing import List
+from typing import Union
 
 import gdb
 from tabulate import tabulate
@@ -125,7 +127,9 @@ def _rx(val: int) -> str:
     return C.red(hex(val))
 
 
-def print_slab(slab: gdb.Value, freelist: Iterator[int], indent, verbose, is_partial):
+def print_slab(
+    slab: gdb.Value, freelist: Union[Iterator[int], List[int]], indent, verbose, is_partial
+):
     page_address = int(slab.address)
     virt_address = pwndbg.gdblib.kernel.page_to_virt(page_address)
     indent.print(f"- {C.green('Slab')} @ {_yx(virt_address)} [{_rx(page_address)}]:")

pwndbg/ida.py

@@ -130,7 +130,7 @@ class withIDA:
         self.fn = fn
         functools.update_wrapper(self, fn)
 
-    def __call__(self, *args: int, **kwargs: Any) -> Optional[Any]:
+    def __call__(self, *args: Any, **kwargs: Any) -> Optional[Any]:
         if _ida is None:
             init_ida_rpc_client()
         if _ida is not None:
@@ -201,9 +201,8 @@ def remote(function) -> None:
 
 @pwndbg.lib.memoize.reset_on_objfile
 def base():
-    segaddr = _ida.get_next_seg(0)
+    segaddr: int = _ida.get_next_seg(0)
-
+    base: int = _ida.get_fileregion_offset(segaddr)
-    base = _ida.get_fileregion_offset(segaddr)
 
     return segaddr - base
 
@@ -386,7 +385,7 @@ def ArgCount(address) -> None:
 
 
 @withIDA
-def SaveBase(path):
+def SaveBase(path: str):
     return _ida.save_database(path)
 
 

pwndbg/lib/arch.py

@@ -1,9 +1,11 @@
 import struct
 import sys
 
+from typing_extensions import Literal
+
 
 class Arch:
-    def __init__(self, arch_name: str, ptrsize: int, endian: str) -> None:
+    def __init__(self, arch_name: str, ptrsize: int, endian: Literal["little", "big"]) -> None:
         self.name = arch_name
         # TODO: `current` is the old name for the arch name, and it's now an
         # alias for `name`. It's used throughout the codebase, do we want to

pyproject.toml

@@ -18,8 +18,6 @@ show_error_codes = true
 incremental = false
 disable_error_code = [
     "misc",
-    # No type stubs for capstone, unicorn, pwnlib, and elftools
-    "import",
     # Lots of dynamic attribute access
     "attr-defined",
     # Issues with indexing Module
@@ -42,6 +40,10 @@ module = [
 ]
 disable_error_code = ["name-defined"]
 
+[[tool.mypy.overrides]]
+module = ["capstone.*", "unicorn.*", "pwnlib.*", "elftools.*", "ipdb.*", "r2pipe", "pt"]
+ignore_missing_imports = true
+
 [tool.isort]
 profile = "black"
 force_single_line = true

requirements.txt

@@ -6,5 +6,7 @@ pyelftools==0.29
 Pygments==2.13.0
 ROPGadget==7.2
 tabulate==0.8.10
+typing-extensions==4.3.0; python_version >= '3.7'
+typing-extensions==4.1.1; python_version < '3.7'
 unicorn==2.0.1.post1; python_version >= '3.7'
 unicorn==2.0.0rc7; python_version < '3.7'