<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import <EOL> import abc <EOL> import socket <EOL> import ssl <EOL> import struct <EOL> import riemann_client . riemann_pb2 <EOL> HOST = '<STR_LIT:localhost>' <EOL> PORT = <NUM_LIT> <EOL> TIMEOUT = None <EOL> def socket_recvall ( socket , length , bufsize = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> data = b"<STR_LIT>" <EOL> while len ( data ) < length : <EOL> data += socket . recv ( bufsize ) <EOL> return data <EOL> class RiemannError ( Exception ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> class Transport ( object ) : <EOL> """<STR_LIT>""" <EOL> __metaclass__ = abc . ABCMeta <EOL> def __enter__ ( self ) : <EOL> self . connect ( ) <EOL> return self <EOL> def __exit__ ( self , exc_type , exc_value , traceback ) : <EOL> self . disconnect ( ) <EOL> @ abc . abstractmethod <EOL> def connect ( self ) : <EOL> pass <EOL> @ abc . abstractmethod <EOL> def disconnect ( self ) : <EOL> pass <EOL> @ abc . abstractmethod <EOL> def send ( self ) : <EOL> pass <EOL> class SocketTransport ( Transport ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , host = HOST , port = PORT ) : <EOL> self . host = host <EOL> self . port = port <EOL> @ property <EOL> def address ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . host , self . port <EOL> @ property <EOL> def socket ( self ) : <EOL> """<STR_LIT>""" <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> return self . _socket <EOL> @ socket . setter <EOL> def socket ( self , value ) : <EOL> self . _socket = value <EOL> class UDPTransport ( SocketTransport ) : <EOL> def connect ( self ) : <EOL> """<STR_LIT>""" <EOL> self . socket = socket . socket ( socket . AF_INET , socket . SOCK_DGRAM ) <EOL> def disconnect ( self ) : <EOL> """<STR_LIT>""" <EOL> self . socket . close ( ) <EOL> def send ( self , message ) : <EOL> """<STR_LIT>""" <EOL> self . socket . sendto ( message . SerializeToString ( ) , self . address ) <EOL> return None <EOL> class TCPTransport ( SocketTransport ) : <EOL> def __init__ ( self , host = HOST , port = PORT , timeout = TIMEOUT ) : <EOL> """<STR_LIT>""" <EOL> super ( TCPTransport , self ) . __init__ ( host , port ) <EOL> self . timeout = timeout <EOL> def connect ( self ) : <EOL> """<STR_LIT>""" <EOL> self . socket = socket . create_connection ( self . address , self . timeout ) <EOL> def disconnect ( self ) : <EOL> """<STR_LIT>""" <EOL> self . socket . close ( ) <EOL> def send ( self , message ) : <EOL> """<STR_LIT>""" <EOL> message = message . SerializeToString ( ) <EOL> self . socket . sendall ( struct . pack ( '<STR_LIT>' , len ( message ) ) + message ) <EOL> length = struct . unpack ( '<STR_LIT>' , self . socket . recv ( <NUM_LIT:4> ) ) [ <NUM_LIT:0> ] <EOL> response = riemann_client . riemann_pb2 . Msg ( ) <EOL> response . ParseFromString ( socket_recvall ( self . socket , length ) ) <EOL> if not response . ok : <EOL> raise RiemannError ( response . error ) <EOL> return response <EOL> class TLSTransport ( TCPTransport ) : <EOL> def __init__ ( self , host = HOST , port = PORT , timeout = TIMEOUT , ca_certs = None ) : <EOL> """<STR_LIT>""" <EOL> super ( TLSTransport , self ) . __init__ ( host , port , timeout ) <EOL> self . ca_certs = ca_certs <EOL> def connect ( self ) : <EOL> """<STR_LIT>""" <EOL> super ( TLSTransport , self ) . connect ( ) <EOL> self . socket = ssl . wrap_socket ( <EOL> self . socket , <EOL> ssl_version = ssl . PROTOCOL_TLSv1 , <EOL> cert_reqs = ssl . CERT_REQUIRED , <EOL> ca_certs = self . ca_certs ) <EOL> class BlankTransport ( Transport ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . events = [ ] <EOL> def connect ( self ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def send ( self , message ) : <EOL> """<STR_LIT>""" <EOL> for event in message . events : <EOL> self . events . append ( event ) <EOL> reply = riemann_client . riemann_pb2 . Msg ( ) <EOL> reply . ok = True <EOL> return reply <EOL> def disconnect ( self ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def __len__ ( self ) : <EOL> return len ( self . events ) </s>
