Completed
Push — master ( a6eb7a...c47a1a )
by Ionel Cristian
27s
created

RemoteStream.flush()   A

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