confluent/confluent_client/confluent/tlvdata.py

# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2014 IBM Corporation
# Copyright 2015 Lenovo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import array
import ctypes
import ctypes.util
import confluent.tlv as tlv
import eventlet.green.socket as socket
import eventlet.green.select as select
from datetime import datetime
import json
import os
import struct

try:
    unicode
except NameError:
    unicode = str

try:
    range = xrange
except NameError:
    pass

class iovec(ctypes.Structure):   # from uio.h
    _fields_ = [('iov_base', ctypes.c_void_p),
                ('iov_len', ctypes.c_size_t)]


iovec_ptr = ctypes.POINTER(iovec)


class cmsghdr(ctypes.Structure):  # also from bits/socket.h
    _fields_ = [('cmsg_len', ctypes.c_size_t),
                ('cmsg_level', ctypes.c_int),
                ('cmsg_type', ctypes.c_int)]

    @classmethod
    def init_data(cls, cmsg_len, cmsg_level, cmsg_type, cmsg_data):
        Data = ctypes.c_ubyte * ctypes.sizeof(cmsg_data)
        class _flexhdr(ctypes.Structure):
            _fields_ = cls._fields_ + [('cmsg_data', Data)]

        datab = Data(*bytearray(cmsg_data))
        return _flexhdr(cmsg_len=cmsg_len, cmsg_level=cmsg_level,
                        cmsg_type=cmsg_type, cmsg_data=datab)


def CMSG_LEN(length):
    sizeof_cmshdr = ctypes.sizeof(cmsghdr)
    return ctypes.c_size_t(CMSG_ALIGN(sizeof_cmshdr).value + length)


SCM_RIGHTS = 1


class msghdr(ctypes.Structure):  # from bits/socket.h
    _fields_ = [('msg_name', ctypes.c_void_p),
                ('msg_namelen', ctypes.c_uint),
                ('msg_iov', ctypes.POINTER(iovec)),
                ('msg_iovlen', ctypes.c_size_t),
                ('msg_control', ctypes.c_void_p),
                ('msg_controllen', ctypes.c_size_t),
                ('msg_flags', ctypes.c_int)]


def CMSG_ALIGN(length):  # bits/socket.h
    ret = (length + ctypes.sizeof(ctypes.c_size_t) - 1
           & ~(ctypes.sizeof(ctypes.c_size_t) - 1))
    return ctypes.c_size_t(ret)


def CMSG_SPACE(length):  # bits/socket.h
    ret = CMSG_ALIGN(length).value + CMSG_ALIGN(ctypes.sizeof(cmsghdr)).value
    return ctypes.c_size_t(ret)


class ClientFile(object):
    def __init__(self, name, mode, fd):
        self.fileobject = os.fdopen(fd, mode)
        self.filename = name

libc = ctypes.CDLL(ctypes.util.find_library('c'))
recvmsg = libc.recvmsg
recvmsg.argtypes = [ctypes.c_int, ctypes.POINTER(msghdr), ctypes.c_int]
recvmsg.restype = ctypes.c_int
sendmsg = libc.sendmsg
sendmsg.argtypes = [ctypes.c_int, ctypes.POINTER(msghdr), ctypes.c_int]
sendmsg.restype = ctypes.c_size_t

def decodestr(value):
    ret = None
    try:
        ret = value.decode('utf-8')
    except UnicodeDecodeError:
        try:
            ret = value.decode('cp437')
        except UnicodeDecodeError:
            ret = value
    except AttributeError:
        return value
    return ret

def unicode_dictvalues(dictdata):
    for key in dictdata:
        if isinstance(dictdata[key], bytes):
            dictdata[key] = decodestr(dictdata[key])
        elif isinstance(dictdata[key], datetime):
            dictdata[key] = dictdata[key].strftime('%Y-%m-%dT%H:%M:%S')
        elif isinstance(dictdata[key], list):
            _unicode_list(dictdata[key])
        elif isinstance(dictdata[key], dict):
            unicode_dictvalues(dictdata[key])


def _unicode_list(currlist):
    for i in range(len(currlist)):
        if isinstance(currlist[i], str):
            currlist[i] = decodestr(currlist[i])
        elif isinstance(currlist[i], dict):
            unicode_dictvalues(currlist[i])
        elif isinstance(currlist[i], list):
            _unicode_list(currlist[i])


def send(handle, data, filehandle=None):
    if isinstance(data, unicode):
        try:
            data = data.encode('utf-8')
        except AttributeError:
            pass
    if isinstance(data, bytes) or isinstance(data, unicode):
        # plain text, e.g. console data
        tl = len(data)
        if tl == 0:
            # if you don't have anything to say, don't say anything at all
            return
        if tl < 16777216:
            # type for string is '0', so we don't need
            # to xor anything in
            handle.sendall(struct.pack("!I", tl))
        else:
            raise Exception("String data length exceeds protocol")
        handle.sendall(data)
    elif isinstance(data, dict):  # JSON currently only goes to 4 bytes
        # Some structured message, like what would be seen in http responses
        unicode_dictvalues(data)  # make everything unicode, assuming UTF-8
        sdata = json.dumps(data, ensure_ascii=False, separators=(',', ':'))
        sdata = sdata.encode('utf-8')
        tl = len(sdata)
        if tl > 16777215:
            raise Exception("JSON data exceeds protocol limits")
        # xor in the type (0b1 << 24)
        if filehandle is None:
            tl |= 16777216
            handle.sendall(struct.pack("!I", tl))
            handle.sendall(sdata)
        else:
            tl |= (2 << 24)
            handle.sendall(struct.pack("!I", tl))
            cdtype = ctypes.c_ubyte * len(sdata)
            cdata = cdtype.from_buffer(bytearray(sdata))
            ciov = iovec(iov_base=ctypes.addressof(cdata),
                         iov_len=ctypes.c_size_t(ctypes.sizeof(cdata)))
            fd = ctypes.c_int(filehandle)
            cmh = cmsghdr.init_data(
                cmsg_len=CMSG_LEN(
                    ctypes.sizeof(fd)), cmsg_level=socket.SOL_SOCKET,
                    cmsg_type=SCM_RIGHTS, cmsg_data=fd)
            mh = msghdr(msg_name=None, msg_len=0, msg_iov=iovec_ptr(ciov),
                        msg_iovlen=1, msg_control=ctypes.addressof(cmh),
                        msg_controllen=ctypes.c_size_t(ctypes.sizeof(cmh)))
            sendmsg(handle.fileno(), mh, 0)


def recvall(handle, size):
    rd = handle.recv(size)
    while len(rd) < size:
        nd = handle.recv(size - len(rd))
        if not nd:
            raise Exception("Error reading data")
        rd += nd
    return rd

def recv(handle):
    tl = handle.recv(4)
    if not tl:
        return None
    while len(tl) < 4:
        ndata = handle.recv(4 - len(tl))
        if not ndata:
            raise Exception("Error reading data")
        tl += ndata
    if len(tl) == 0:
        return None
    tl = struct.unpack("!I", tl)[0]
    if tl & 0b10000000000000000000000000000000:
        raise Exception("Protocol Violation, reserved bit set")
    # 4 byte tlv
    dlen = tl & 16777215  # grab lower 24 bits
    datatype = (tl & 2130706432) >> 24  # grab 7 bits from near beginning
    if dlen == 0:
        return None
    if datatype == tlv.Types.filehandle:
        filehandles = array.array('i')
        rawbuffer = bytearray(2048)
        pkttype = ctypes.c_ubyte * 2048
        data = pkttype.from_buffer(rawbuffer)
        cmsgsize = CMSG_SPACE(ctypes.sizeof(ctypes.c_int)).value
        cmsgarr = bytearray(cmsgsize)
        cmtype = ctypes.c_ubyte * cmsgsize
        cmsg = cmtype.from_buffer(cmsgarr)
        cmsg.cmsg_level = socket.SOL_SOCKET
        cmsg.cmsg_type = SCM_RIGHTS
        cmsg.cmsg_len = CMSG_LEN(ctypes.sizeof(ctypes.c_int))
        iov = iovec()
        iov.iov_base = ctypes.addressof(data)
        iov.iov_len = 2048
        msg = msghdr()
        msg.msg_iov = ctypes.pointer(iov)
        msg.msg_iovlen = 1
        msg.msg_control = ctypes.addressof(cmsg)
        msg.msg_controllen = ctypes.sizeof(cmsg)
        select.select([handle], [], [])
        i = recvmsg(handle.fileno(), ctypes.pointer(msg), 0)
        cdata = cmsgarr[CMSG_LEN(0).value:]
        data = rawbuffer[:i]
        if cmsg.cmsg_level == socket.SOL_SOCKET and cmsg.cmsg_type == SCM_RIGHTS:
            filehandles.fromstring(bytes(
                cdata[:len(cdata) - len(cdata) % filehandles.itemsize]))
        data = json.loads(bytes(data))
        return ClientFile(data['filename'], data['mode'], filehandles[0])
    else:
        data = handle.recv(dlen)
        while len(data) < dlen:
            ndata = handle.recv(dlen - len(data))
            if not ndata:
                raise Exception("Error reading data")
            data += ndata
    if datatype == tlv.Types.text:
        return data
    elif datatype == tlv.Types.json:
        return json.loads(data)