More type fixes (#1473)

pwndbg/commands/context.py

@@ -6,6 +6,7 @@ from collections import defaultdict
 from io import open
 from typing import DefaultDict
 from typing import Dict
+from typing import List
 
 import gdb
 
@@ -828,7 +829,7 @@ def context_args(with_banner=True, target=sys.stdout, width=None):
     return args
 
 
-last_signal = []
+last_signal: List[str] = []
 
 
 def save_signal(signal) -> None:

pwndbg/commands/cymbol.py

@@ -21,6 +21,7 @@ import os
 import subprocess
 import sys
 import tempfile
+from typing import Dict
 
 import gdb
 
@@ -46,7 +47,7 @@ cymbol_editor = pwndbg.gdblib.config.add_param(
 )
 
 # Remeber loaded symbols. This would be useful for 'remove-symbol-file'.
-loaded_symbols = {}
+loaded_symbols: Dict[str, str] = {}
 
 # Where generated symbol source files are saved.
 pwndbg_cachedir = pwndbg.lib.tempfile.cachedir("custom-symbols")

pwndbg/commands/heap.py

@@ -294,13 +294,15 @@ def malloc_chunk(addr, fake=False, verbose=False, simple=False) -> None:
         if not chunk.is_top_chunk and arena:
 
             bins_list = [
-                allocator.fastbins(arena.address) or {},
-                allocator.smallbins(arena.address) or {},
-                allocator.largebins(arena.address) or {},
-                allocator.unsortedbin(arena.address) or {},
+                allocator.fastbins(arena.address),
+                allocator.smallbins(arena.address),
+                allocator.largebins(arena.address),
+                allocator.unsortedbin(arena.address),
             ]
             if allocator.has_tcache():
                 bins_list.append(allocator.tcachebins(None))
+
+            bins_list = [x for x in bins_list if x is not None]
             no_match = True
             for bins in bins_list:
                 if bins.contains_chunk(chunk.real_size, chunk.address):
@@ -897,7 +899,7 @@ def try_free(addr) -> None:
         print(message.notice("__libc_free: Doing munmap_chunk"))
         return
 
-    errors_found = False
+    errors_found = 0
     returned_before_error = False
 
     # chunk doesn't overlap memory

pwndbg/commands/leakfind.py

@@ -5,6 +5,7 @@ Find a chain of leaks given some starting address.
 import argparse
 import queue
 from typing import Dict
+from typing import List
 
 import gdb
 
@@ -130,7 +131,7 @@ def leakfind(
     # We need to store both so that we can nicely create our leak chain.
     visited_map = {}
     visited_set = {int(address)}
-    address_queue = queue.Queue()
+    address_queue: queue.Queue[int] = queue.Queue()
     address_queue.put(int(address))
     depth = 0
     time_to_depth_increase = 0

pwndbg/disasm/__init__.py

@@ -4,6 +4,7 @@ address +/- a few instructions.
 """
 
 import collections
+from typing import DefaultDict
 
 import capstone
 import gdb
@@ -54,7 +55,7 @@ VariableInstructionSizeMax = {
     "mips": 8,
 }
 
-backward_cache = collections.defaultdict(lambda: None)
+backward_cache: DefaultDict = collections.defaultdict(lambda: None)
 
 
 @pwndbg.lib.memoize.reset_on_objfile

pwndbg/gdblib/elf.py

@@ -10,6 +10,7 @@ import ctypes
 import importlib
 import sys
 from collections import namedtuple
+from typing import List
 
 import gdb
 from elftools.elf.constants import SH_FLAGS

pwndbg/gdblib/qemu.py

@@ -60,8 +60,6 @@ def is_qemu_kernel() -> bool:
 @start
 @pwndbg.lib.memoize.reset_on_stop
 def root() -> Optional[Any]:
-    global binfmt_root
-
     if not is_qemu_usermode():
         return None
 

pwndbg/gdblib/regs.py

@@ -226,7 +226,7 @@ sys.modules[__name__] = module(__name__, "")
 @pwndbg.gdblib.events.cont
 @pwndbg.gdblib.events.stop
 def update_last() -> None:
-    M: module = sys.modules[__name__]
+    M = sys.modules[__name__]
     M.previous = M.last
     M.last = {k: M[k] for k in M.common}
     if pwndbg.gdblib.config.show_retaddr_reg:

pwndbg/gdblib/symbol.py

@@ -7,6 +7,7 @@ information available.
 """
 import os
 import re
+from typing import Dict
 
 import elftools.common.exceptions
 import elftools.elf.constants
@@ -63,7 +64,7 @@ if "/usr/lib/debug" not in _get_debug_file_directory():
     _add_debug_file_directory("/usr/lib/debug")
 
 
-_remote_files = {}
+_remote_files: Dict[str, str] = {}
 
 
 @pwndbg.gdblib.events.exit

pwndbg/gdblib/vmmap.py

@@ -8,6 +8,7 @@ system has /proc/$$/maps, which backs 'info proc mapping'.
 import bisect
 import os
 from typing import Any
+from typing import List
 from typing import Optional
 from typing import Tuple
 
@@ -30,10 +31,10 @@ import pwndbg.lib.memoize
 
 # List of manually-explored pages which were discovered
 # by analyzing the stack or register context.
-explored_pages = []
+explored_pages: List[pwndbg.lib.memory.Page] = []
 
 # List of custom pages that can be managed manually by vmmap_* commands family
-custom_pages = []
+custom_pages: List[pwndbg.lib.memory.Page] = []
 
 
 kernel_vmmap_via_pt = pwndbg.gdblib.config.add_param(
@@ -79,7 +80,7 @@ def is_corefile() -> bool:
 
 @pwndbg.lib.memoize.reset_on_start
 @pwndbg.lib.memoize.reset_on_stop
-def get() -> Tuple[pwndbg.lib.memory.Page, str]:
+def get() -> Tuple[pwndbg.lib.memory.Page, ...]:
     """
     Returns a tuple of `Page` objects representing the memory mappings of the
     target, sorted by virtual address ascending.
@@ -417,7 +418,7 @@ def proc_pid_maps():
 def kernel_vmmap_via_page_tables():
     import pt
 
-    retpages = []
+    retpages: List[pwndbg.lib.memory.Page] = []
 
     p = pt.PageTableDump()
     try:

pwndbg/heap/ptmalloc.py

@@ -2,6 +2,10 @@ import copy
 import importlib
 from collections import OrderedDict
 from enum import Enum
+from typing import Any
+from typing import Dict
+from typing import Set
+from typing import Union
 
 import gdb
 
@@ -71,7 +75,8 @@ class Bin:
 
 class Bins:
     def __init__(self, bin_type) -> None:
-        self.bins = OrderedDict()
+        # `typing.OrderedDict` requires Python 3.7
+        self.bins = OrderedDict()  # type: OrderedDict[Union[int, str], Bin]
         self.bin_type = bin_type
 
     # TODO: There's a bunch of bin-specific logic in here, maybe we should
@@ -1880,7 +1885,7 @@ class HeuristicHeap(GlibcMemoryAllocator):
                     pass
             elif pwndbg.gdblib.arch.current == "aarch64" and self.possible_page_of_symbols:
                 base_offset = self.possible_page_of_symbols.vaddr
-                regs = set()
+                regs: Set[str] = set()
                 found = False
                 for instr in __libc_free_instructions:
                     if found:
@@ -1899,7 +1904,7 @@ class HeuristicHeap(GlibcMemoryAllocator):
                         regs.add(instr.operands[0].str)
             elif pwndbg.gdblib.arch.current == "arm":
                 regs = {}
-                ldr = {}
+                ldr: Dict[str, Any] = {}
                 found = False
                 for instr in __libc_free_instructions:
                     if found:

pwndbg/ida.py

@@ -11,6 +11,8 @@ import time
 import traceback
 import xmlrpc.client
 from typing import Any
+from typing import Dict
+from typing import List
 from typing import Optional
 
 import gdb
@@ -252,16 +254,16 @@ def Jump(addr):
 @takes_address
 @pwndbg.lib.memoize.reset_on_objfile
 def Anterior(addr):
-    hexrays_prefix = "\x01\x04; "
+    hexrays_prefix = b"\x01\x04; "
     lines = []
     for i in range(10):
-        r = _ida.get_extra_cmt(addr, 0x3E8 + i)  # E_PREV
+        r: Optional[bytes] = _ida.get_extra_cmt(addr, 0x3E8 + i)  # E_PREV
         if not r:
             break
         if r.startswith(hexrays_prefix):
             r = r[len(hexrays_prefix) :]
         lines.append(r)
-    return "\n".join(lines)
+    return b"\n".join(lines)
 
 
 @withIDA
@@ -281,7 +283,7 @@ def GetBptEA(i):
     return _ida.get_bpt_ea(i)
 
 
-_breakpoints = []
+_breakpoints: List[gdb.Breakpoint] = []
 
 
 @pwndbg.gdblib.events.cont
@@ -488,7 +490,7 @@ class IDC:
 
     def __init__(self) -> None:
         if available():
-            data = _ida.eval(self.query)
+            data: Dict = _ida.eval(self.query)
             self.__dict__.update(data)
 
 

pwndbg/lib/config.py

@@ -1,6 +1,8 @@
 import collections
 from functools import total_ordering
+from typing import Callable
 from typing import DefaultDict
+from typing import Dict
 from typing import List
 
 import gdb
@@ -127,7 +129,7 @@ class Parameter:
 class Config:
     def __init__(self) -> None:
         self.params: Dict[str, Parameter] = {}
-        self.triggers: DefaultDict[str, Callable] = collections.defaultdict(lambda: [])
+        self.triggers: DefaultDict[str, List[Callable]] = collections.defaultdict(lambda: [])
 
     def add_param(
         self,

pwndbg/lib/kernel/kconfig.py

@@ -50,7 +50,7 @@ class Kconfig(UserDict):
 
         raise KeyError(f"Key {name} not found")
 
-    def __contains__(self, name: str) -> bool:
+    def __contains__(self, name: str) -> bool:  # type: ignore[override]
         return self.get_key(name) is not None
 
     def __getattr__(self, name: str):

pwndbg/lib/which.py

@@ -62,7 +62,9 @@ def which(name: str, all: bool = False):
     try:
         path = os.environ["PATH"]
     except KeyError:
-        log.exception("Environment variable $PATH is not set")
+        print("Environment variable $PATH is not set")
+        return None
+
     for p in path.split(os.pathsep):
         p = os.path.join(p, name)
         if os.access(p, os.X_OK):

pwndbg/profiling.py

@@ -14,7 +14,7 @@ def init(p: cProfile.Profile, _start_time: Optional[float]) -> None:
 class Profiler:
     def __init__(self, p: cProfile.Profile) -> None:
         self._profiler = p
-        self._start_time = None
+        self._start_time: Optional[float] = None
 
     def print_time_elapsed(self):
         assert self._start_time is not None

pyproject.toml

@@ -11,33 +11,31 @@ check_untyped_defs = true
 pretty = true
 show_error_codes = true
 incremental = false
-follow_imports = "skip"
-exclude = [
-    #"^pwndbg/commands",
-    "^pwndbg/heap",
-    #'^pwndbg/[^/]*\.py',
-    '^pwndbg/ida.py',
-    '^pwndbg/gdblib/vmmap.py',
-]
 disable_error_code = [
     "misc",
+    # No type stubs for capstone, unicorn, pwnlib, and elftools
     "import",
-    "name-defined",
+    # Lots of dynamic attribute access
     "attr-defined",
-    # "var-annotated",
-    # "index",
-    # "arg-type",
+    # Issues with indexing Module
+    "index",
+    # https://github.com/python/mypy/issues/6232
     "assignment",
-    # "operator",
-    "override",
-    # "call-overload",
-    "union-attr",
-    # "str-format",
-    # "call-arg",
-    # "return-value",
-    # "has-type",
-    # "no-redef",
+    # Issues with @property
+    "operator",
+]
+
+[[tool.mypy.overrides]]
+module = [
+   # Capstone constants
+   "pwndbg.disasm.*",
+   "pwndbg.commands.nearpc",
+   # Module fields
+   "pwndbg.gdblib.typeinfo",
+   "pwndbg.gdblib.elf",
+   "pwndbg.auxv",
 ]
+disable_error_code = ["name-defined"]
 
 [tool.isort]
 profile = "black"