A tool to give us literacy of TCompactProtocol

lreader.py

# Input file must be a binary file (I recommend to use op_to_bin.py to convert int sequence to binary)
# This code partly uses the original thrift code. https://github.com/apache/thrift/blob/master/LICENSE

# This tool is made for easily understanding the complex TCompactProtocol format.
# Because of my laziness and LINE barely using, there're no implementations of MAP and LIST.

# How to use
# ./lreader.py output.bin

#!/usr/bin/env python

import struct
import binascii
import sys

def makeZigZag(n, bits):
    return (n << 1) ^ (n >> (bits - 1))

def fromZigZag(n):
    return (n >> 1) ^ -(n & 1)

class CompactType(object):
    PROTOCOL_ID = 0x82

    STOP = 0x00
    TRUE = 0x01
    FALSE = 0x02
    BYTE = 0x03
    I16 = 0x04
    I32 = 0x05
    I64 = 0x06
    DOUBLE = 0x07
    BINARY = 0x08
    LIST = 0x09
    SET = 0x0A
    MAP = 0x0B
    STRUCT = 0x0C

    types = {
            0:"STOP",
            1:"TRUE",
            2:"FALSE",
            3:"BYTE",
            4:"I16",
            5:"I32",
            6:"I64",
            7:"DOUBLE",
            8:"BINARY",
            9:"LIST",
            10:"SET",
            11:"MAP",
            12:"STRUCT",
            }

class node(object):
    
    def __init__(self, t, fid, payload):
        self.type = t
        self.fid = fid
        self.payload = payload

    def __repr__(self, level=0):
        #return "<lreader.node -> fid {} type={} {}>".format(self.fid, CompactType.types[self.type], self.payload) 
        return "<lreader.node -> fid {} type={} {}>".format(self.fid, CompactType.types[self.type], self.payload) 

    def __str__(self, level=0):
        return "<lreader.node -> fid {} type={} {}>".format(self.fid, CompactType.types[self.type], self.payload) 

class lreader(object):

    def __init__(self, fn):
        
        self.fp = open(fn, "r")
        
        if self.__ra_ubyte() != CompactType.PROTOCOL_ID:
            raise Exception("proto id is not matched, seems not to be compact protocol format")

        ver_type = self.__ra_ubyte()
        self.version = (ver_type >> 5) & 0x7
        self.type = ver_type & 0x1f

        self.seqid = self.__ra_varint()
        self.method_name = self.__ra_binary()
        self.__last_fid = 0

        self.tree = self.evaluate()
        self.fp.close()


    def __ra_varint(self):
        r = 0 
        shift = 0
        while 1:
            x = self.fp.read(1)
            byte = ord(x)
            r |= (byte & 0x7f) << shift
            if byte >> 7 == 0:
                return r
            shift += 7

    def __ra_i16(self):
        return fromZigZag(self.__ra_varint())

    def __ra_i32(self):
        return fromZigZag(self.__ra_varint())

    def __ra_i64(self):
        return fromZigZag(self.__ra_varint())

    def __ra_double(self):
        r = struct.unpack("!d", self.fp.read(8))[0]
        return r

    def __ra_ubyte(self):
        r = struct.unpack("!B", self.fp.read(1))[0]
        return r

    def __ra_byte(self):
        r = struct.unpack("!b", self.fp.read(1))[0]
        return r

    def __ra_binary(self):
        size = self.__ra_varint()
        if size < 0:
            raise Exception("length is negative")
        return self.fp.read(size)
   
    def evaluate(self):
        tree = []

        while 1:
            fid = 0
            field_type = self.__ra_ubyte()
            
            if (field_type & 0x0f) == CompactType.STOP:
                break
            
            delta = field_type >> 4
            if delta == 0:
                fid = self.__ra_i16()
            else:
                fid = self.__last_fid + delta

            self.__last_fid = fid
            field_type = field_type & 0x0f

            if field_type == CompactType.TRUE:
                n = self.true(fid)
            elif field_type == CompactType.FALSE:
                n = self.false(fid)
                tree.append(n)
            elif field_type == CompactType.BYTE:
                n = self.byte(fid)
                tree.append(n)
            elif field_type == CompactType.I16:
                n = self.I16(fid)
                tree.append(n)
            elif field_type == CompactType.I32:
                n = self.I32(fid)
                tree.append(n)
            elif field_type == CompactType.I64:
                n = self.I64(fid)
                tree.append(n)
            elif field_type == CompactType.DOUBLE:
                n = self.double(fid)
                tree.append(n)
            elif field_type == CompactType.BINARY:
                n = self.binary(fid)
                tree.append(n)
            elif field_type == CompactType.LIST:
                n = self.list(fid)
                tree.append(n)
            elif field_type == CompactType.SET:
                n = self.set(fid)
                tree.append(n)
            elif field_type == CompactType.MAP:
                n = self.map(fid)
                tree.append(n)
            elif field_type == CompactType.STRUCT:
                n = self.struct(fid)
                tree.append(n)
            else:
                raise Exception("cant recognize this type")
        
        return tree

    
    def true(self, fid):
        n = node(CompactType.TRUE, fid, True)
        return n

    def false(self, fid):
        n = node(CompactType.FALSE, fid, False)
        return n

    def I16(self, fid):
        r = self.__ra_i16()
        n = node(CompactType.I16, fid, r)
        return n
    
    def I32(self, fid):
        r = self. __ra_i32()
        n = node(CompactType.I32, fid, r)
        return n

    def I64(self, fid):
        r = self.__ra_i64()
        n = node(CompactType.I64, fid, r)
        return n
    
    def double(self, fid):
        r = self.__ra_double()
        n = node(CompactType.DOUBLE, fid, r)
        return n

    def binary(self, fid):
        r = self.__ra_binary()
        n = node(CompactType.BINARY, fid, r)
        return n

    def list(self, fid):
        raise NotImplementedError()

    def set(self, fid):
        raise NotImplementedError()

    def map(self, fid):
        raise NotImplementedError()

    def struct(self, fid):
        tree = self.evaluate()
        n = node(CompactType.STRUCT, fid, tree)
        return n
    
if __name__ == '__main__':
    l = lreader(sys.argv[1])
    print l.method_name
    print l.tree 

op_to_bin.p

# How to use
# echo "8221 0006 ......." | ./op_to_bin.py

#!/usr/bin/env python
import sys
import struct

def str_to_bin(p):
        payload = [int(n+i,16) for (n,i) in zip(p[::2], p[1::2])]
        return payload

if __name__ == '__main__':
    buf = sys.stdin.read().replace(" ","")
    byte_buf = str_to_bin(buf)
    print map(lambda x: hex(x),byte_buf)
    if len(sys.argv) == 2:
        f = open(sys.argv[1], "w")
    else:
        f = open("output.bin", "w")

    for i in byte_buf:
        f.write(struct.pack("B", i))

    f.close()