Misc run_in_new_terminal improvements.

pwnlib/util/misc.py

@@ -3,12 +3,14 @@
 import base64
 import errno
 import os
-import platform
 import re
+import signal
 import socket
 import stat
 import string
+import subprocess
 
+from pwnlib import atexit
 from pwnlib.context import context
 from pwnlib.log import getLogger
 from pwnlib.util import fiddling
@@ -178,27 +180,31 @@ def which(name, all = False):
     else:
         return None
 
-def run_in_new_terminal(command, terminal = None, args = None):
-    """run_in_new_terminal(command, terminal = None) -> None
+def run_in_new_terminal(command, terminal = None, args = None, kill_at_exit = True):
+    """run_in_new_terminal(command, terminal = None, args = None, kill_at_exit = True) -> int
 
     Run a command in a new terminal.
 
     When ``terminal`` is not set:
         - If ``context.terminal`` is set it will be used.
           If it is an iterable then ``context.terminal[1:]`` are default arguments.
         - If a ``pwntools-terminal`` command exists in ``$PATH``, it is used
-        - If ``$TERM_PROGRAM`` is set, that is used.
-        - If X11 is detected (by the presence of the ``$DISPLAY`` environment
-          variable), ``x-terminal-emulator`` is used.
         - If tmux is detected (by the presence of the ``$TMUX`` environment
           variable), a new pane will be opened.
         - If GNU Screen is detected (by the presence of the ``$STY`` environment
           variable), a new screen will be opened.
+        - If ``$TERM_PROGRAM`` is set, that is used.
+        - If X11 is detected (by the presence of the ``$DISPLAY`` environment
+          variable), ``x-terminal-emulator`` is used.
+
+    If `kill_at_exit` is :const:`True`, try to close the command/terminal when the
+    current process exits. This may not work for all terminal types.
 
     Arguments:
         command (str): The command to run.
         terminal (str): Which terminal to use.
         args (list): Arguments to pass to the terminal
+        kill_at_exit (bool): Whether to close the command/terminal on process exit.
 
     Note:
         The command is opened with ``/dev/null`` for stdin, stdout, stderr.
@@ -214,18 +220,18 @@ def run_in_new_terminal(command, terminal = None, args = None):
         elif which('pwntools-terminal'):
             terminal = 'pwntools-terminal'
             args     = []
+        elif 'TMUX' in os.environ and which('tmux'):
+            terminal = 'tmux'
+            args     = ['splitw', '-F' '#{pane_pid}', '-P']
+        elif 'STY' in os.environ and which('screen'):
+            terminal = 'screen'
+            args     = ['-t','pwntools-gdb','bash','-c']
         elif 'TERM_PROGRAM' in os.environ:
             terminal = os.environ['TERM_PROGRAM']
             args     = []
         elif 'DISPLAY' in os.environ and which('x-terminal-emulator'):
             terminal = 'x-terminal-emulator'
             args     = ['-e']
-        elif 'TMUX' in os.environ and which('tmux'):
-            terminal = 'tmux'
-            args     = ['splitw']
-        elif 'STY' in os.environ and which('screen'):
-            terminal = 'screen'
-            args     = ['-t','pwntools-gdb','bash','-c']
 
     if not terminal:
         log.error('Could not find a terminal binary to use. Set context.terminal to your terminal.')
@@ -248,17 +254,23 @@ def run_in_new_terminal(command, terminal = None, args = None):
 
     log.debug("Launching a new terminal: %r" % argv)
 
-    pid = os.fork()
-
-    if pid == 0:
-        # Closing the file descriptors makes everything fail under tmux on OSX.
-        if platform.system() != 'Darwin':
-            devnull = open(os.devnull, 'rwb')
-            os.dup2(devnull.fileno(), 0)
-            os.dup2(devnull.fileno(), 1)
-            os.dup2(devnull.fileno(), 2)
-        os.execv(argv[0], argv)
-        os._exit(1)
+    stdin = stdout = stderr = open(os.devnull, 'rwb')
+    if terminal == 'tmux':
+        stdout = subprocess.PIPE
+
+    p = subprocess.Popen(argv, stdin=stdin, stdout=stdout, stderr=stderr)
+
+    if terminal == 'tmux':
+        out, _ = p.communicate()
+        pid = int(out)
+    else:
+        pid = p.pid
+
+    if kill_at_exit:
+        if terminal == 'tmux':
+            atexit.register(lambda: os.kill(pid, signal.SIGTERM))
+        else:
+            atexit.register(lambda: p.terminate())
 
     return pid