<s> from __future__ import division <EOL> from base64 import b64encode <EOL> from collections import deque <EOL> from io import BytesIO <EOL> import logging <EOL> from os import makedirs , open as os_open , write as os_write , close as os_close , O_CREAT , O_APPEND , O_WRONLY <EOL> from os . path import dirname , isfile <EOL> from select import select <EOL> from socket import create_connection , SHUT_RDWR , error as SocketError <EOL> from struct import pack as struct_pack , unpack as struct_unpack , unpack_from as struct_unpack_from <EOL> import errno <EOL> from . constants import DEFAULT_PORT , DEFAULT_USER_AGENT , KNOWN_HOSTS , MAGIC_PREAMBLE , TRUST_DEFAULT , TRUST_ON_FIRST_USE <EOL> from . compat import hex2 <EOL> from . exceptions import ProtocolError <EOL> from . packstream import Packer , Unpacker <EOL> from . ssl_compat import SSL_AVAILABLE , HAS_SNI , SSLError <EOL> INIT = b"<STR_LIT>" <EOL> RESET = b"<STR_LIT>" <EOL> RUN = b"<STR_LIT>" <EOL> DISCARD_ALL = b"<STR_LIT>" <EOL> PULL_ALL = b"<STR_LIT>" <EOL> SUCCESS = b"<STR_LIT>" <EOL> RECORD = b"<STR_LIT>" <EOL> IGNORED = b"<STR_LIT>" <EOL> FAILURE = b"<STR_LIT>" <EOL> DETAIL = { RECORD } <EOL> SUMMARY = { SUCCESS , IGNORED , FAILURE } <EOL> message_names = { <EOL> INIT : "<STR_LIT>" , <EOL> RESET : "<STR_LIT>" , <EOL> RUN : "<STR_LIT>" , <EOL> DISCARD_ALL : "<STR_LIT>" , <EOL> PULL_ALL : "<STR_LIT>" , <EOL> SUCCESS : "<STR_LIT>" , <EOL> RECORD : "<STR_LIT>" , <EOL> IGNORED : "<STR_LIT>" , <EOL> FAILURE : "<STR_LIT>" , <EOL> } <EOL> log = logging . getLogger ( "<STR_LIT>" ) <EOL> log_debug = log . debug <EOL> log_info = log . info <EOL> log_warning = log . warning <EOL> log_error = log . error <EOL> class ChunkChannel ( object ) : <EOL> """<STR_LIT>""" <EOL> max_chunk_size = <NUM_LIT> <EOL> def __init__ ( self , sock ) : <EOL> self . socket = sock <EOL> self . raw = BytesIO ( ) <EOL> self . output_buffer = [ ] <EOL> self . output_size = <NUM_LIT:0> <EOL> self . _recv_buffer = b"<STR_LIT>" <EOL> def write ( self , b ) : <EOL> """<STR_LIT>""" <EOL> max_chunk_size = self . max_chunk_size <EOL> output_buffer = self . output_buffer <EOL> while b : <EOL> size = len ( b ) <EOL> future_size = self . output_size + size <EOL> if future_size >= max_chunk_size : <EOL> end = max_chunk_size - self . output_size <EOL> output_buffer . append ( b [ : end ] ) <EOL> self . output_size = max_chunk_size <EOL> b = b [ end : ] <EOL> self . flush ( ) <EOL> else : <EOL> output_buffer . append ( b ) <EOL> self . output_size = future_size <EOL> b = b"<STR_LIT>" <EOL> def flush ( self , end_of_message = False ) : <EOL> """<STR_LIT>""" <EOL> output_buffer = self . output_buffer <EOL> if output_buffer : <EOL> lines = [ struct_pack ( "<STR_LIT>" , self . output_size ) ] + output_buffer <EOL> else : <EOL> lines = [ ] <EOL> if end_of_message : <EOL> lines . append ( b"<STR_LIT>" ) <EOL> if lines : <EOL> self . raw . writelines ( lines ) <EOL> self . raw . flush ( ) <EOL> del output_buffer [ : ] <EOL> self . output_size = <NUM_LIT:0> <EOL> def send ( self ) : <EOL> """<STR_LIT>""" <EOL> data = self . raw . getvalue ( ) <EOL> if __debug__ : <EOL> log_debug ( "<STR_LIT>" , "<STR_LIT::>" . join ( map ( hex2 , data ) ) ) <EOL> self . socket . sendall ( data ) <EOL> self . raw . seek ( self . raw . truncate ( <NUM_LIT:0> ) ) <EOL> def _recv ( self , size ) : <EOL> remaining = size - len ( self . _recv_buffer ) <EOL> ready_to_read = None <EOL> while remaining > <NUM_LIT:0> : <EOL> b = self . socket . recv ( <NUM_LIT> ) <EOL> if b : <EOL> if __debug__ : log_debug ( "<STR_LIT>" , "<STR_LIT::>" . join ( map ( hex2 , b ) ) ) <EOL> else : <EOL> if ready_to_read is not None : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> remaining -= len ( b ) <EOL> self . _recv_buffer += b <EOL> if remaining > <NUM_LIT:0> : <EOL> ready_to_read , _ , _ = select ( ( self . socket , ) , ( ) , ( ) , <NUM_LIT:0> ) <EOL> while not ready_to_read : <EOL> ready_to_read , _ , _ = select ( ( self . socket , ) , ( ) , ( ) , <NUM_LIT:0> ) <EOL> data , self . _recv_buffer = self . _recv_buffer [ : size ] , self . _recv_buffer [ size : ] <EOL> return data <EOL> def chunk_reader ( self ) : <EOL> chunk_size = - <NUM_LIT:1> <EOL> while chunk_size != <NUM_LIT:0> : <EOL> chunk_header = self . _recv ( <NUM_LIT:2> ) <EOL> chunk_size , = struct_unpack_from ( "<STR_LIT>" , chunk_header ) <EOL> if chunk_size > <NUM_LIT:0> : <EOL> data = self . _recv ( chunk_size ) <EOL> yield data <EOL> class Response ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , connection ) : <EOL> self . connection = connection <EOL> self . complete = False <EOL> def on_record ( self , values ) : <EOL> pass <EOL> def on_success ( self , metadata ) : <EOL> pass <EOL> def on_failure ( self , metadata ) : <EOL> pass <EOL> def on_ignored ( self , metadata = None ) : <EOL> pass <EOL> class Connection ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , sock , ** config ) : <EOL> self . defunct = False <EOL> self . channel = ChunkChannel ( sock ) <EOL> self . packer = Packer ( self . channel ) <EOL> self . responses = deque ( ) <EOL> self . closed = False <EOL> user_agent = config . get ( "<STR_LIT>" , DEFAULT_USER_AGENT ) <EOL> if isinstance ( user_agent , bytes ) : <EOL> user_agent = user_agent . decode ( "<STR_LIT>" ) <EOL> self . user_agent = user_agent <EOL> try : <EOL> self . auth_dict = vars ( config [ "<STR_LIT>" ] ) <EOL> except ( KeyError , TypeError ) : <EOL> self . auth_dict = { } <EOL> self . der_encoded_server_certificate = config . get ( "<STR_LIT>" ) <EOL> def on_failure ( metadata ) : <EOL> raise ProtocolError ( metadata . get ( "<STR_LIT:message>" , "<STR_LIT>" ) ) <EOL> response = Response ( self ) <EOL> response . on_failure = on_failure <EOL> self . append ( INIT , ( self . user_agent , self . auth_dict ) , response = response ) <EOL> self . send ( ) <EOL> while not response . complete : <EOL> self . fetch ( ) <EOL> def __del__ ( self ) : <EOL> self . close ( ) <EOL> def append ( self , signature , fields = ( ) , response = None ) : <EOL> """<STR_LIT>""" <EOL> if __debug__ : <EOL> log_info ( "<STR_LIT>" , message_names [ signature ] , "<STR_LIT:U+0020>" . join ( map ( repr , fields ) ) ) <EOL> self . packer . pack_struct_header ( len ( fields ) , signature ) <EOL> for field in fields : <EOL> self . packer . pack ( field ) <EOL> self . channel . flush ( end_of_message = True ) <EOL> self . responses . append ( response ) <EOL> def reset ( self ) : <EOL> """<STR_LIT>""" <EOL> response = Response ( self ) <EOL> def on_failure ( metadata ) : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> response . on_failure = on_failure <EOL> self . append ( RESET , response = response ) <EOL> self . send ( ) <EOL> fetch = self . fetch <EOL> while not response . complete : <EOL> fetch ( ) <EOL> def send ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . closed : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> if self . defunct : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> self . channel . send ( ) <EOL> def fetch ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . closed : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> if self . defunct : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> raw = BytesIO ( ) <EOL> unpack = Unpacker ( raw ) . unpack <EOL> try : <EOL> raw . writelines ( self . channel . chunk_reader ( ) ) <EOL> except ProtocolError : <EOL> self . defunct = True <EOL> self . close ( ) <EOL> raise <EOL> raw . seek ( <NUM_LIT:0> ) <EOL> response = self . responses [ <NUM_LIT:0> ] <EOL> for signature , fields in unpack ( ) : <EOL> if __debug__ : <EOL> log_info ( "<STR_LIT>" , message_names [ signature ] , "<STR_LIT:U+0020>" . join ( map ( repr , fields ) ) ) <EOL> if signature in SUMMARY : <EOL> response . complete = True <EOL> self . responses . popleft ( ) <EOL> if signature == FAILURE : <EOL> self . reset ( ) <EOL> handler_name = "<STR_LIT>" % message_names [ signature ] . lower ( ) <EOL> try : <EOL> handler = getattr ( response , handler_name ) <EOL> except AttributeError : <EOL> pass <EOL> else : <EOL> handler ( * fields ) <EOL> raw . close ( ) <EOL> def close ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . closed : <EOL> if __debug__ : <EOL> log_info ( "<STR_LIT>" ) <EOL> self . channel . socket . close ( ) <EOL> self . closed = True <EOL> class CertificateStore ( object ) : <EOL> def match_or_trust ( self , host , der_encoded_certificate ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> class PersonalCertificateStore ( CertificateStore ) : <EOL> def __init__ ( self , path = None ) : <EOL> self . path = path or KNOWN_HOSTS <EOL> def match_or_trust ( self , host , der_encoded_certificate ) : <EOL> base64_encoded_certificate = b64encode ( der_encoded_certificate ) <EOL> if isfile ( self . path ) : <EOL> with open ( self . path ) as f_in : <EOL> for line in f_in : <EOL> known_host , _ , known_cert = line . strip ( ) . partition ( "<STR_LIT::>" ) <EOL> known_cert = known_cert . encode ( "<STR_LIT:utf-8>" ) <EOL> if host == known_host : <EOL> return base64_encoded_certificate == known_cert <EOL> try : <EOL> makedirs ( dirname ( self . path ) ) <EOL> except OSError : <EOL> pass <EOL> f_out = os_open ( self . path , O_CREAT | O_APPEND | O_WRONLY , <NUM_LIT> ) <EOL> if isinstance ( host , bytes ) : <EOL> os_write ( f_out , host ) <EOL> else : <EOL> os_write ( f_out , host . encode ( "<STR_LIT:utf-8>" ) ) <EOL> os_write ( f_out , b"<STR_LIT::>" ) <EOL> os_write ( f_out , base64_encoded_certificate ) <EOL> os_write ( f_out , b"<STR_LIT:\n>" ) <EOL> os_close ( f_out ) <EOL> return True <EOL> def connect ( host , port = None , ssl_context = None , ** config ) : <EOL> """<STR_LIT>""" <EOL> port = port or DEFAULT_PORT <EOL> if __debug__ : log_info ( "<STR_LIT>" , host , port ) <EOL> try : <EOL> s = create_connection ( ( host , port ) ) <EOL> except SocketError as error : <EOL> if error . errno == <NUM_LIT> or error . errno == <NUM_LIT> : <EOL> raise ProtocolError ( "<STR_LIT>" % ( host , port ) ) <EOL> else : <EOL> raise <EOL> if ssl_context and SSL_AVAILABLE : <EOL> if __debug__ : log_info ( "<STR_LIT>" , host ) <EOL> try : <EOL> s = ssl_context . wrap_socket ( s , server_hostname = host if HAS_SNI else None ) <EOL> except SSLError as cause : <EOL> error = ProtocolError ( "<STR_LIT>" % cause . args [ <NUM_LIT:1> ] ) <EOL> error . __cause__ = cause <EOL> raise error <EOL> else : <EOL> der_encoded_server_certificate = s . getpeercert ( binary_form = True ) <EOL> if der_encoded_server_certificate is None : <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> trust = config . get ( "<STR_LIT>" , TRUST_DEFAULT ) <EOL> if trust == TRUST_ON_FIRST_USE : <EOL> store = PersonalCertificateStore ( ) <EOL> if not store . match_or_trust ( host , der_encoded_server_certificate ) : <EOL> raise ProtocolError ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( host , KNOWN_HOSTS ) ) <EOL> else : <EOL> der_encoded_server_certificate = None <EOL> supported_versions = [ <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> handshake = [ MAGIC_PREAMBLE ] + supported_versions <EOL> if __debug__ : log_info ( "<STR_LIT>" , MAGIC_PREAMBLE , supported_versions ) <EOL> data = b"<STR_LIT>" . join ( struct_pack ( "<STR_LIT>" , num ) for num in handshake ) <EOL> if __debug__ : log_debug ( "<STR_LIT>" , "<STR_LIT::>" . join ( map ( hex2 , data ) ) ) <EOL> s . sendall ( data ) <EOL> ready_to_read , _ , _ = select ( ( s , ) , ( ) , ( ) , <NUM_LIT:0> ) <EOL> while not ready_to_read : <EOL> ready_to_read , _ , _ = select ( ( s , ) , ( ) , ( ) , <NUM_LIT:0> ) <EOL> data = s . recv ( <NUM_LIT:4> ) <EOL> data_size = len ( data ) <EOL> if data_size == <NUM_LIT:0> : <EOL> log_error ( "<STR_LIT>" ) <EOL> raise ProtocolError ( "<STR_LIT>" ) <EOL> if data_size == <NUM_LIT:4> : <EOL> if __debug__ : log_debug ( "<STR_LIT>" , "<STR_LIT::>" . join ( map ( hex2 , data ) ) ) <EOL> else : <EOL> log_error ( "<STR_LIT>" ) <EOL> raise ProtocolError ( "<STR_LIT>" % data ) <EOL> agreed_version , = struct_unpack ( "<STR_LIT>" , data ) <EOL> if __debug__ : log_info ( "<STR_LIT>" , agreed_version ) <EOL> if agreed_version == <NUM_LIT:0> : <EOL> if __debug__ : log_info ( "<STR_LIT>" ) <EOL> s . shutdown ( SHUT_RDWR ) <EOL> s . close ( ) <EOL> else : <EOL> return Connection ( s , der_encoded_server_certificate = der_encoded_server_certificate , ** config ) </s>
