RemoteStream.flush()   A
last analyzed

Complexity

Conditions 1

Size

Total Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 0
Metric Value
cc 1
c 2
b 0
f 0
dl 0
loc 2
rs 10
1
from __future__ import print_function
2
3
import argparse
4
import errno
5
import json
6
import os
7
import signal
8
import socket
9
import sys
10
import time
11
from contextlib import closing
12
from contextlib import contextmanager
13
from subprocess import check_call
14
15
import manhole
16
from manhole import get_peercred
17
from manhole.cli import parse_signal
18
19
from . import actions
20
from . import stop
21
from . import trace
22
23
24
def install(**kwargs):
25
    kwargs.setdefault('oneshot_on', 'URG')
26
    kwargs.setdefault('connection_handler', 'exec')
27
    manhole.install(**kwargs)
28
29
30
class RemoteStream(object):
31
    def __init__(self, path, isatty, encoding):
32
        self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
33
        self._sock.connect(path)
34
        self._isatty = isatty
35
        self._encoding = encoding
36
37
    def isatty(self):
38
        return self._isatty
39
40
    def write(self, data):
41
        try:
42
            self._sock.send(data.encode(self._encoding))
43
        except Exception as exc:
44
            print("Hunter failed to send trace output: %s. Stopping tracer." % exc, file=sys.stderr)
45
            stop()
46
47
    def flush(self):
48
        pass
49
50
51
@contextmanager
52
def manhole_bootstrap(args, activation_payload, deactivation_payload):
53
    activation_payload += '\nexit()\n'
54
    deactivation_payload += '\nexit()\n'
55
56
    activation_payload = activation_payload.encode('utf-8')
57
    deactivation_payload = deactivation_payload.encode('utf-8')
58
59
    with closing(connect_manhole(args.pid, args.timeout, args.signal)) as manhole:
60
        manhole.send(activation_payload)
61
    try:
62
        yield
63
    finally:
64
        with closing(connect_manhole(args.pid, args.timeout, args.signal)) as manhole:
65
            manhole.send(deactivation_payload)
66
67
68
@contextmanager
69
def gdb_bootstrap(args, activation_payload, deactivation_payload):
70
    print('WARNING: Using GDB may deadlock the process or create unpredictable results!')
71
    activation_command = [
72
        'gdb', '-p', str(args.pid), '-batch',
73
        '-ex', 'call Py_AddPendingCall(PyRun_SimpleString, %s)' % json.dumps(activation_payload),
74
    ]
75
    deactivation_command = [
76
        'gdb', '-p', str(args.pid), '-batch',
77
        '-ex', 'call Py_AddPendingCall(PyRun_SimpleString, %s)' % json.dumps(deactivation_payload),
78
    ]
79
    check_call(activation_command)
80
    try:
81
        yield
82
    finally:
83
        check_call(deactivation_command)
84
85
86
def connect_manhole(pid, timeout, signal):
87
    os.kill(pid, signal)
88
89
    start = time.time()
90
    uds_path = '/tmp/manhole-%s' % pid
91
    manhole = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
92
    manhole.settimeout(timeout)
93
    while time.time() - start < timeout:
94
        try:
95
            manhole.connect(uds_path)
96
        except Exception as exc:
97
            if exc.errno not in (errno.ENOENT, errno.ECONNREFUSED):
98
                print("Failed to connect to %r: %r" % (uds_path, exc), file=sys.stderr)
99
        else:
100
            break
101
    else:
102
        print("Failed to connect to %r: Timeout" % uds_path, file=sys.stderr)
103
        sys.exit(5)
104
    return manhole
105
106
107
def activate(sink_path, isatty, encoding, options):
108
    stream = actions.DEFAULT_STREAM = RemoteStream(sink_path, isatty, encoding)
109
    try:
110
        stream.write("Output stream active. Starting tracer ...\n\n")
111
        eval("trace(%s)" % options)
112
    except Exception as exc:
113
        stream.write("Failed to activate: %s. %s\n" % (
114
            exc,
115
            'Tracer options where: %s.' % options if options else 'No tracer options.'
116
        ))
117
        actions.DEFAULT_STREAM = sys.stderr
118
        raise
119
120
121
trace  # used in eval above
122
123
124
def deactivate():
125
    actions.DEFAULT_STREAM = sys.stderr
126
    stop()
127
128
129
parser = argparse.ArgumentParser(description='Trace a process.')
130
parser.add_argument('-p', '--pid', metavar='PID', type=int, required=True,
131
                    help='A numerical process id.')
132
parser.add_argument('options', metavar='OPTIONS', nargs='*')
133
parser.add_argument('-t', '--timeout', dest='timeout', default=1, type=float,
134
                    help='Timeout to use. Default: %(default)s seconds.')
135
parser.add_argument('--gdb', dest='gdb', action='store_true',
136
                    help='Use GDB to activate tracing. WARNING: it may deadlock the process!')
137
parser.add_argument('-s', '--signal', dest='signal', type=parse_signal, metavar="SIGNAL", default=signal.SIGURG,
138
                    help='Send the given SIGNAL to the process before connecting.')
139
140
141
def main():
142
    args = parser.parse_args()
143
144
    sink = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
145
    sink_path = '/tmp/hunter-%s' % os.getpid()
146
    sink.bind(sink_path)
147
    sink.listen(1)
148
    os.chmod(sink_path, 0o777)
149
150
    stdout = os.fdopen(sys.stdout.fileno(), 'wb', 0)
151
    encoding = getattr(sys.stdout, 'encoding', 'utf-8') or 'utf-8'
152
    bootstrapper = gdb_bootstrap if args.gdb else manhole_bootstrap
153
    payload = 'from hunter import remote; remote.activate(%r, %r, %r, %r)' % (
154
        sink_path,
155
        sys.stdout.isatty(),
156
        encoding,
157
        ','.join(i.strip(',') for i in args.options)
158
    )
159
    with bootstrapper(args, payload, 'from hunter import remote; remote.deactivate()'):
160
        conn, _ = sink.accept()
161
        os.unlink(sink_path)
162
        pid, _, _ = get_peercred(conn)
163
        if pid != args.pid:
164
            raise Exception("Unexpected pid %r connected to output socket. Was expecting %s." % (pid, args.pid))
165
        data = conn.recv(1024)
166
        while data:
167
            stdout.write(data)
168
            data = conn.recv(1024)
169