def udp_loop():
    """
    Event loop. Wait for messages and then pretty-print them
    """
    while True:
        print("Waiting for UDP message")
        raw_message, address = udp_sock.recvfrom(2**16)
        # All UDP packets contain compact GPB messages
        decode_gpb_compact(raw_message, args)
def udp_loop():
    """
    Event loop. Wait for messages and then pretty-print them
    """
    while True:
        print("Waiting for UDP message")
        raw_message, address = udp_sock.recvfrom(2**16)
        # All UDP packets contain compact GPB messages
        decode_gpb_compact(raw_message, args)
def get_message(conn, deco):
    """
    Handle a receved TCP message
    Argument: conn
      TCP connection
    Argument: deco
      ZLIB decompression object
    Return: updated decompression object in the case where compression was
            reset
    """
    print("Getting TCP message")
    # v1 message header (from XR6.0) consists of just a 4-byte length
    # v2 message header (from XR6.1 onwards) consists of 3 4-byte fields:
    #     Type,Flags,Length
    # If the first 4 bytes read is <=4 then it is too small to be a 
    # valid length. Assume it is v2 instead
    # 
    t = conn.recv(4)
    msg_type = unpack_int(t)
    if msg_type > 4:
        # V1 message - compressed JSON
        flags = TCP_FLAG_ZLIB_COMPRESSION
        msg_type_str = "JSONv1 (COMPRESSED)"
        length = msg_type
        msg_type = TCPMsgType.JSON
        print("  Message Type: {}".format(msg_type_str))
        return get_v1_message(length, conn)
    else:
        # V2 message
        try:
            msg_type_str = TCPMsgType.to_string(msg_type)
            print("  Message Type: {})".format(msg_type_str))
        except:
            print("  Invalid Message type: {}".format(msg_type))
    
        t = conn.recv(4)
        flags = unpack_int(t)
        print("  Flags: {}".format(tcp_flags_to_string(flags)))
        t = conn.recv(4)
        length = unpack_int(t)
    print("  Length: {}".format(length))
   
    # Read all the bytes of the message according to the length in the header 
    data = b""
    while len(data) < length:
        data += conn.recv(length - len(data))

    # Decompress the message if necessary. Otherwise use as-is
    if flags & TCP_FLAG_ZLIB_COMPRESSION != 0:
        try:
            print("Decompressing message")
            msg = deco.decompress(data)
        except Exception as e:
            print("ERROR: failed to decompress message: {}".format(e))
            msg = None
    else:
        msg = data

    # Decode the data according to the message type in the header
    print("Decoding message")
    try:
        if msg_type == TCPMsgType.GPB_COMPACT:
            decode_gpb_compact(msg, args)
        elif msg_type == TCPMsgType.GPB_KEY_VALUE:
            decode_gpb_kv(msg, args)
        elif msg_type == TCPMsgType.JSON:
            if args.json_dump:
                # Print the message as-is
                print(msg)
            else:
                # Decode and pretty-print the message
                decode_json(msg, args)
        elif msg_type == TCPMsgType.RESET_COMPRESSOR:
            deco = zlib.decompressobj()
    except Exception as e:
        print("ERROR: failed to decode TCP message: {}".format(e))

    return deco
def get_message(conn, deco):
    """
    Handle a receved TCP message
    Argument: conn
      TCP connection
    Argument: deco
      ZLIB decompression object
    Return: updated decompression object in the case where compression was
            reset
    """
    print("Getting TCP message")
    # v1 message header (from XR6.0) consists of just a 4-byte length
    # v2 message header (from XR6.1 onwards) consists of 3 4-byte fields:
    #     Type,Flags,Length
    # If the first 4 bytes read is <=4 then it is too small to be a
    # valid length. Assume it is v2 instead
    #
    t = conn.recv(4)
    msg_type = unpack_int(t)
    if msg_type > 4:
        # V1 message - compressed JSON
        flags = TCP_FLAG_ZLIB_COMPRESSION
        msg_type_str = "JSONv1 (COMPRESSED)"
        length = msg_type
        msg_type = TCPMsgType.JSON
        print("  Message Type: {}".format(msg_type_str))
        return get_v1_message(length, conn)
    else:
        # V2 message
        try:
            msg_type_str = TCPMsgType.to_string(msg_type)
            print("  Message Type: {})".format(msg_type_str))
        except:
            print("  Invalid Message type: {}".format(msg_type))

        t = conn.recv(4)
        flags = unpack_int(t)
        print("  Flags: {}".format(tcp_flags_to_string(flags)))
        t = conn.recv(4)
        length = unpack_int(t)
    print("  Length: {}".format(length))

    # Read all the bytes of the message according to the length in the header
    data = b""
    while len(data) < length:
        data += conn.recv(length - len(data))

    # Decompress the message if necessary. Otherwise use as-is
    if flags & TCP_FLAG_ZLIB_COMPRESSION != 0:
        try:
            print("Decompressing message")
            msg = deco.decompress(data)
        except Exception as e:
            print("ERROR: failed to decompress message: {}".format(e))
            msg = None
    else:
        msg = data

    # Decode the data according to the message type in the header
    print("Decoding message")
    try:
        if msg_type == TCPMsgType.GPB_COMPACT:
            decode_gpb_compact(msg, args)
        elif msg_type == TCPMsgType.GPB_KEY_VALUE:
            decode_gpb_kv(msg, args)
        elif msg_type == TCPMsgType.JSON:
            if args.json_dump:
                # Print the message as-is
                print(msg)
            else:
                # Decode and pretty-print the message
                decode_json(msg, args)
        elif msg_type == TCPMsgType.RESET_COMPRESSOR:
            deco = zlib.decompressobj()
    except Exception as e:
        print("ERROR: failed to decode TCP message: {}".format(e))

    return deco