def thrift_obj_in_bytes(thrift_obj):
"""
Encodes a Thrift object using TBinaryProtocol.
:param thrift_obj
:returns: TBinaryProtocol bytes represenation of thrift_obj.
"""
trans = TMemoryBuffer()
thrift_obj.write(TBinaryProtocol(trans))
return trans.getvalue()
return bytes(trans.getvalue())
def read_it(stream):
stream_bytes = stream.read()
transport_in = TMemoryBuffer(stream_bytes)
protocol_in = TBinaryProtocol(transport_in)
topology = storm_thrift.StormTopology()
topology.read(protocol_in)
cls._topology = topology
cls.thrift_bolts = topology.bolts
cls.thrift_spouts = topology.spouts
# Can't reconstruct Python specs from Thrift.
cls.specs = []
def thrift_obj_in_bytes(thrift_obj): # pragma: no cover
"""
Returns TBinaryProtocol encoded Thrift object.
:param thrift_obj: thrift object to encode
:returns: thrift object in TBinaryProtocol format bytes.
"""
trans = TMemoryBuffer()
thrift_obj.write(TBinaryProtocol(trans))
return bytes(trans.getvalue())
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 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 test_skip_struct():
b = TCyBufferedTransport(TMemoryBuffer())
p = proto.TCyBinaryProtocol(b)
item = TItem(id=123, phones=["123456", "abcdef"])
p.write_struct(item)
p.write_message_end()
proto.write_val(b, TType.I32, 123)
b.flush()
proto.skip(b, TType.STRUCT)
assert 123 == proto.read_val(b, TType.I32)
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())
def thrift_objs_in_bytes(thrift_obj_list): # pragma: no cover
"""
Returns TBinaryProtocol encoded Thrift objects.
:param thrift_obj_list: thrift objects list to encode
:returns: thrift objects in TBinaryProtocol format bytes.
"""
transport = TMemoryBuffer()
protocol = TBinaryProtocol(transport)
write_list_begin(transport, TType.STRUCT, len(thrift_obj_list))
for thrift_obj in thrift_obj_list:
thrift_obj.write(protocol)
return bytes(transport.getvalue())
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_write_wrong_arg_type():
trans = TCyBufferedTransport(TMemoryBuffer())
b = proto.TCyBinaryProtocol(trans)
item = TItem(id="wrong type", phones=["123456", "abcdef"])
try:
b.write_struct(item)
except Exception:
pass
b.write_message_end()
item2 = TItem(id=123, phones=["123456", "abcdef"])
b.write_struct(item2)
b.write_message_end()
assert ("08 00 01 00 00 00 7b 0f 00 02 0b 00 00 00 02 00 00 00 "
"06 31 32 33 34 35 36 00 00 00 06 61 62 63 64 65 66 00") == \
hexlify(trans.getvalue())
def test_read_huge_args():
class Hello(TPayload):
thrift_spec = {
1: (TType.STRING, "name"),
2: (TType.STRING, "world"),
}
default_spec = [("name", None), ("world", None)]
b = TMemoryBuffer()
b = TCyBufferedTransport(b)
item = Hello(name='我' * 326, world='你' * 1365)
p = proto.TCyBinaryProtocol(b)
p.write_struct(item)
p.write_message_end()
item2 = Hello()
p.read_struct(item2)
def test_unicode_string():
class Foo(TPayload):
thrift_spec = {
1: (TType.STRING, "name")
}
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.
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 span builders
:rtype: list
"""
span_builders = []
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))
span_builders.append(self._decode_thrift_span(span))
return span_builders
def test_unpack_i8():
b = TMemoryBuffer(b"{")
p = proto.TCyBinaryProtocol(b)
assert_equal(123, p.read_val(TType.I08))
def write_it(stream):
transport_out = TMemoryBuffer()
protocol_out = TBinaryProtocol(transport_out)
cls._topology.write(protocol_out)
transport_bytes = transport_out.getvalue()
stream.write(transport_bytes)
def test_unpack_string():
b = TMemoryBuffer(b'\x00\x00\x00\x0c'
b'\xe4\xbd\xa0\xe5\xa5\xbd\xe4\xb8\x96\xe7\x95\x8c')
p = proto.TCyBinaryProtocol(b)
assert_equal(u("你好世界"), p.read_val(TType.STRING))
def _decode_binary_thrift_obj(obj):
trans = TMemoryBuffer(obj)
span = zipkin_core.Span()
span.read(TBinaryProtocol(trans))
return span
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, level=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 test_pack_double():
b = TMemoryBuffer()
p = proto.TCyBinaryProtocol(b)
p.write_val(TType.DOUBLE, 1234567890.1234567890)
p.write_message_end()
assert_equal("41 d2 65 80 b4 87 e6 b7", hexlify(b.getvalue()))
def test_unpack_double():
b = TMemoryBuffer(b'A\xd2e\x80\xb4\x87\xe6\xb7')
p = proto.TCyBinaryProtocol(b)
assert_equal(1234567890.1234567890, p.read_val(TType.DOUBLE))
def test_pack_i64():
b = TMemoryBuffer()
p = proto.TCyBinaryProtocol(b)
p.write_val(TType.I64, 1234567890123456789)
p.write_message_end()
assert_equal("11 22 10 f4 7d e9 81 15", hexlify(b.getvalue()))
def test_unpack_i64():
b = TMemoryBuffer(b'\x11"\x10\xf4}\xe9\x81\x15')
p = proto.TCyBinaryProtocol(b)
assert_equal(1234567890123456789, p.read_val(TType.I64))
def test_unpack_i32():
b = TMemoryBuffer(b'I\x96\x02\xd2')
p = proto.TCyBinaryProtocol(b)
assert_equal(1234567890, p.read_val(TType.I32))
def test_pack_i32():
b = TMemoryBuffer()
p = proto.TCyBinaryProtocol(b)
p.write_val(TType.I32, 1234567890)
p.write_message_end()
assert_equal("49 96 02 d2", hexlify(b.getvalue()))
def test_unpack_i16():
b = TMemoryBuffer(b"09")
p = proto.TCyBinaryProtocol(b)
assert_equal(12345, p.read_val(TType.I16))
def test_read_message_begin():
b = TMemoryBuffer(b'\x80\x01\x00\x0b\x00\x00\x00\x04test\x00\x00\x00\x01')
res = proto.TCyBinaryProtocol(b).read_message_begin()
assert_equal(res, ("test", TType.STRING, 1))
def serialize(thrift_object, protocol_factory=TBinaryProtocolFactory()):
transport = TMemoryBuffer()
protocol = protocol_factory.get_protocol(transport)
thrift_object.write(protocol)
return transport.getvalue()
def deserialize(base, buf, protocol_factory=TBinaryProtocolFactory()):
transport = TMemoryBuffer(buf)
protocol = protocol_factory.get_protocol(transport)
base.read(protocol)
return base
def test_pack_i16():
b = TMemoryBuffer()
p = proto.TCyBinaryProtocol(b)
p.write_val(TType.I16, 12345)
p.write_message_end()
assert_equal("30 39", hexlify(b.getvalue()))
