def __init__(self, protocol_id, frame_type):
self.protocol_id = protocol_id
self.frame_type = frame_type
self.id = "%X" % id(self)
self._tracing = set()
self.tracing(*os.environ.get("AMQP_TRACE", "").split())
self.multiline = False
self.input = Buffer()
self.output = Buffer(struct.pack(PROTO_HDR_FMT, "AMQP", self.protocol_id, 1, 0, 0))
self.state = self.__proto_header
def __init__(self, channel):
self.channel = channel
self.new_connection = self.raise_channel_not_ready
self.close_connection = self.raise_channel_not_ready
self.register_connection = self.raise_channel_not_ready
self.overt_connection = None
self._covert_buffers = {}
self._buffer = Buffer()
self.pending_connections = {}
self.connections = {}
self.on_channel_ready = Deferred()
class OvertGateway(protocol.ClientFactory):
protocol = OvertConnection
def __init__(self, channel):
self.channel = channel
self.new_connection = self.raise_channel_not_ready
self.close_connection = self.raise_channel_not_ready
self.register_connection = self.raise_channel_not_ready
self.overt_connection = None
self._covert_buffers = {}
self._buffer = Buffer()
self.pending_connections = {}
self.connections = {}
self.on_channel_ready = Deferred()
def start_channel(self):
reactor.connectTCP(self.channel.host, self.channel.port, self)
def channel_connected(self, channel):
self.overt_connection = channel
log.info("Initializing overt channel...")
d = self.overt_connection.establish_tls()
d.addCallback(self.init_relay)
def channel_ready(self, *args):
log.info("Channel Ready")
self.new_connection = self._new_connection
self.register_connection = self._register_connection
self.close_connection = self._close_connection
self.on_channel_ready.callback(None)
def connection_established(self, conn_uuid, conn_id):
d = self.pending_connections[conn_uuid]
d.callback(conn_id)
del self.pending_connections[conn_uuid]
def connection_data_received(self, connid, data):
self.connections[connid].data_received(data)
def connection_closed(self, connid):
self.connections[connid].remote_closed()
# TODO:
# del self.connections[connid]
def init_relay(self, *args):
log.info("Initializing relay...")
# if not self.channel.support_upstream:
# log.info("Non need, channel doesn't support upstream")
# self.on_channel_ready.callback(None)
# return
data = self.overt_connection.channel_uuid
self.send_covert_command(CommandFactory.initialize_relay(data),
wait_for_overt=False)
def _new_connection(self, addr, port, conn_uuid):
command = CommandFactory.new_connection(addr, port, conn_uuid)
self.send_covert_command(command)
d = Deferred()
self.pending_connections[conn_uuid] = d
return d
def _close_connection(self, conn_id):
command = CommandFactory.close_connection(conn_id)
self.send_covert_command(command)
def _register_connection(self, connid, connection):
self.connections[connid] = connection
def send_overt_request(self,
request,
use_as_covert=False,
expected_response_size=None):
use_as_covert = use_as_covert and self.channel.support_upstream
covert_size = self.channel.calculate_sendable_covert_data(
len(request)) if use_as_covert else 0
if use_as_covert and covert_size and self._buffer.has_data():
covert_data = self._buffer.read(covert_size)
wrapped_data = self.channel.wrap_message(covert_data)
log.debug("Sending {} Covert data on {}, total size: {}".format(
len(covert_data), self.channel.host, len(wrapped_data)))
self.overt_connection.send(
wrapped_data, expected_response_size=expected_response_size)
else:
log.debug("Sending {} Overt data on {}".format(
len(request), self.channel.host))
self.overt_connection.send(request)
def send_covert_data(self, data, connid, wait_for_overt=True):
# log.debug('Sending covert data for connection: {}'.format(connid))
message = CommandFactory.data(data, connid)
self.send_covert_command(message, wait_for_overt=wait_for_overt)
def send_covert_command(self, command, wait_for_overt=True):
if wait_for_overt:
self._buffer.write(command)
else:
self.overt_connection.send(self.channel.wrap_message(command))
def raise_channel_not_ready(self, *args, **kwargs):
raise ChannelNotReadyError("Channel has not been initialized yet")
class Dispatcher:
def __init__(self, protocol_id, frame_type):
self.protocol_id = protocol_id
self.frame_type = frame_type
self.id = "%X" % id(self)
self._tracing = set()
self.tracing(*os.environ.get("AMQP_TRACE", "").split())
self.multiline = False
self.input = Buffer()
self.output = Buffer(struct.pack(PROTO_HDR_FMT, "AMQP", self.protocol_id, 1, 0, 0))
self.state = self.__proto_header
def tracing(self, *args, **kwargs):
names = set(args)
for n in kwargs:
if kwargs[n]:
names.add(n)
if "err" not in kwargs:
names.add("err")
self._tracing = names
def trace(self, categories, format, *args):
if isinstance(categories, basestring):
categories = (categories,)
for category in categories:
if category in self._tracing:
prefix = "[%s %s]" % (self.id, category)
if args:
message = format % args
else:
message = format
print >>sys.stderr, prefix, message.replace(os.linesep, "%s%s " % (os.linesep, prefix))
break
def write(self, bytes):
self.trace("raw", "RECV: %r", bytes)
self.input.write(bytes)
self.state = parse(self.state)
def closed(self):
self.trace(("raw", "frm"), "CLOSED")
def error(self, exc):
pass
def __proto_header(self):
if self.input.pending() >= PROTO_HDR_SIZE:
hdr = self.input.read(PROTO_HDR_SIZE)
magic, proto, major, minor, revision = struct.unpack(PROTO_HDR_FMT, hdr)
if (magic, proto, major, minor, revision) == ("AMQP", self.protocol_id, 1, 0, 0):
return self.__framing
else:
raise ValueError("bad protocol header")
def __framing(self):
while True:
f, n = decode(self.input.peek())
if f:
self.input.read(n)
state = self.process_frame(f)
if state is not None:
return state
else:
break
def process_frame(self, f):
body, remainder = self.type_decoder.decode(f.payload)
body.payload = remainder
self.trace("frm", "RECV[%s]: %s", f.channel, body.format(self.multiline))
return getattr(self, "do_%s" % body.NAME, self.unhandled)(f.channel, body)
def post_frame(self, channel, body):
self.trace("frm", "SENT[%s]: %s", channel, body.format(self.multiline))
encoded = self.type_encoder.encode(body)
if body.payload:
encoded += body.payload
f = Frame(self.frame_type, channel, None, encoded)
self.output.write(encode(f))
def read(self, n=None):
self.tick()
result = self.output.read(n)
self.trace("raw", "SENT: %r", result)
return result
def peek(self, n=None):
return self.output.peek(n)
def pending(self):
self.tick()
return self.output.pending()
