From 2296999ed9c8004557a4098de59ef84ce2fb06f1 Mon Sep 17 00:00:00 2001
From: syheliel <45957390+syheliel@users.noreply.github.com>
Date: Mon, 12 Sep 2022 21:50:06 +0800
Subject: [PATCH] add type for `./pwndbg/lib` (#1135)

* add type for `./pwndbg/lib`

* add more type hints

Co-authored-by: syheliel <syheliel@gmail.com>
---
 pwndbg/lib/abi.py        | 10 ++++++----
 pwndbg/lib/android.py    |  2 +-
 pwndbg/lib/arch.py       |  6 +++---
 pwndbg/lib/elftypes.py   |  3 ++-
 pwndbg/lib/funcparser.py |  2 +-
 pwndbg/lib/gcc.py        |  9 +++++----
 pwndbg/lib/memoize.py    | 17 +++++++++--------
 pwndbg/lib/memory.py     |  2 +-
 pwndbg/lib/stdio.py      |  3 ++-
 pwndbg/lib/tips.py       |  3 ++-
 pwndbg/lib/version.py    |  2 +-
 11 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/pwndbg/lib/abi.py b/pwndbg/lib/abi.py
index fb76d699..bc58f7a9 100644
--- a/pwndbg/lib/abi.py
+++ b/pwndbg/lib/abi.py
@@ -1,3 +1,5 @@
+from typing import List
+
 import pwndbg.gdblib.arch
 
 
@@ -8,7 +10,7 @@ class ABI:
 
     #: List or registers which should be filled with arguments before
     #: spilling onto the stack.
-    register_arguments = []
+    register_arguments: List[str] = []
 
     #: Minimum alignment of the stack.
     #: The value used is min(context.bytes, stack_alignment)
@@ -29,7 +31,7 @@ class ABI:
         self.stack_minimum = minimum
 
     @staticmethod
-    def default():
+    def default():  # type: () -> ABI
         return {
             (32, "i386", "linux"): linux_i386,
             (64, "x86-64", "linux"): linux_amd64,
@@ -42,7 +44,7 @@ class ABI:
         }[(8 * pwndbg.gdblib.arch.ptrsize, pwndbg.gdblib.arch.current, "linux")]
 
     @staticmethod
-    def syscall():
+    def syscall():  # type: () -> ABI
         return {
             (32, "i386", "linux"): linux_i386_syscall,
             (64, "x86-64", "linux"): linux_amd64_syscall,
@@ -55,7 +57,7 @@ class ABI:
         }[(8 * pwndbg.gdblib.arch.ptrsize, pwndbg.gdblib.arch.current, "linux")]
 
     @staticmethod
-    def sigreturn():
+    def sigreturn():  # type: () -> SigreturnABI
         return {
             (32, "i386", "linux"): linux_i386_sigreturn,
             (64, "x86-64", "linux"): linux_amd64_sigreturn,
diff --git a/pwndbg/lib/android.py b/pwndbg/lib/android.py
index 1dd0501f..42188f7b 100644
--- a/pwndbg/lib/android.py
+++ b/pwndbg/lib/android.py
@@ -79,7 +79,7 @@ KNOWN_AIDS = {
 }
 
 
-def aid_name(uid):
+def aid_name(uid):  # types: (int) -> str
     if uid in KNOWN_AIDS:
         return KNOWN_AIDS[uid]
 
diff --git a/pwndbg/lib/arch.py b/pwndbg/lib/arch.py
index 6d102ba8..7adcc1b7 100644
--- a/pwndbg/lib/arch.py
+++ b/pwndbg/lib/arch.py
@@ -35,8 +35,8 @@ class Arch:
     def unpack(self, data: bytearray) -> int:
         return struct.unpack(self.fmt, data)[0]
 
-    def signed(self, integer):
-        return self.unpack(self.pack(integer), signed=True)
+    def signed(self, integer: int) -> int:
+        return self.unpack(self.pack(integer), signed=True)  # type: ignore
 
-    def unsigned(integer):
+    def unsigned(self, integer: int) -> int:
         return self.unpack(self.pack(integer))
diff --git a/pwndbg/lib/elftypes.py b/pwndbg/lib/elftypes.py
index 23542c44..b5f9fe8c 100644
--- a/pwndbg/lib/elftypes.py
+++ b/pwndbg/lib/elftypes.py
@@ -28,6 +28,7 @@
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #
 import ctypes
+from typing import Dict
 
 import pwndbg.gdblib.ctypes
 
@@ -81,7 +82,7 @@ AT_CONSTANTS = {
     35: "AT_L1D_CACHESHAPE",
     36: "AT_L2_CACHESHAPE",
     37: "AT_L3_CACHESHAPE",
-}
+}  # type: Dict[int,str]
 
 
 class constants:
diff --git a/pwndbg/lib/funcparser.py b/pwndbg/lib/funcparser.py
index 528392a4..6ea4703b 100644
--- a/pwndbg/lib/funcparser.py
+++ b/pwndbg/lib/funcparser.py
@@ -1,6 +1,6 @@
 import collections
 
-from pycparser import CParser
+from pycparser import CParser  # type: ignore
 from pycparser import c_ast
 
 
diff --git a/pwndbg/lib/gcc.py b/pwndbg/lib/gcc.py
index 65cd538a..8627f190 100644
--- a/pwndbg/lib/gcc.py
+++ b/pwndbg/lib/gcc.py
@@ -6,13 +6,14 @@ GCC and any flags it should be executed with.
 import glob
 import os
 import platform
+from typing import List
 
 from pwndbg.lib.arch import Arch
 
 printed_message = False
 
 
-def which(arch: Arch):
+def which(arch):  # type: (Arch) -> List[str]
     gcc = _which_binutils("g++", arch)
 
     if not gcc:
@@ -21,9 +22,9 @@ def which(arch: Arch):
             printed_message = True
             print("Can't find appropriate GCC, using default version")
 
-        if ptrsize == 32:
+        if arch.ptrsize == 32:
             return ["g++", "-m32"]
-        elif ptrsize == 64:
+        elif arch.ptrsize == 64:
             return ["g++", "-m32"]
 
     return [gcc] + _flags(arch.name)
@@ -72,7 +73,7 @@ def _which_binutils(util, arch, **kwargs):
                     return res[0]
 
 
-def _flags(arch_name):
+def _flags(arch_name):  # type: (str) -> List[str]
     if arch_name == "i386":
         return ["-m32"]
     if arch_name.endswith("x86-64"):
diff --git a/pwndbg/lib/memoize.py b/pwndbg/lib/memoize.py
index 535bbbcf..da1b35d0 100644
--- a/pwndbg/lib/memoize.py
+++ b/pwndbg/lib/memoize.py
@@ -8,6 +8,7 @@ import functools
 import sys
 from typing import Any
 from typing import Callable
+from typing import List
 
 try:
     # Python >= 3.10
@@ -75,7 +76,7 @@ class forever(memoize):
     Memoizes forever - for a pwndbg session or until `_reset` is called explicitly.
     """
 
-    caches = []
+    caches = []  # type: List[forever]
 
     @staticmethod
     def _reset():
@@ -84,7 +85,7 @@ class forever(memoize):
 
 
 class reset_on_stop(memoize):
-    caches = []
+    caches = []  # type: List[reset_on_stop]
     kind = "stop"
 
     @staticmethod
@@ -96,7 +97,7 @@ class reset_on_stop(memoize):
 
 
 class reset_on_prompt(memoize):
-    caches = []
+    caches = []  # type: List[reset_on_prompt]
     kind = "prompt"
 
     @staticmethod
@@ -108,7 +109,7 @@ class reset_on_prompt(memoize):
 
 
 class reset_on_exit(memoize):
-    caches = []
+    caches = []  # type: List[reset_on_exit]
     kind = "exit"
 
     @staticmethod
@@ -120,7 +121,7 @@ class reset_on_exit(memoize):
 
 
 class reset_on_objfile(memoize):
-    caches = []
+    caches = []  # type: List[reset_on_objfile]
     kind = "objfile"
 
     @staticmethod
@@ -132,7 +133,7 @@ class reset_on_objfile(memoize):
 
 
 class reset_on_start(memoize):
-    caches = []
+    caches = []  # type: List[reset_on_start]
     kind = "start"
 
     @staticmethod
@@ -144,7 +145,7 @@ class reset_on_start(memoize):
 
 
 class reset_on_cont(memoize):
-    caches = []
+    caches = []  # type: List[reset_on_cont]
     kind = "cont"
 
     @staticmethod
@@ -156,7 +157,7 @@ class reset_on_cont(memoize):
 
 
 class while_running(memoize):
-    caches = []
+    caches = []  # type: List[while_running]
     kind = "running"
     caching = False
 
diff --git a/pwndbg/lib/memory.py b/pwndbg/lib/memory.py
index a657d1a1..0fab298e 100644
--- a/pwndbg/lib/memory.py
+++ b/pwndbg/lib/memory.py
@@ -42,7 +42,7 @@ def page_size_align(address: int) -> int:
     return round_up(address, PAGE_SIZE)
 
 
-def page_offset(address):
+def page_offset(address: int) -> int:
     return address & (PAGE_SIZE - 1)
 
 
diff --git a/pwndbg/lib/stdio.py b/pwndbg/lib/stdio.py
index decd7184..c4cf0799 100644
--- a/pwndbg/lib/stdio.py
+++ b/pwndbg/lib/stdio.py
@@ -4,10 +4,11 @@ which prevent output from appearing on-screen inside of certain event handlers.
 """
 
 import sys
+from typing import *  # noqa note: TextIO is not abaliable in low python version
 
 
 class Stdio:
-    queue = []
+    queue = []  # type: List[Tuple[TextIO, TextIO, TextIO]]
 
     def __enter__(self, *a, **kw):
         self.queue.append((sys.stdin, sys.stdout, sys.stderr))
diff --git a/pwndbg/lib/tips.py b/pwndbg/lib/tips.py
index dc8c2035..eecdc267 100644
--- a/pwndbg/lib/tips.py
+++ b/pwndbg/lib/tips.py
@@ -1,4 +1,5 @@
 from random import choice
+from typing import List
 
 TIPS = [
     # GDB hints
@@ -22,7 +23,7 @@ TIPS = [
     "Want to display each context panel in a separate tmux window? See https://github.com/pwndbg/pwndbg/blob/dev/FEATURES.md#splitting--layouting-context",
     "The $heap_base GDB variable can be used to refer to the starting address of the heap after running the `heap` command",
     "Use the `errno` (or `errno <number>`) command to see the name of the last or provided (libc) error",
-]
+]  # type: List[str]
 
 
 def get_tip_of_the_day() -> str:
diff --git a/pwndbg/lib/version.py b/pwndbg/lib/version.py
index 4ddb7d97..bc5fce55 100644
--- a/pwndbg/lib/version.py
+++ b/pwndbg/lib/version.py
@@ -2,7 +2,7 @@ import os
 import subprocess
 
 
-def build_id():
+def build_id():  # type: () -> str
     """
     Returns pwndbg commit id if git is available.
     """