<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2009 Las Cumbres Observatory (www.lcogt.net)
# Copyright (c) 2010 <NAME>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
'''
channel.py - Socket implementation of Google's Protocol Buffers RPC
service interface.

This package contains classes providing a socket implementation of the
RPCChannel abstract class.

Original Authors: <NAME> (<EMAIL>)
         <NAME> (<EMAIL>)
         <NAME> (<EMAIL>)

May 2009, Nov 2010

Modified for snakebite: W<NAME> (<EMAIL>)

May 2012

'''

# Standard library imports
import socket
import os
import math

# Third party imports
from google.protobuf.service import RpcChannel

# Protobuf imports
from snakebite.protobuf.RpcHeader_pb2 import RpcRequestHeaderProto, RpcResponseHeaderProto
from snakebite.protobuf.IpcConnectionContext_pb2 import IpcConnectionContextProto
from snakebite.protobuf.ProtobufRpcEngine_pb2 import RequestHeaderProto
from snakebite.protobuf.datatransfer_pb2 import OpReadBlockProto, BlockOpResponseProto, PacketHeaderProto, ClientReadStatusProto

from snakebite.platformutils import get_current_username
from snakebite.formatter import format_bytes
from snakebite.errors import RequestError
from snakebite.crc32c import crc

import google.protobuf.internal.encoder as encoder
import google.protobuf.internal.decoder as decoder

# Module imports

import logger
import logging
import struct
import uuid

# Configure package logging
log = logger.getLogger(__name__)


def log_protobuf_message(header, message):
    log.debug("%s:\n\n\033[92m%s\033[0m" % (header, message))


def get_delimited_message_bytes(byte_stream, nr=4):
    ''' Parse a delimited protobuf message. This is done by first getting a protobuf varint from
    the stream that represents the length of the message, then reading that amount of
    from the message and then parse it.
    Since the int can be represented as max 4 bytes, first get 4 bytes and try to decode.
    The decoder returns the value and the position where the value was found, so we need
    to rewind the buffer to the position, because the remaining bytes belong to the message
    after.
    '''

    (length, pos) = decoder._DecodeVarint32(byte_stream.read(nr), 0)
    if log.getEffectiveLevel() == logging.DEBUG:
        log.debug("Delimited message length (pos %d): %d" % (pos, length))

    delimiter_bytes = nr - pos

    byte_stream.rewind(delimiter_bytes)
    message_bytes = byte_stream.read(length)
    if log.getEffectiveLevel() == logging.DEBUG:
        log.debug("Delimited message bytes (%d): %s" % (len(message_bytes), format_bytes(message_bytes)))

    total_len = length + pos
    return (total_len, message_bytes)


class RpcBufferedReader(object):
    '''Class that wraps a socket and provides some utility methods for reading
    and rewinding of the buffer. This comes in handy when reading protobuf varints.
    '''
    MAX_READ_ATTEMPTS = 100

    def __init__(self, socket):
        self.socket = socket
        self.reset()

    def read(self, n):
        '''Reads n bytes into the internal buffer'''
        bytes_wanted = n - self.buffer_length + self.pos + 1
        if bytes_wanted > 0:
            self._buffer_bytes(bytes_wanted)

        end_pos = self.pos + n
        ret = self.buffer[self.pos + 1:end_pos + 1]
        self.pos = end_pos
        return ret

    def _buffer_bytes(self, n):
        to_read = n
        for _ in xrange(self.MAX_READ_ATTEMPTS):
            bytes_read = self.socket.recv(to_read)
            self.buffer += bytes_read
            to_read -= len(bytes_read)
            if to_read == 0:
                log.debug("Bytes read: %d, total: %d" % (len(bytes_read), self.buffer_length))
                return n
        if len(bytes_read) < n:
            raise Exception("RpcBufferedReader only managed to read %s out of %s bytes" % (len(bytes_read), n))

    def rewind(self, places):
        '''Rewinds the current buffer to a position. Needed for reading varints,
        because we might read bytes that belong to the stream after the varint.
        '''
        log.debug("Rewinding pos %d with %d places" % (self.pos, places))
        self.pos -= places
        log.debug("Reset buffer to pos %d" % self.pos)

    def reset(self):
        self.buffer = ""
        self.pos = -1  # position of last byte read

    @property
    def buffer_length(self):
        '''Returns the length of the current buffer.'''
        return len(self.buffer)


class SocketRpcChannel(RpcChannel):
    ERROR_BYTES = 18446744073709551615L
    RPC_HEADER = "hrpc"
    RPC_SERVICE_CLASS = 0x00
    AUTH_PROTOCOL_NONE = 0x00
    AUTH_PROTOCOL_SASL = 0xDF
    RPC_PROTOCOL_BUFFFER = 0x02


    '''Socket implementation of an RpcChannel.
    '''

    def __init__(self, host, port, version, effective_user=None, use_sasl=False):
        '''SocketRpcChannel to connect to a socket server on a user defined port.
           It possible to define version and effective user for the communication.'''
        self.host = host
        self.port = port
        self.sock = None
        self.call_id = -3  # First time (when the connection context is sent, the call_id should be -3, otherwise start with 0 and increment)
        self.version = version
        self.client_id = str(uuid.uuid4())
        self.use_sasl = use_sasl
        if self.use_sasl:
            try:
                # Only import if SASL is enabled
                from snakebite.rpc_sasl import SaslRpcClient
                from snakebite.kerberos import Kerberos
            except ImportError:
                raise Exception("Kerberos libs not found. Please install snakebite using 'pip install snakebite[kerberos]'")

            kerberos = Kerberos()
            self.effective_user = effective_user or kerberos.user_principal().name
        else: 
            self.effective_user = effective_user or get_current_username()

    def validate_request(self, request):
        '''Validate the client request against the protocol file.'''

        # Check the request is correctly initialized
        if not request.IsInitialized():
            raise Exception("Client request (%s) is missing mandatory fields" % type(request))

    def get_connection(self, host, port):
        '''Open a socket connection to a given host and port and writes the Hadoop header
        The Hadoop RPC protocol looks like this when creating a connection:<commit_msg>Do conditional import higher up to make sure import stays in scope<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2009 Las Cumbres Observatory (www.lcogt.net)
# Copyright (c) 2010 <NAME>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
'''
channel.py - Socket implementation of Google's Protocol Buffers RPC
service interface.

This package contains classes providing a socket implementation of the
RPCChannel abstract class.

Original Authors: <NAME> (<EMAIL>)
         <NAME> (<EMAIL>)
         <NAME> (<EMAIL>)

May 2009, Nov 2010

Modified for snakebite: <NAME> (<EMAIL>)

May 2012

'''

# Standard library imports
import socket
import os
import math

# Third party imports
from google.protobuf.service import RpcChannel

# Protobuf imports
from snakebite.protobuf.RpcHeader_pb2 import RpcRequestHeaderProto, RpcResponseHeaderProto
from snakebite.protobuf.IpcConnectionContext_pb2 import IpcConnectionContextProto
from snakebite.protobuf.ProtobufRpcEngine_pb2 import RequestHeaderProto
from snakebite.protobuf.datatransfer_pb2 import OpReadBlockProto, BlockOpResponseProto, PacketHeaderProto, ClientReadStatusProto

from snakebite.platformutils import get_current_username
from snakebite.formatter import format_bytes
from snakebite.errors import RequestError
from snakebite.crc32c import crc

import google.protobuf.internal.encoder as encoder
import google.protobuf.internal.decoder as decoder

# Module imports

import logger
import logging
import struct
import uuid

_kerberos_available = False
try:
    from snakebite.rpc_sasl import SaslRpcClient
    from snakebite.kerberos import Kerberos
    _kerberos_available = True
except ImportError:
    pass

# Configure package logging
log = logger.getLogger(__name__)


def log_protobuf_message(header, message):
    log.debug("%s:\n\n\033[92m%s\033[0m" % (header, message))


def get_delimited_message_bytes(byte_stream, nr=4):
    ''' Parse a delimited protobuf message. This is done by first getting a protobuf varint from
    the stream that represents the length of the message, then reading that amount of
    from the message and then parse it.
    Since the int can be represented as max 4 bytes, first get 4 bytes and try to decode.
    The decoder returns the value and the position where the value was found, so we need
    to rewind the buffer to the position, because the remaining bytes belong to the message
    after.
    '''

    (length, pos) = decoder._DecodeVarint32(byte_stream.read(nr), 0)
    if log.getEffectiveLevel() == logging.DEBUG:
        log.debug("Delimited message length (pos %d): %d" % (pos, length))

    delimiter_bytes = nr - pos

    byte_stream.rewind(delimiter_bytes)
    message_bytes = byte_stream.read(length)
    if log.getEffectiveLevel() == logging.DEBUG:
        log.debug("Delimited message bytes (%d): %s" % (len(message_bytes), format_bytes(message_bytes)))

    total_len = length + pos
    return (total_len, message_bytes)


class RpcBufferedReader(object):
    '''Class that wraps a socket and provides some utility methods for reading
    and rewinding of the buffer. This comes in handy when reading protobuf varints.
    '''
    MAX_READ_ATTEMPTS = 100

    def __init__(self, socket):
        self.socket = socket
        self.reset()

    def read(self, n):
        '''Reads n bytes into the internal buffer'''
        bytes_wanted = n - self.buffer_length + self.pos + 1
        if bytes_wanted > 0:
            self._buffer_bytes(bytes_wanted)

        end_pos = self.pos + n
        ret = self.buffer[self.pos + 1:end_pos + 1]
        self.pos = end_pos
        return ret

    def _buffer_bytes(self, n):
        to_read = n
        for _ in xrange(self.MAX_READ_ATTEMPTS):
            bytes_read = self.socket.recv(to_read)
            self.buffer += bytes_read
            to_read -= len(bytes_read)
            if to_read == 0:
                log.debug("Bytes read: %d, total: %d" % (len(bytes_read), self.buffer_length))
                return n
        if len(bytes_read) < n:
            raise Exception("RpcBufferedReader only managed to read %s out of %s bytes" % (len(bytes_read), n))

    def rewind(self, places):
        '''Rewinds the current buffer to a position. Needed for reading varints,
        because we might read bytes that belong to the stream after the varint.
        '''
        log.debug("Rewinding pos %d with %d places" % (self.pos, places))
        self.pos -= places
        log.debug("Reset buffer to pos %d" % self.pos)

    def reset(self):
        self.buffer = ""
        self.pos = -1  # position of last byte read

    @property
    def buffer_length(self):
        '''Returns the length of the current buffer.'''
        return len(self.buffer)


class SocketRpcChannel(RpcChannel):
    ERROR_BYTES = 18446744073709551615L
    RPC_HEADER = "hrpc"
    RPC_SERVICE_CLASS = 0x00
    AUTH_PROTOCOL_NONE = 0x00
    AUTH_PROTOCOL_SASL = 0xDF
    RPC_PROTOCOL_BUFFFER = 0x02


    '''Socket implementation of an RpcChannel.
    '''

    def __init__(self, host, port, version, effective_user=None, use_sasl=False):
        '''SocketRpcChannel to connect to a socket server on a user defined port.
           It possible to define version and effective user for the communication.'''
        self.host = host
        self.port = port
        self.sock = None
        self.call_id = -3  # First time (when the connection context is sent, the call_id should be -3, otherwise start with 0 and increment)
        self.version = version
        self.client_id = str(uuid.uuid4())
        self.use_sasl = use_sasl
        if self.use_sasl:
            if not _kerberos_available:
                raise Exception("Kerberos libs not found. Please install snakebite using 'pip install snakebite[kerberos]'")

            kerberos = Kerberos()
            self.effective_user = effective_user or kerberos.user_principal().name
        else: 
            self.effective_user = effective_user or get_current_username()

    def validate_request(self, request):
        '''Validate the client request against the protocol file.'''

        # Check the request is correctly initialized
        if not request.IsInitialized():
            raise Exception("Client request (%s) is missing mandatory fields" % type(request))

    def get_connection(self, host, port):
        '''Open a socket connection to a given host and port and writes the Hadoop header
        The Hadoop RPC protocol looks like this when creating a connection:
{'fraction_non_alphanumeric': 0.07504363001745201, 'fraction_numerical': 0.011410927641294134, 'mean_word_length': 3.2471493728620295, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2931069', 'n_tokens_mistral': 4336, 'n_tokens_neox': 4078, 'n_words': 1576}