import os
import sys
import functools
from twisted.python import usage, log
from twisted.internet import defer, reactor
from zope.interface import implementer
import txtorcon
import humanize
from carml import util
def attach_streams_per_process(state):
print("Exiting (e.g. Ctrl-C) will cause Tor to resume choosing circuits.")
print("Giving each new PID we see its own Circuit (until they're gone).")
@implementer(txtorcon.IStreamAttacher)
class Attacher(object):
def __init__(self):
self.pid_to_circuits = {}
def choose_new_circuit(self, stream, circuits):
for circ in circuits.values():
if circ in self.pid_to_circuits.values():
continue
if circ.state != 'BUILT':
continue
return circ
raise RuntimeError("Ran out of circuits to select.")
def attach_stream(self, stream, circuits):
src_addr, src_port = stream.flags['SOURCE_ADDR'].split(':')
pid = txtorcon.util.process_from_address(src_addr, src_port)
procname = os.path.realpath('/proc/%d/exe' % pid)
try:
circ = self.pid_to_circuits[pid]
except KeyError:
circ = self.choose_new_circuit(stream, circuits)
self.pid_to_circuits[pid] = circ
print('Selected circuit %d for process %d (%s).' % (circ.id, pid, procname))
print(' ', '->'.join([p.name if p.name_is_unique else ('{%s}' % p.name) for p in circ.path]))
# if stream.state == 'NEWRESOLVE':
# print " attaching %d (resolve %s)" % (stream.id, stream.target_host)
# else:
print(" attaching stream %d to circuit %d for %s:%d (%s)" % (stream.id, circ.id, stream.target_host, stream.target_port, procname))
return circ
state.set_attacher(Attacher(), reactor)
def attach_streams_to_circuit(circid, state):
try:
circ = state.circuits[circid]
except KeyError:
print("Circuit {} doesn't exist.".format(circid))
return None
print("Exiting (e.g. Ctrl-C) will cause Tor to resume choosing circuits.")
print("Attaching all new streams to Circuit %d." % circ.id)
print(" ", '->'.join([p.name if p.name_is_unique else ('~%s' % p.name) for p in circ.path]))
@implementer(txtorcon.IStreamAttacher)
class Attacher(txtorcon.CircuitListenerMixin):
def circuit_closed(self, this_circ, **kw):
if circ == this_circ:
print("Circuit {} vanished (REASON={}, REMOTE_REASON={})".format(
circ.id,
kw.get('REASON', 'not specified'),
kw.get('REMOTE_REASON', 'not specified'),
))
# should we just exit now?
# Pro: kind-of makes sense
# Con: if you're expecting streams to go via "your"
# circuit, maybe you want them to "fail closed" and
# not work at all -- which is what I'm doing right now
# so -> probably want exiting to be an option, and not the default
def attach_stream(self, stream, circuits):
if stream.flags.get('PURPOSE', 'unknown') in ['DIR_FETCH', 'DIR_UPLOAD', 'DIRPORT_TEST']:
print(" tor-internal directory stream ({})".format(stream.flags['PURPOSE']))
return None
if circ.state == 'CLOSED':
print(" target circuit is closed, not attaching")
return txtorcon.TorState.DO_NOT_ATTACH
if stream.state == 'NEWRESOLVE':
print(" attaching %d (resolve %s)" % (stream.id, stream.target_host))
else:
print(" attaching %d %s:%d" % (stream.id, stream.target_host,
stream.target_port))
return circ
attacher = Attacher()
state.set_attacher(attacher, reactor)
state.add_circuit_listener(attacher)
# FIXME doesn't exit on control-c? :(
d = defer.Deferred()
d.addBoth(lambda x: print('foo', x))
return d
async def list_streams(state, verbose):
print("Streams:")
for stream in state.streams.values():
flags = str(stream.flags) if stream.flags else 'no flags'
state = stream.state
state_to_color = dict(SUCCEEDED=util.colors.green,
FAILED=util.colors.red)
if state in state_to_color:
state = state_to_color[state](state)
print(" %d: %s on circuit %d (%s)" % (stream.id, state, stream.circuit.id,
flags))
if verbose:
h = stream.target_addr if stream.target_addr else stream.target_host
source = txtorcon.util.process_from_address(stream.source_addr, stream.source_port)
if source is None:
source = 'unknown'
print(" to %s:%s, from %s" % (h, stream.target_port, source))
async def close_stream(state, streamid):
class DetermineStreamClosure(object):
def __init__(self, target_id, done_d):
self.circ_id = str(target_id)
self.stream_gone = False
self.already_deleted = False
self.completed_d = done_d
def __call__(self, text):
cid, what, _ = text.split(' ', 2)
if what in ['CLOSED', 'FAILED']:
if self.circ_id == cid:
self.stream_gone = True
print("gone (%s)..." % self.circ_id,)
sys.stdout.flush()
if self.already_deleted:
self.completed_d.callback(self)
if streamid not in state.streams:
print('No such stream "%s".' % streamid)
return
print('Closing stream "%s"...' % (streamid, ))
gone_d = defer.Deferred()
monitor = DetermineStreamClosure(streamid, gone_d)
state.protocol.add_event_listener('STREAM', monitor)
sys.stdout.flush()
try:
status = await state.streams[streamid].close()
status = status.state
monitor.already_deleted = True
except txtorcon.TorProtocolError as e:
print(util.colors.red('Error: ') + e.what())
return
if monitor.stream_gone:
print(status)
return
print('%s (waiting for CLOSED)...' % status)
sys.stdout.flush()
await gone_d
# we're now awaiting a callback via CIRC events indicating
# that our stream has entered state CLOSED
class StreamBandwidth(object):
"""
The bandwidth-events of a single stream
"""
# __slots__ = ['_events']
def __init__(self, max_live=20, roll_up=5):
self._events = [] # list of 3-tuples
self._history = []
self.max_live = max_live
self.roll_up = roll_up
# XXX we could recursively roll-up too, i.e. spill from one
# bucket to the next. but, for now, there are precisely two
# bucks: the "current", and the first one containing (avg,
# min, max) etc
def add_bandwidth(self, epoch, read, write):
self._events.append((epoch, read, write))
self._maybe_truncate()
def _maybe_truncate(self):
"""
If we've gone past our max_live amount by at least roll_up events,
we push it into the history (possibly also truncating that).
"""
# XXX should we examine seconds here instead? i.e. have a
# max-seconds (instead of going by event-count)?
if len(self._events) > self.max_live + self.roll_up:
rolling = self._events[:self.roll_up]
self._events = self._events[self.roll_up:]
duration = float(self._events[0][0] - rolling[0][0])
mean_r = sum(x[1] for x in rolling) / duration
mean_w = sum(x[2] for x in rolling) / duration
# XXX should append some smarter-er object instead of tuple?
# (start, duration, mean_r, mean_w, max_r, max_w)
self._history.append(
(
rolling[0][0],
duration,
mean_r, mean_w,
max(x[1] for x in rolling),
max(x[2] for x in rolling),
)
)
self._history = self._history[-10:]
print("HISTORY NOW", self._history)
print("age {}, total bw {}".format(
self._events[-1][0] - (self._history[-1][0] + self._history[-1][1]),
sum([sum(x[1:]) for x in self._history]),
))
def bytes_read(self):
return sum(event[1] for event in self._events)
def bytes_written(self):
return sum(event[2] for event in self._events)
def duration(self):
if not self._events:
return 0.0
if len(self._events) == 1:
return 1.0
return float(self._events[-1][0] - self._events[0][0]) + 1.0
def rate(self):
span = self.duration()
if span == 0.0:
return (0.0, 0.0) # mmm...pragmatism
return (self.bytes_read() / span, self.bytes_written() / span)
class BandwidthMonitor(txtorcon.StreamListenerMixin):
@staticmethod
async def create(reactor, state):
bw = BandwidthMonitor(reactor, state)
await bw._setup()
return bw
def __init__(self, reactor, state):
self._reactor = reactor # just IReactorClock required?
self._state = state
self._active = {} # maps stream ID -> list-of-tuples
def stream_new(self, stream):
print("new", stream)
self._active[stream.id] = StreamBandwidth()
def stream_succeeded(self, stream):
# i think this happens when it *starts* passing data?
print("succeeded", stream, stream.target_host, stream.target_addr)
def stream_attach(self, stream, circuit):
pass
def stream_detach(self, stream, **kw):
pass
def stream_closed(self, stream, **kw):
# print("closed", stream, self._active)
if stream.id not in self._active:
print(
"Previously unknown stream to {stream.target_host} died".format(
stream=stream,
)
)
else:
bw = self._active[stream.id]
print(
"Stream {stream.id} to {stream.target_host}: {read} read, {written} written in {duration:.1f}s ({read_rate})".format(
stream=stream,
read=util.colors.green(humanize.naturalsize(bw.bytes_read())),
written=util.colors.red(humanize.naturalsize(bw.bytes_written())),
read_rate=humanize.naturalsize(sum(bw.rate())) + '/s',
duration=bw.duration(),
)
)
def stream_failed(self, stream, **kw):
pass
def _stream_bw(self, bw):
sid, wr, rd = [int(x) for x in bw.split()]
try:
bandwidth = self._active[sid]
except KeyError:
bandwidth = self._active[sid] = StreamBandwidth()
bandwidth.add_bandwidth(self._reactor.seconds(), rd, wr)
async def _setup(self):
self._state.add_stream_listener(self)
await self._state.protocol.add_event_listener('STREAM_BW', self._stream_bw)
async def monitor_streams(state, verbose):
print("monitor", state, verbose)
from twisted.internet import reactor
bw = await BandwidthMonitor.create(reactor, state)
print(bw)
await defer.Deferred()
async def run(reactor, cfg, tor, list, follow, attach, close, verbose):
state = await tor.create_state()
if attach:
await attach_streams_to_circuit(attach, state)
elif list:
await list_streams(state, verbose)
elif close:
await close_stream(state, close)
elif follow:
d = defer.succeed(None)
await monitor_streams(state, verbose)
await defer.Deferred()
