viztracer 1.1.0__cp314-cp314-win32.whl

viztracer/attach_process/add_code_to_python_process.py

@@ -0,0 +1,582 @@
+# type: ignore
+
+r'''
+Copyright: Brainwy Software Ltda.
+
+License: EPL.
+=============
+
+Works for Windows by using an executable that'll inject a dll to a process and call a function.
+
+Note: https://github.com/fabioz/winappdbg is used just to determine if the target process is 32 or 64 bits.
+
+Works for Linux relying on gdb.
+
+Limitations:
+============
+
+    Linux:
+    ------
+
+        1. It possible that ptrace is disabled: /etc/sysctl.d/10-ptrace.conf
+
+        Note that even enabling it in /etc/sysctl.d/10-ptrace.conf (i.e.: making the
+        ptrace_scope=0), it's possible that we need to run the application that'll use ptrace (or
+        gdb in this case) as root (so, we must sudo the python which'll run this module).
+
+        2. It currently doesn't work in debug builds (i.e.: python_d)
+
+
+Other implementations:
+- pyrasite.com:
+    GPL
+    Windows/linux (in Linux it also uses gdb to connect -- although specifics are different as we use a dll to execute
+    code with other threads stopped). It's Windows approach is more limited because it doesn't seem to deal properly with
+    Python 3 if threading is disabled.
+
+- https://github.com/google/pyringe:
+    Apache v2.
+    Only linux/Python 2.
+
+- http://pytools.codeplex.com:
+    Apache V2
+    Windows Only (but supports mixed mode debugging)
+    Our own code relies heavily on a part of it: http://pytools.codeplex.com/SourceControl/latest#Python/Product/PyDebugAttach/PyDebugAttach.cpp
+    to overcome some limitations of attaching and running code in the target python executable on Python 3.
+    See: attach.cpp
+
+Linux: References if we wanted to use a pure-python debugger:
+    https://bitbucket.org/haypo/python-ptrace/
+    http://stackoverflow.com/questions/7841573/how-to-get-an-error-message-for-errno-value-in-python
+    Jugaad:
+        https://www.defcon.org/images/defcon-19/dc-19-presentations/Jakhar/DEFCON-19-Jakhar-Jugaad-Linux-Thread-Injection.pdf
+        https://github.com/aseemjakhar/jugaad
+
+Something else (general and not Python related):
+- http://www.codeproject.com/Articles/4610/Three-Ways-to-Inject-Your-Code-into-Another-Proces
+
+Other references:
+- https://github.com/haypo/faulthandler
+- http://nedbatchelder.com/text/trace-function.html
+- https://github.com/python-git/python/blob/master/Python/sysmodule.c (sys_settrace)
+- https://github.com/python-git/python/blob/master/Python/ceval.c (PyEval_SetTrace)
+- https://github.com/python-git/python/blob/master/Python/thread.c (PyThread_get_key_value)
+
+
+To build the dlls needed on windows, visual studio express 13 was used (see compile_dll.bat)
+
+See: attach_pydevd.py to attach the pydev debugger to a running python process.
+'''
+
+# Note: to work with nasm compiling asm to code and decompiling to see asm with shellcode:
+# x:\nasm\nasm-2.07-win32\nasm-2.07\nasm.exe
+# nasm.asm&x:\nasm\nasm-2.07-win32\nasm-2.07\ndisasm.exe -b arch nasm
+import ctypes
+import os
+import struct
+import subprocess
+import sys
+import time
+from contextlib import contextmanager
+import platform
+import traceback
+
+try:
+    TimeoutError = TimeoutError  # @ReservedAssignment
+except NameError:
+
+    class TimeoutError(RuntimeError):  # @ReservedAssignment
+        pass
+
+
+@contextmanager
+def _create_win_event(name):
+    from winappdbg.win32.kernel32 import CreateEventA, WaitForSingleObject, CloseHandle
+
+    manual_reset = False  # i.e.: after someone waits it, automatically set to False.
+    initial_state = False
+    if not isinstance(name, bytes):
+        name = name.encode('utf-8')
+    event = CreateEventA(None, manual_reset, initial_state, name)
+    if not event:
+        raise ctypes.WinError()
+
+    class _WinEvent(object):
+
+        def wait_for_event_set(self, timeout=None):
+            '''
+            :param timeout: in seconds
+            '''
+            if timeout is None:
+                timeout = 0xFFFFFFFF
+            else:
+                timeout = int(timeout * 1000)
+            ret = WaitForSingleObject(event, timeout)
+            if ret in (0, 0x80):
+                return True
+            elif ret == 0x102:
+                # Timed out
+                return False
+            else:
+                raise ctypes.WinError()
+
+    try:
+        yield _WinEvent()
+    finally:
+        CloseHandle(event)
+
+
+IS_WINDOWS = sys.platform == 'win32'
+IS_LINUX = sys.platform in ('linux', 'linux2')
+IS_MAC = sys.platform == 'darwin'
+
+
+def is_python_64bit():
+    return (struct.calcsize('P') == 8)
+
+
+def get_target_filename(is_target_process_64=None, prefix=None, extension=None):
+    # Note: we have an independent (and similar -- but not equal) version of this method in
+    # `pydevd_tracing.py` which should be kept synchronized with this one (we do a copy
+    # because the `pydevd_attach_to_process` is mostly independent and shouldn't be imported in the
+    # debugger -- the only situation where it's imported is if the user actually does an attach to
+    # process, through `attach_pydevd.py`, but this should usually be called from the IDE directly
+    # and not from the debugger).
+    libdir = os.path.dirname(__file__)
+
+    if is_target_process_64 is None:
+        if IS_WINDOWS:
+            # i.e.: On windows the target process could have a different bitness (32bit is emulated on 64bit).
+            raise AssertionError("On windows it's expected that the target bitness is specified.")
+
+        # For other platforms, just use the the same bitness of the process we're running in.
+        is_target_process_64 = is_python_64bit()
+
+    arch = ''
+    if IS_WINDOWS:
+        # prefer not using platform.machine() when possible (it's a bit heavyweight as it may
+        # spawn a subprocess).
+        arch = os.environ.get("PROCESSOR_ARCHITEW6432", os.environ.get('PROCESSOR_ARCHITECTURE', ''))
+
+    if not arch:
+        arch = platform.machine()
+        if not arch:
+            print('platform.machine() did not return valid value.')  # This shouldn't happen...
+            return None
+
+    if IS_WINDOWS:
+        if not extension:
+            extension = '.dll'
+        suffix_64 = 'amd64'
+        suffix_32 = 'x86'
+
+    elif IS_LINUX:
+        if not extension:
+            extension = '.so'
+        suffix_64 = 'amd64'
+        suffix_32 = 'x86'
+
+    elif IS_MAC:
+        if not extension:
+            extension = '.dylib'
+        suffix_64 = 'x86_64'
+        suffix_32 = 'x86'
+
+    else:
+        print('Unable to attach to process in platform: %s', sys.platform)
+        return None
+
+    if arch.lower() not in ('amd64', 'x86', 'x86_64', 'i386', 'x86'):
+        # We don't support this processor by default. Still, let's support the case where the
+        # user manually compiled it himself with some heuristics.
+        #
+        # Ideally the user would provide a library in the format: "attach_<arch>.<extension>"
+        # based on the way it's currently compiled -- see:
+        # - windows/compile_windows.bat
+        # - linux_and_mac/compile_linux.sh
+        # - linux_and_mac/compile_mac.sh
+
+        try:
+            found = [name for name in os.listdir(libdir) if name.startswith('attach_') and name.endswith(extension)]
+        except:
+            print('Error listing dir: %s' % (libdir,))
+            traceback.print_exc()
+            return None
+
+        if prefix:
+            expected_name = prefix + arch + extension
+            expected_name_linux = prefix + 'linux_' + arch + extension
+        else:
+            # Default is looking for the attach_ / attach_linux
+            expected_name = 'attach_' + arch + extension
+            expected_name_linux = 'attach_linux_' + arch + extension
+
+        filename = None
+        if expected_name in found:  # Heuristic: user compiled with "attach_<arch>.<extension>"
+            filename = os.path.join(libdir, expected_name)
+
+        elif IS_LINUX and expected_name_linux in found:  # Heuristic: user compiled with "attach_linux_<arch>.<extension>"
+            filename = os.path.join(libdir, expected_name_linux)
+
+        elif len(found) == 1:  # Heuristic: user removed all libraries and just left his own lib.
+            filename = os.path.join(libdir, found[0])
+
+        else:  # Heuristic: there's one additional library which doesn't seem to be our own. Find the odd one.
+            filtered = [name for name in found if not name.endswith((suffix_64 + extension, suffix_32 + extension))]
+            if len(filtered) == 1:  # If more than one is available we can't be sure...
+                filename = os.path.join(libdir, found[0])
+
+        if filename is None:
+            print(
+                'Unable to attach to process in arch: %s (did not find %s in %s).' % (
+                    arch, expected_name, libdir
+                )
+            )
+            return None
+
+        print('Using %s in arch: %s.' % (filename, arch))
+
+    else:
+        if is_target_process_64:
+            suffix = suffix_64
+        else:
+            suffix = suffix_32
+
+        if not prefix:
+            # Default is looking for the attach_ / attach_linux
+            if IS_WINDOWS or IS_MAC:  # just the extension changes
+                prefix = 'attach_'
+            elif IS_LINUX:
+                prefix = 'attach_linux_'  # historically it has a different name
+            else:
+                print('Unable to attach to process in platform: %s' % (sys.platform,))
+                return None
+
+        filename = os.path.join(libdir, '%s%s%s' % (prefix, suffix, extension))
+
+    if not os.path.exists(filename):
+        print('Expected: %s to exist.' % (filename,))
+        return None
+
+    return filename
+
+
+def run_python_code_windows(pid, python_code, connect_debugger_tracing=False, show_debug_info=0):
+    assert '\'' not in python_code, 'Having a single quote messes with our command.'
+    from winappdbg.process import Process
+    if not isinstance(python_code, bytes):
+        python_code = python_code.encode('utf-8')
+
+    process = Process(pid)
+    bits = process.get_bits()
+    is_target_process_64 = bits == 64
+
+    # Note: this restriction no longer applies (we create a process with the proper bitness from
+    # this process so that the attach works).
+    # if is_target_process_64 != is_python_64bit():
+    #     raise RuntimeError("The architecture of the Python used to connect doesn't match the architecture of the target.\n"
+    #     "Target 64 bits: %s\n"
+    #     "Current Python 64 bits: %s" % (is_target_process_64, is_python_64bit()))
+
+    with _acquire_mutex('_pydevd_pid_attach_mutex_%s' % (pid,), 10):
+        # print('--- Connecting to %s bits target (current process is: %s) ---' % (bits, 64 if is_python_64bit() else 32))
+
+        with _win_write_to_shared_named_memory(python_code, pid):
+
+            target_executable = get_target_filename(is_target_process_64, 'inject_dll_', '.exe')
+            if not target_executable:
+                raise RuntimeError('Could not find expected .exe file to inject dll in attach to process.')
+
+            target_dll = get_target_filename(is_target_process_64)
+            if not target_dll:
+                raise RuntimeError('Could not find expected .dll file in attach to process.')
+
+            # print('\n--- Injecting attach dll: %s into pid: %s ---' % (os.path.basename(target_dll), pid))
+            args = [target_executable, str(pid), target_dll]
+            subprocess.check_call(args)
+
+            # Now, if the first injection worked, go on to the second which will actually
+            # run the code.
+            target_dll_run_on_dllmain = get_target_filename(is_target_process_64, 'run_code_on_dllmain_', '.dll')
+            if not target_dll_run_on_dllmain:
+                raise RuntimeError('Could not find expected .dll in attach to process.')
+
+            with _create_win_event('_pydevd_pid_event_%s' % (pid,)) as event:
+                # print('\n--- Injecting run code dll: %s into pid: %s ---' % (os.path.basename(target_dll_run_on_dllmain), pid))
+                args = [target_executable, str(pid), target_dll_run_on_dllmain]
+                subprocess.check_call(args)
+
+                if not event.wait_for_event_set(10):
+                    print('Timeout error: the attach may not have completed.')
+            # print('--- Finished dll injection ---\n')
+
+    return 0, b"", b""
+
+
+@contextmanager
+def _acquire_mutex(mutex_name, timeout):
+    '''
+    Only one process may be attaching to a pid, so, create a system mutex
+    to make sure this holds in practice.
+    '''
+    from winappdbg.win32.kernel32 import CreateMutex, GetLastError, CloseHandle
+    from winappdbg.win32.defines import ERROR_ALREADY_EXISTS
+
+    initial_time = time.time()
+    while True:
+        mutex = CreateMutex(None, True, mutex_name)
+        acquired = GetLastError() != ERROR_ALREADY_EXISTS
+        if acquired:
+            break
+        if time.time() - initial_time > timeout:
+            raise TimeoutError('Unable to acquire mutex to make attach before timeout.')
+        time.sleep(.2)
+
+    try:
+        yield
+    finally:
+        CloseHandle(mutex)
+
+
+@contextmanager
+def _win_write_to_shared_named_memory(python_code, pid):
+    # Use the definitions from winappdbg when possible.
+    from winappdbg.win32 import defines
+    from winappdbg.win32.kernel32 import (
+        CreateFileMapping,
+        MapViewOfFile,
+        CloseHandle,
+        UnmapViewOfFile,
+    )
+
+    memmove = ctypes.cdll.msvcrt.memmove
+    memmove.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_void_p,
+        defines.SIZE_T,
+    ]
+    memmove.restype = ctypes.c_void_p
+
+    # Note: BUFSIZE must be the same from run_code_in_memory.hpp
+    BUFSIZE = 2048
+    assert isinstance(python_code, bytes)
+    assert len(python_code) > 0, 'Python code must not be empty.'
+    # Note: -1 so that we're sure we'll add a \0 to the end.
+    assert len(python_code) < BUFSIZE - 1, 'Python code must have at most %s bytes (found: %s)' % (BUFSIZE - 1, len(python_code))
+
+    python_code += b'\0' * (BUFSIZE - len(python_code))
+    assert python_code.endswith(b'\0')
+
+    INVALID_HANDLE_VALUE = -1
+    PAGE_READWRITE = 0x4
+    FILE_MAP_WRITE = 0x2
+    filemap = CreateFileMapping(
+        INVALID_HANDLE_VALUE, 0, PAGE_READWRITE, 0, BUFSIZE, u"__pydevd_pid_code_to_run__%s" % (pid,))
+
+    if filemap == INVALID_HANDLE_VALUE or filemap is None:
+        raise Exception("Failed to create named file mapping (ctypes: CreateFileMapping): %s" % (filemap,))
+    try:
+        view = MapViewOfFile(filemap, FILE_MAP_WRITE, 0, 0, 0)
+        if not view:
+            raise Exception("Failed to create view of named file mapping (ctypes: MapViewOfFile).")
+
+        try:
+            memmove(view, python_code, BUFSIZE)
+            yield
+        finally:
+            UnmapViewOfFile(view)
+    finally:
+        CloseHandle(filemap)
+
+
+def run_python_code_linux(pid, python_code, connect_debugger_tracing=False, show_debug_info=0):
+    assert '\'' not in python_code, 'Having a single quote messes with our command.'
+
+    target_dll = get_target_filename()
+    if not target_dll:
+        raise RuntimeError('Could not find .so for attach to process.')
+    target_dll_name = os.path.splitext(os.path.basename(target_dll))[0]
+
+    # Note: we currently don't support debug builds
+    is_debug = 0
+    # Note that the space in the beginning of each line in the multi-line is important!
+    cmd = [
+        'gdb',
+        '--nw',  # no gui interface
+        '--nh',  # no ~/.gdbinit
+        '--nx',  # no .gdbinit
+#         '--quiet',  # no version number on startup
+        '--pid',
+        str(pid),
+        '--batch',
+#         '--batch-silent',
+    ]
+
+    cmd.extend(["--eval-command='set scheduler-locking off'"])  # If on we'll deadlock.
+
+    # Leave auto by default (it should do the right thing as we're attaching to a process in the
+    # current host).
+    cmd.extend(["--eval-command='set architecture auto'"])
+
+    cmd.extend([
+        "--eval-command='call (void*)dlopen(\"%s\", 2)'" % target_dll,
+        "--eval-command='sharedlibrary %s'" % target_dll_name,
+        "--eval-command='call (int)DoAttach(%s, \"%s\", %s)'" % (
+            is_debug, python_code, show_debug_info)
+    ])
+
+    # print ' '.join(cmd)
+
+    env = os.environ.copy()
+    # Remove the PYTHONPATH (if gdb has a builtin Python it could fail if we
+    # have the PYTHONPATH for a different python version or some forced encoding).
+    env.pop('PYTHONIOENCODING', None)
+    env.pop('PYTHONPATH', None)
+    # print('Running: %s' % (' '.join(cmd)))
+    p = subprocess.Popen(
+        ' '.join(cmd),
+        shell=True,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    # print('Running gdb in target process.')
+    out, err = p.communicate()
+    # print('stdout: %s' % (out,))
+    # print('stderr: %s' % (err,))
+    return p.returncode, out, err
+
+
+def find_helper_script(filedir, script_name):
+    target_filename = os.path.join(filedir, 'linux_and_mac', script_name)
+    target_filename = os.path.normpath(target_filename)
+    if not os.path.exists(target_filename):
+        raise RuntimeError('Could not find helper script: %s' % target_filename)
+
+    return target_filename
+
+
+def run_python_code_mac(pid, python_code, connect_debugger_tracing=False, show_debug_info=0):
+    assert '\'' not in python_code, 'Having a single quote messes with our command.'
+
+    target_dll = get_target_filename()
+    if not target_dll:
+        raise RuntimeError('Could not find .dylib for attach to process.')
+
+    libdir = os.path.dirname(__file__)
+    lldb_prepare_file = find_helper_script(libdir, 'lldb_prepare.py')
+    # Note: we currently don't support debug builds
+
+    is_debug = 0
+    # Note that the space in the beginning of each line in the multi-line is important!
+    cmd = [
+        'lldb',
+        '--no-lldbinit',  # Do not automatically parse any '.lldbinit' files.
+        # '--attach-pid',
+        # str(pid),
+        # '--arch',
+        # arch,
+        '--script-language',
+        'Python'
+        #         '--batch-silent',
+    ]
+
+    cmd.extend([
+        "-o 'process attach --pid %d'" % pid,
+        "-o 'command script import \"%s\"'" % (lldb_prepare_file,),
+        "-o 'load_lib_and_attach \"%s\" %s \"%s\" %s'" % (target_dll,
+            is_debug, python_code, show_debug_info),
+    ])
+
+    cmd.extend([
+        "-o 'process detach'",
+        "-o 'script import os; os._exit(0)'",
+    ])
+
+    # print ' '.join(cmd)
+
+    env = os.environ.copy()
+    # Remove the PYTHONPATH (if gdb has a builtin Python it could fail if we
+    # have the PYTHONPATH for a different python version or some forced encoding).
+    env.pop('PYTHONIOENCODING', None)
+    env.pop('PYTHONPATH', None)
+    # print('Running: %s' % (' '.join(cmd)))
+    p = subprocess.Popen(
+        ' '.join(cmd),
+        shell=True,
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        )
+    # print('Running lldb in target process.')
+    out, err = p.communicate()
+    # print('stdout: %s' % (out,))
+    # print('stderr: %s' % (err,))
+    return p.returncode, out, err
+
+
+if IS_WINDOWS:
+    run_python_code = run_python_code_windows
+elif IS_MAC:
+    run_python_code = run_python_code_mac
+elif IS_LINUX:
+    run_python_code = run_python_code_linux
+else:
+
+    def run_python_code(*args, **kwargs):
+        print('Unable to attach to process in platform: %s', sys.platform)
+
+
+sys.path.append(os.path.dirname(__file__))
+
+
+def test():
+    print('Running with: %s' % (sys.executable,))
+    code = '''
+import os, time, sys
+print(os.getpid())
+#from threading import Thread
+#Thread(target=str).start()
+if __name__ == '__main__':
+    while True:
+        time.sleep(.5)
+        sys.stdout.write('.\\n')
+        sys.stdout.flush()
+'''
+
+    p = subprocess.Popen([sys.executable, '-u', '-c', code])
+    try:
+        code = 'print("It worked!")\n'
+
+        # Real code will be something as:
+        # code = '''import sys;sys.path.append(r'X:\winappdbg-code\examples'); import imported;'''
+        run_python_code(p.pid, python_code=code)
+        print('\nRun a 2nd time...\n')
+        run_python_code(p.pid, python_code=code)
+
+        time.sleep(3)
+    finally:
+        p.kill()
+
+
+def main(args):
+    # Otherwise, assume the first parameter is the pid and anything else is code to be executed
+    # in the target process.
+    pid = int(args[0])
+    del args[0]
+    python_code = ';'.join(args)
+
+    # Note: on Linux the python code may not have a single quote char: '
+    run_python_code(pid, python_code)
+
+
+if __name__ == '__main__':
+    args = sys.argv[1:]
+    if not args:
+        print('Expected pid and Python code to execute in target process.')
+    else:
+        if '--test' == args[0]:
+            test()
+        else:
+            main(args)
+

viztracer/attach_process/linux_and_mac/lldb_prepare.py

@@ -0,0 +1,54 @@
+# This file is meant to be run inside lldb
+# It registers command to load library and invoke attach function
+# Also it marks process threads to to distinguish them from debugger
+# threads later while settings trace in threads
+
+def load_lib_and_attach(debugger, command, result, internal_dict):
+    import shlex
+    args = shlex.split(command)
+
+    dll = args[0]
+    is_debug = args[1]
+    python_code = args[2]
+    show_debug_info = args[3]
+
+    import lldb
+    options = lldb.SBExpressionOptions()
+    options.SetFetchDynamicValue()
+    options.SetTryAllThreads(run_others=False)
+    options.SetTimeoutInMicroSeconds(timeout=60000000)
+
+    print(dll)
+    target = debugger.GetSelectedTarget()
+    res = target.EvaluateExpression("(void*)dlopen(\"%s\", 2);" % (
+        dll), options)
+    error = res.GetError()
+    if error:
+        print(error)
+
+    print(python_code)
+    res = target.EvaluateExpression("(int)DoAttach(%s, \"%s\", %s);" % (
+        is_debug, python_code.replace('"', "'"), show_debug_info), options)
+    error = res.GetError()
+    if error:
+        print(error)
+
+def __lldb_init_module(debugger, internal_dict):
+    import lldb
+
+    debugger.HandleCommand('command script add -f lldb_prepare.load_lib_and_attach load_lib_and_attach')
+
+    try:
+        target = debugger.GetSelectedTarget()
+        if target:
+            process = target.GetProcess()
+            if process:
+                for thread in process:
+                    # print('Marking process thread %d'%thread.GetThreadID())
+                    internal_dict['_thread_%d' % thread.GetThreadID()] = True
+                    # thread.Suspend()
+    except:
+        import traceback;traceback.print_exc()
+
+
+