def base64_thrift(thrift_obj):
trans = TMemoryBuffer()
tbp = TBinaryProtocol(trans)
thrift_obj.write(tbp)
return b64encode(bytes(trans.getvalue())).strip()
def span_to_bytes(thrift_span):
"""
Returns a TBinaryProtocol encoded Thrift span.
:param thrift_span: thrift object to encode.
:returns: thrift object in TBinaryProtocol format bytes.
"""
transport = TMemoryBuffer()
protocol = TBinaryProtocol(transport)
thrift_span.write(protocol)
return bytes(transport.getvalue())
def encode_bytes_list(binary_thrift_obj_list): # pragma: no cover
"""
Returns a TBinaryProtocol encoded list of Thrift objects.
:param binary_thrift_obj_list: list of TBinaryProtocol objects to encode.
:returns: bynary object representing the encoded list.
"""
transport = TMemoryBuffer()
write_list_begin(transport, TType.STRUCT, len(binary_thrift_obj_list))
for thrift_bin in binary_thrift_obj_list:
transport.write(thrift_bin)
return bytes(transport.getvalue())
async def _send_data(self, data: List[Record]) -> bool:
if not data:
return True
batch = self.jaeger_thrift.Batch()
process = self.jaeger_thrift.Process()
process.serviceName = data[0].service_name
batch.process = process
spans = []
for record in data:
span = record.asthrift(self.jaeger_thrift)
spans.append(span)
batch.spans = spans
otrans = TMemoryBuffer()
self._binary.get_protocol(otrans).write_struct(batch)
resp = await self._session.post(self._address, data=otrans.getvalue())
try:
resp.raise_for_status()
except ClientResponseError:
text = await resp.text()
logger.exception("Can not send spans to jaeger: %r", text)
return False
return True
def test_buffered_read():
s = TMemoryBuffer()
t = TCyBufferedTransport(s)
t.write(b"ping")
t.flush()
assert t.read(4) == b"ping"
def _decode_binary_thrift_objs(obj):
spans = []
trans = TMemoryBuffer(obj)
_, size = read_list_begin(trans)
for _ in range(size):
span = zipkin_core.Span()
span.read(TBinaryProtocol(trans))
spans.append(span)
return spans
def do_POST(self):
content_len = int(self.headers['Content-Length'])
buf = BytesIO(self.rfile.read(content_len))
itrans = TFileObjectTransport(buf)
itrans = thttpserver.itrans_factory.get_transport(itrans)
iprot = thttpserver.iprot_factory.get_protocol(itrans)
otrans = TMemoryBuffer()
oprot = thttpserver.oprot_factory.get_protocol(otrans)
try:
thttpserver.processor.process(iprot, oprot)
except ResponseException as exn:
exn.handler(self)
else:
self.send_response(200)
self.send_header("content-type", "application/x-thrift")
self.end_headers()
self.wfile.write(otrans.getvalue())
def base64_thrift_formatter_many(parent_trace):
"""
Returns a TBinaryProtocol encoded list of Thrift objects.
:param binary_thrift_obj_list: list of TBinaryProtocol objects to encode.
:returns: bynary object representing the encoded list.
"""
traces = list(parent_trace.children())
transport = TMemoryBuffer()
write_list_begin(transport, TType.STRUCT, len(traces))
for trace in traces:
thrift_annotations = []
binary_annotations = []
for annotation in trace.annotations:
host = None
if annotation.endpoint:
host = ttypes.Endpoint(
ipv4=ipv4_to_int(annotation.endpoint.ip),
port=annotation.endpoint.port,
service_name=annotation.endpoint.service_name,
)
if annotation.annotation_type == "timestamp":
thrift_annotations.append(
ttypes.Annotation(timestamp=annotation.value,
value=annotation.name,
host=host))
else:
binary_annotations.append(
binary_annotation_formatter(annotation, host))
thrift_trace = ttypes.Span(
name=trace.name,
trace_id=u64_as_i64(trace.trace_id),
id=u64_as_i64(trace.span_id),
parent_id=u64_as_i64(trace.parent_span_id),
annotations=thrift_annotations,
binary_annotations=binary_annotations,
)
transport.write(span_to_bytes(thrift_trace))
return bytes(transport.getvalue())
def test_json_proto_api_read():
obj = TItem(id=13, phones=["5234", "12346456"])
trans = TMemoryBuffer()
p = TJSONProtocol(trans)
p.write_struct(obj)
obj2 = TItem()
obj2 = p.read_struct(obj2)
assert obj.id == 13 and obj.phones == ["5234", "12346456"]
def read_it(stream):
stream_bytes = stream.read()
transport_in = TMemoryBuffer(stream_bytes)
protocol_in = TBinaryProtocol(transport_in)
topology = StormTopology()
topology.read(protocol_in)
cls.thrift_topology = topology
cls.thrift_bolts = topology.bolts
cls.thrift_spouts = topology.spouts
# Can't reconstruct Python specs from Thrift.
cls.specs = []
def test_transport_mismatch():
s = TMemoryBuffer()
t1 = TCyBufferedTransport(s)
t1.write(b"\x80\x01\x00\x01\x00\x00\x00\x04ping hello world")
t1.flush()
with pytest.raises(TTransportException) as exc:
t2 = TCyFramedTransport(s)
t2.read(4)
assert "No frame" in str(exc.value)
def test_json_proto_api_write():
obj = TItem(id=13, phones=["5234", "12346456"])
trans = TMemoryBuffer()
p = TJSONProtocol(trans)
p.write_struct(obj)
data = trans.getvalue().decode("utf-8")
length = data[0:4]
import json
data = json.loads(data[4:])
assert length == "\x00\x00\x00S" and data == {
"metadata": {
"version": 1
},
"payload": {
"phones": ["5234", "12346456"],
"id": 13
}
}
def test_unicode_string():
class Foo(TPayload):
thrift_spec = {1: (TType.STRING, "name", False)}
default_spec = [("name", None)]
trans = TMemoryBuffer()
p = TJSONProtocol(trans)
foo = Foo(name=u('pão de açúcar'))
foo.write(p)
foo2 = Foo()
foo2.read(p)
assert foo == foo2
def do_POST(self):
# Don't care about the request path.
# Pre-read all of the data into a BytesIO. Buffered transport
# was previously configured to read everything on the first
# consumption, but that was a hack relying on the internal
# mechanism and prevents other transports from working, so
# replicate that properly to prevent timeout issues
content_len = int(self.headers['Content-Length'])
buf = BytesIO(self.rfile.read(content_len))
itrans = TFileObjectTransport(buf)
itrans = thttpserver.itrans_factory.get_transport(itrans)
iprot = thttpserver.iprot_factory.get_protocol(itrans)
otrans = TMemoryBuffer()
oprot = thttpserver.oprot_factory.get_protocol(otrans)
try:
thttpserver.processor.process(iprot, oprot)
except ResponseException as exn:
exn.handler(self)
else:
self.send_response(200)
self.send_header("content-type", "application/x-thrift")
self.end_headers()
self.wfile.write(otrans.getvalue())
def do_POST(self):
# Don't care about the request path.
itrans = TFileObjectTransport(self.rfile)
otrans = TFileObjectTransport(self.wfile)
itrans = TBufferedTransport(
itrans, int(self.headers['Content-Length']))
otrans = TMemoryBuffer()
iprot = thttpserver.iprot_factory.get_protocol(itrans)
oprot = thttpserver.oprot_factory.get_protocol(otrans)
try:
thttpserver.processor.process(iprot, oprot)
except ResponseException as exn:
exn.handler(self)
else:
self.send_response(200)
self.send_header("content-type", "application/x-thrift")
self.end_headers()
self.wfile.write(otrans.getvalue())
def decode_spans(self, spans):
"""Decodes an encoded list of spans.
:param spans: encoded list of spans
:type spans: bytes
:return: list of spans
:rtype: list of Span
"""
decoded_spans = []
transport = TMemoryBuffer(spans)
if six.byte2int(spans) == TType.STRUCT:
_, size = read_list_begin(transport)
else:
size = 1
for _ in range(size):
span = zipkin_core.Span()
span.read(TBinaryProtocol(transport))
decoded_spans.append(self._decode_thrift_span(span))
return decoded_spans
def test_thrift_transport():
test_thrift = thriftpy2.load("apache_json_test.thrift",
module_name="test_thrift")
Test = test_thrift.Test
Foo = test_thrift.Foo
test_object = Test(
tbool=False,
tbyte=16,
tdouble=1.234567,
tlong=123123123,
tshort=123,
tint=12345678,
tstr="Testing String",
tsetofints={1, 2, 3, 4, 5},
tmap_of_int2str={
1: "one",
2: "two",
3: "three"
},
tlist_of_strings=["how", "do", "i", "test", "this?"],
tmap_of_str2foo={
'first': Foo("first"),
"2nd": Foo("baz")
},
tmap_of_str2foolist={'test': [Foo("test list entry")]},
tmap_of_str2mapofstring2foo={"first": {
"second": Foo("testing")
}},
tmap_of_str2stringlist={
"words": ["dog", "cat", "pie"],
"other": ["test", "foo", "bar", "baz", "quux"]
},
tfoo=Foo("test food"),
tlist_of_foo=[Foo("1"), Foo("2"), Foo("3")],
tlist_of_maps2int=[{
"one": 1,
"two": 2,
"three": 3
}],
tmap_of_int2foo={
1: Foo("One"),
2: Foo("Two"),
5: Foo("Five")
},
tbinary=b"\x01\x0fabc123\x00\x02")
# A request generated by apache thrift that matches the above object
request_data = b"""[1,"test",1,0,{"1":{"rec":{"1":{"tf":0},"2":{"i8":16},
"3":{"i16":123},"4":{"i32":12345678},"5":{"i64":123123123},"6":
{"dbl":1.234567},"7":{"str":"Testing String"},"8":{"lst":["str",5,
"how","do","i","test","this?"]},"9":{"map":["i32","str",3,{"1":"one",
"2":"two","3":"three"}]},"10":{"set":["i32",5,1,2,3,4,5]},
"11":{"map":["str","rec",2,{"first":{"1":{"str":"first"}},"2nd":
{"1":{"str":"baz"}}}]},"12":{"map":["str","lst",
2,{"words":["str",3,"dog","cat","pie"],"other":["str",5,"test",
"foo","bar","baz","quux"]}]},"13":{"map":["str",
"map",1,{"first":["str","rec",1,{"second":{"1":{"str":"testing"}}}]}]},
"14":{"lst":["rec",3,{"1":{"str":"1"}},
{"1":{"str":"2"}},{"1":{"str":"3"}}]},"15":{"rec":{"1":{
"str":"test food"}}},"16":{"lst":["map",1,["str","i32",
3,{"one":1,"two":2,"three":3}]]},"17":{"map":["str","lst",1,{"test":
["rec",1,{"1":{"str":"test list entry"}}]}]},
"18":{"map":["i32","rec",3,{"1":{"1":{"str":"One"}},"2":{"1":
{"str":"Two"}},"5":{"1":{"str":"Five"}}}]},
"19":{"str":"AQ9hYmMxMjMAAg=="}}}}]"""
class Handler:
@staticmethod
def test(t):
# t should match the object above
assert recursive_vars(t) == recursive_vars(test_object)
return t
tp2_thrift_processor = TProcessor(test_thrift.TestService, Handler())
tp2_factory = TApacheJSONProtocolFactory()
iprot = tp2_factory.get_protocol(TMemoryBuffer(request_data))
obuf = TMemoryBuffer()
oprot = tp2_factory.get_protocol(obuf)
tp2_thrift_processor.process(iprot, oprot)
# output buffers should be the same
final_data = obuf.getvalue()
assert json.loads(request_data.decode('utf8'))[4]['1'] == \
json.loads(final_data.decode('utf8'))[4]['0']
def build_package(cls, fbns_auth: FBNSAuth, clean_session, keepalive, protocol, will_message=None, **kwargs):
keepalive = 900
connect_payload = thrift.Connect()
connect_payload.clientIdentifier = fbns_auth.clientId
client_info = thrift.ClientInfo()
client_info.userId = fbns_auth.userId
client_info.userAgent = '[FBAN/MQTT;FBAV/64.0.0.14.96;FBBV/125398467;FBDM/{density=4.0,width=1440,height=2392};FBLC/en_US;FBCR/;FBMF/LGE;FBBD/lge;FBPN/com.instagram.android;FBDV/RS988;FBSV/6.0.1;FBLR/0;FBBK/1;FBCA/armeabi-v7a:armeabi;]'
client_info.clientCapabilities = 439
client_info.endpointCapabilities = 128
client_info.publishFormat = 1
client_info.noAutomaticForeground = True
client_info.makeUserAvailableInForeground = False
client_info.deviceId = fbns_auth.deviceId
client_info.isInitiallyForeground = False
client_info.networkType = 1
client_info.networkSubtype = 0
last_monday = datetime.now() - relativedelta(weekday=calendar.MONDAY, hour=0, minute=0, second=0, microsecond=0)
last_monday = last_monday.timestamp()
session_id = int((time.time() - last_monday) * 1000)
client_info.clientMqttSessionId = session_id
client_info.subscribeTopics = [int(FBNSMQTTClient.MESSAGE_TOPIC_ID), int(FBNSMQTTClient.REG_RESP_TOPIC_ID)]
client_info.clientType = 'device_auth'
client_info.appId = 567310203415052 # Const
client_info.deviceSecret = fbns_auth.deviceSecret
client_info.clientStack = 3
connect_payload.clientInfo = client_info
connect_payload.password = fbns_auth.password
trans = TMemoryBuffer()
p = TCompactProtocol(trans)
p.write_struct(connect_payload)
data = trans.getvalue()
prop_bytes = zlib.compress(data, 9)
remaining_length = 2 + len(protocol.proto_name) + 1 + 1 + 2
connect_flags = 0
clean_session = True
if clean_session:
connect_flags |= 0x02
connect_flags |= 0x80 # username
connect_flags |= 0x40 # password
command = MQTTCommands.CONNECT
packet = bytearray()
packet.append(command)
remaining_length += len(prop_bytes)
packet.extend(pack_variable_byte_integer(remaining_length))
packet.extend(struct.pack("!H" + str(len(protocol.proto_name)) + "sBBH",
len(protocol.proto_name),
protocol.proto_name,
protocol.proto_ver,
connect_flags,
keepalive))
packet.extend(prop_bytes)
return packet
def write_it(stream):
transport_out = TMemoryBuffer()
protocol_out = TBinaryProtocol(transport_out)
cls.thrift_topology.write(protocol_out)
transport_bytes = transport_out.getvalue()
stream.write(transport_bytes)