Support x64dbg and set debugger paths through context

pwnlib/commandline/checksec.py

@@ -2,9 +2,9 @@
 from __future__ import division
 
 import argparse
-import sys
 
 from pwn import *
+from pwnlib.binary import Binary
 from pwnlib.commandline import common
 
 parser = common.parser_commands.add_parser(
@@ -34,7 +34,7 @@ def main(args):
 
     for f in files:
         try:
-            e = ELF(f)
+            e = Binary.from_path(f)
         except Exception as e:
             print("{name}: {error}".format(name=f, error=e))
 

pwnlib/context/__init__.py

@@ -364,6 +364,9 @@ class ContextType(object):
         'endian': 'little',
         'gdbinit': "",
         'gdb_binary': "",
+        'windbg_binary': "",
+        'x64dbg_binary': "",
+        'debugger_selection': "",
         'kernel': None,
         'local_libcdb': "/var/lib/libc-database",
         'log_level': logging.INFO,
@@ -453,6 +456,9 @@ class ContextType(object):
 
     valid_signed = sorted(signednesses)
 
+    #: Valid values for :attr:`debugger_selection`
+    debugger_choices = ['gdb', 'windbg', 'x64dbg']
+
     def __init__(self, **kwargs):
         """
         Initialize the ContextType structure.
@@ -1564,6 +1570,54 @@ def gdb_binary(self, value):
         """
         return str(value)
 
+    @_validator
+    def windbg_binary(self, value):
+        """Path to the binary that is used when running WinDBG locally.
+
+        This is useful when you have multiple versions of WinDBG installed or the WinDBG binary is
+        called something different.
+
+        If set to an empty string, pwntools will try to search for a reasonable WinDBG binary from 
+        the path.
+
+        Default value is ``""``.
+        """
+        return str(value)
+
+    @_validator
+    def x64dbg_binary(self, value):
+        """Path to the binary that is used when running x64dbg locally.
+
+        Should be set to the x96dbg.exe launcher binary to handle 32-bit and 64-bit binaries.
+
+        This is useful when you have multiple versions of x64dbg installed or the x64dbg binary is
+        called something different.
+
+        If set to an empty string, pwntools will try to search for a reasonable x64dbg binary from 
+        the path.
+
+        Default value is ``""``.
+        """
+        return str(value)
+
+    @_validator
+    def debugger_selection(self, value):
+        """Type of debugger to use when running locally.
+
+        Possible values are:
+
+        - ``gdb``: Use GDB as the debugger.
+        - ``windbg``: Use WinDBG as the debugger.
+        - ``x64dbg``: Use x64dbg as the debugger.
+
+        Defaults to ``windbg`` on Windows and ``gdb`` otherwise.
+        
+        Default value is ``""``.
+        """
+        if value not in self.debugger_choices:
+            raise AttributeError("debugger_selection must be one of %r" % sorted(self.debugger_choices))
+        return str(value)
+
     @_validator
     def cyclic_alphabet(self, alphabet):
         """Cyclic alphabet.

pwnlib/data/templates/pwnup.mako

@@ -5,16 +5,18 @@ import sys
 
 from pwnlib.context import context as ctx
 from pwnlib.elf.elf import ELF
+from pwnlib.binary import Binary
 from pwnlib.util.sh_string import sh_string
 from elftools.common.exceptions import ELFError
+from pefile import PEFormatError
 
 argv = list(sys.argv)
 argv[0] = os.path.basename(argv[0])
 
 try:
     if binary:
-       ctx.binary = ELF(binary, checksec=False)
-except ELFError:
+       ctx.binary = Binary.from_path(binary, checksec=False)
+except (ELFError, PEFormatError):
     pass
 
 if not binary:
@@ -45,7 +47,11 @@ from pwn import *
 # Set up pwntools for the correct architecture
 %endif
 %if ctx.binary or not host:
-exe = context.binary = ELF(args.EXE or ${binary_repr})
+%if not ctx.binary or isinstance(ctx.binary, ELF):
+    exe = context.binary = ELF(args.EXE or ${binary_repr})
+%else:
+    exe = context.binary = PE(args.EXE or ${binary_repr})
+%endif
 <% binary_repr = 'exe.path' %>
 %else:
 context.update(arch='i386')
@@ -113,7 +119,11 @@ else:
 def start_local(argv=[], *a, **kw):
     '''Execute the target binary locally'''
     if args.GDB:
+%if not ctx.binary or isinstance(ctx.binary, ELF)
         return gdb.debug([${binary_repr}] + argv, gdbscript=gdbscript, *a, **kw)
+%else:
+        return windbg.debug([${binary_repr}] + argv, dbgscript=gdbscript, *a, **kw)
+%endif
     else:
         return process([${binary_repr}] + argv, *a, **kw)
 

pwnlib/pe/pe.py

@@ -21,7 +21,7 @@
 __all__ = ['PE']
 
 class PE(PEFile, Binary):
-    def __init__(self, path, checksec=True, load_pdb=True):
+    def __init__(self, path, checksec=True, load_pdb=False):
         super(PE,self).__init__(path)
 
         # File handle to mimic ELF class

pwnlib/windbg.py

@@ -77,14 +77,14 @@
 CREATE_SUSPENDED = 0x00000004
 
 @LocalContext
-def debug(args, windbgscript=None, exe=None, env=None, creationflags=0, **kwargs):
-    """debug(args, windbgscript=None, exe=None, env=None, creationflags=0) -> tube
+def debug(args, dbgscript=None, exe=None, env=None, creationflags=0, **kwargs):
+    """debug(args, dbgscript=None, exe=None, env=None, creationflags=0) -> tube
 
     Launch a process in suspended state, attach debugger and resume process.
 
     Arguments:
         args(list): Arguments to the process, similar to :class:`.process`.
-        windbgscript(str): windbg script to run.
+        dbgscript(str): windbg script to run.
         exe(str): Path to the executable on disk.
         env(dict): Environment to start the binary in.
         creationflags(int): Flags to pass to :func:`.process.process`.
@@ -102,7 +102,7 @@ def debug(args, windbgscript=None, exe=None, env=None, creationflags=0, **kwargs
             go
             ''')
 
-        When WinDbg opens via :func:`.debug`, it will initially be stopped on the very first
+        When the debugger opens via :func:`.debug`, it will initially be stopped on the very first
         instruction of the entry point.
     """
     if isinstance(
@@ -114,43 +114,84 @@ def debug(args, windbgscript=None, exe=None, env=None, creationflags=0, **kwargs
         log.warn_once("Skipping debugger since context.noptrace==True")
         return tubes.process.process(args, executable=exe, env=env, creationflags=creationflags)
 
-    windbgscript = windbgscript or ''
-    if isinstance(windbgscript, six.string_types):
-        windbgscript = windbgscript.split('\n')
+    dbgscript = dbgscript or ''
+    if isinstance(dbgscript, six.string_types):
+        dbgscript = dbgscript.split('\n')
     # resume main thread
-    windbgscript = ['~0m'] + windbgscript
+    if context.debugger_selection in ('windbg', ''):
+        dbgscript = ['~0m'] + dbgscript
     creationflags |= CREATE_SUSPENDED
     io = tubes.process.process(args, executable=exe, env=env, creationflags=creationflags)
-    attach(target=io, windbgscript=windbgscript, **kwargs)
+    attach(target=io, dbgscript=dbgscript, **kwargs)
 
     return io
 
 def binary():
     """binary() -> str
 
-    Returns the path to the WinDbg binary.
+    Returns the path to the debugger binary depending on the context.
+    :attr:`.context.debugger_selection` is used to determine which debugger to use.
 
     Returns:
-        str: Path to the appropriate ``windbg`` binary to use.
+        str: Path to the appropriate debugger binary to use.
     """
-    windbg = misc.which('windbgx.exe') or misc.which('windbg.exe')
-    if not windbg:
-        log.error('windbg is not installed or in system PATH')
-    return windbg
+    if context.debugger_selection == 'x64dbg':
+        return _lookup_x64dbg()
+
+    if not context.debugger_selection or context.debugger_selection == 'windbg':
+        if context.windbg_binary:
+            windbg = misc.which(context.windbg_binary)
+            if not windbg:
+                log.warn_once('Path to WinDBG binary `{}` not found'.format(context.windbg_binary))
+            return windbg
+
+        windbg = misc.which('windbgx.exe') or misc.which('windbg.exe')
+        if not windbg:
+            log.error('windbg is not installed or in system PATH')
+        return windbg
+    log.error('Invalid debugger selection: %s', context.debugger_selection)
+
+def _lookup_x64dbg():
+    if context.x64dbg_binary:
+        x64dbg = misc.which(context.x64dbg_binary)
+        if not x64dbg:
+            log.warn_once('Path to x64dbg binary `{}` not found'.format(context.x64dbg_binary))
+        return x64dbg
+
+    x64dbg = misc.which('x96dbg.exe')
+    if x64dbg:
+        return x64dbg
+
+    # See if the "Debug with x64dbg" shell extension is installed
+    try:
+        import winreg
+        with winreg.OpenKey(winreg.HKEY_CLASSES_ROOT, r'exefile\shell\Debug with x64dbg\Command') as key:
+            regcmd = winreg.QueryValueEx(key, None)
+            # ('"C:\\Users\\User\\Downloads\\x64dbg\\bin\\x96dbg.exe" "%1"', 2)
+            if regcmd[1] != winreg.REG_EXPAND_SZ:
+                log.error('x64dbg registry key is not REG_EXPAND_SZ')
+            command = regcmd[0].split('"')[1]
+            if not os.path.exists(command):
+                log.error('x64dbg path from registry does not exist')
+            return command
+    except FileNotFoundError:
+        pass
+
+    log.error('x64dbg is not installed or in system PATH')
 
 @LocalContext
-def attach(target, windbgscript=None, windbg_args=[]):
-    """attach(target, windbgscript=None, windbg_args=[]) -> int
+def attach(target, dbgscript=None, dbg_args=[]):
+    """attach(target, dbgscript=None, dbg_args=[]) -> int
 
-    Attach to a running process with WinDbg.
+    Attach to a running process with WinDbg or x64dbg.
 
     Arguments:
         target(int, str, process): Process to attach to.
-        windbgscript(str, list): WinDbg script to run after attaching.
-        windbg_args(list): Additional arguments to pass to WinDbg.
+        dbgscript(str, list): Debugger script to run after attaching.
+        dbg_args(list): Additional arguments to pass to the debugger.
 
     Returns:
-        int: PID of the WinDbg process.
+        int: PID of the debugger process.
 
     Notes:
 
@@ -176,7 +217,7 @@ def attach(target, windbgscript=None, windbg_args=[]):
         Attach a debugger to a :class:`.process` tube and automate interaction
 
         >>> io = process('cmd') # doctest: +SKIP
-        >>> pid = windbg.attach(io, windbgscript='''
+        >>> pid = windbg.attach(io, dbgscript='''
         ... bp kernelbase!WriteFile
         ... g
         ... ''') # doctest: +SKIP
@@ -207,36 +248,36 @@ def attach(target, windbgscript=None, windbg_args=[]):
         log.error('could not find target process')
 
     cmd = [binary()]
-    if windbg_args:
-        cmd.extend(windbg_args)
+    if dbg_args:
+        cmd.extend(dbg_args)
 
     cmd.extend(['-p', str(pid)])
 
-    windbgscript = windbgscript or ''
-    if isinstance(windbgscript, six.string_types):
-        windbgscript = windbgscript.split('\n')
-    if isinstance(windbgscript, list):
-        windbgscript = ';'.join(script.strip() for script in windbgscript if script.strip())
-    if windbgscript:
-        cmd.extend(['-c', windbgscript])
+    dbgscript = dbgscript or ''
+    if isinstance(dbgscript, six.string_types):
+        dbgscript = dbgscript.split('\n')
+    if isinstance(dbgscript, list):
+        dbgscript = ';'.join(script.strip() for script in dbgscript if script.strip())
+    if dbgscript:
+        cmd.extend(['-c', dbgscript])
 
     log.info("Launching a new process: %r" % cmd)
 
     io = subprocess.Popen(cmd)
-    windbg_pid = io.pid
+    debugger_pid = io.pid
 
     def kill():
         try:
-            os.kill(windbg_pid, signal.SIGTERM)
+            os.kill(debugger_pid, signal.SIGTERM)
         except OSError:
             pass
 
     atexit.register(kill)
 
     if context.native:
-        proc.wait_for_debugger(pid, windbg_pid)
+        proc.wait_for_debugger(pid, debugger_pid)
 
-    return windbg_pid
+    return debugger_pid
 
 def minidump_on_crash(process, dump_type='mini'):
     if context.noptrace: