<s> from __future__ import ( absolute_import , division , generators , nested_scopes , print_function , <EOL> unicode_literals , with_statement ) <EOL> import struct <EOL> import six <EOL> class ChunkType ( object ) : <EOL> """<STR_LIT>""" <EOL> ARGUMENT = b'<STR_LIT:A>' <EOL> ENVIRONMENT = b'<STR_LIT:E>' <EOL> WORKING_DIR = b'<STR_LIT:D>' <EOL> COMMAND = b'<STR_LIT:C>' <EOL> PID = b'<STR_LIT:P>' <EOL> STDIN = b'<STR_LIT:0>' <EOL> STDOUT = b'<STR_LIT:1>' <EOL> STDERR = b'<STR_LIT:2>' <EOL> START_READING_INPUT = b'<STR_LIT:S>' <EOL> STDIN_EOF = b'<STR_LIT:.>' <EOL> EXIT = b'<STR_LIT:X>' <EOL> REQUEST_TYPES = ( ARGUMENT , ENVIRONMENT , WORKING_DIR , COMMAND ) <EOL> EXECUTION_TYPES = ( PID , STDIN , STDOUT , STDERR , START_READING_INPUT , STDIN_EOF , EXIT ) <EOL> VALID_TYPES = REQUEST_TYPES + EXECUTION_TYPES <EOL> class NailgunProtocol ( object ) : <EOL> """<STR_LIT>""" <EOL> ENVIRON_SEP = '<STR_LIT:=>' <EOL> TTY_ENV_TMPL = '<STR_LIT>' <EOL> HEADER_FMT = b'<STR_LIT>' <EOL> HEADER_BYTES = <NUM_LIT:5> <EOL> class ProtocolError ( Exception ) : <EOL> """<STR_LIT>""" <EOL> class TruncatedRead ( ProtocolError ) : <EOL> """<STR_LIT>""" <EOL> class TruncatedHeaderError ( TruncatedRead ) : <EOL> """<STR_LIT>""" <EOL> class TruncatedPayloadError ( TruncatedRead ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def send_request ( cls , sock , working_dir , command , * arguments , ** environment ) : <EOL> """<STR_LIT>""" <EOL> for argument in arguments : <EOL> cls . write_chunk ( sock , ChunkType . ARGUMENT , argument ) <EOL> for item_tuple in environment . items ( ) : <EOL> cls . write_chunk ( sock , ChunkType . ENVIRONMENT , cls . ENVIRON_SEP . join ( item_tuple ) ) <EOL> cls . write_chunk ( sock , ChunkType . WORKING_DIR , working_dir ) <EOL> cls . write_chunk ( sock , ChunkType . COMMAND , command ) <EOL> @ classmethod <EOL> def parse_request ( cls , sock ) : <EOL> """<STR_LIT>""" <EOL> command = None <EOL> working_dir = None <EOL> arguments = [ ] <EOL> environment = { } <EOL> while not all ( ( working_dir , command ) ) : <EOL> chunk_type , payload = cls . read_chunk ( sock ) <EOL> if chunk_type == ChunkType . ARGUMENT : <EOL> arguments . append ( payload ) <EOL> elif chunk_type == ChunkType . ENVIRONMENT : <EOL> key , val = payload . split ( cls . ENVIRON_SEP , <NUM_LIT:1> ) <EOL> environment [ key ] = val <EOL> elif chunk_type == ChunkType . WORKING_DIR : <EOL> working_dir = payload <EOL> elif chunk_type == ChunkType . COMMAND : <EOL> command = payload <EOL> else : <EOL> raise cls . ProtocolError ( '<STR_LIT>' ) <EOL> return working_dir , command , arguments , environment <EOL> @ classmethod <EOL> def write_chunk ( cls , sock , chunk_type , payload = b'<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> chunk = cls . construct_chunk ( chunk_type , payload ) <EOL> sock . sendall ( chunk ) <EOL> @ classmethod <EOL> def construct_chunk ( cls , chunk_type , payload , encoding = '<STR_LIT:utf-8>' ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( payload , six . text_type ) : <EOL> payload = payload . encode ( encoding ) <EOL> elif not isinstance ( payload , six . binary_type ) : <EOL> raise TypeError ( '<STR_LIT>' . format ( type ( payload ) ) ) <EOL> header = struct . pack ( cls . HEADER_FMT , len ( payload ) , chunk_type ) <EOL> return header + payload <EOL> @ classmethod <EOL> def _read_until ( cls , sock , desired_size ) : <EOL> """<STR_LIT>""" <EOL> buf = b'<STR_LIT>' <EOL> while len ( buf ) < desired_size : <EOL> recv_bytes = sock . recv ( desired_size - len ( buf ) ) <EOL> if not recv_bytes : <EOL> raise cls . TruncatedRead ( '<STR_LIT>' <EOL> . format ( desired_size , len ( buf ) ) ) <EOL> buf += recv_bytes <EOL> return buf <EOL> @ classmethod <EOL> def read_chunk ( cls , sock ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> header = cls . _read_until ( sock , cls . HEADER_BYTES ) <EOL> except cls . TruncatedRead as e : <EOL> raise cls . TruncatedHeaderError ( '<STR_LIT>' . format ( e ) ) <EOL> payload_len , chunk_type = struct . unpack ( cls . HEADER_FMT , header ) <EOL> try : <EOL> payload = cls . _read_until ( sock , payload_len ) <EOL> except cls . TruncatedRead as e : <EOL> raise cls . TruncatedPayloadError ( '<STR_LIT>' . format ( e ) ) <EOL> if chunk_type not in ChunkType . VALID_TYPES : <EOL> raise cls . ProtocolError ( '<STR_LIT>' . format ( chunk_type ) ) <EOL> return chunk_type , payload <EOL> @ classmethod <EOL> def iter_chunks ( cls , sock ) : <EOL> """<STR_LIT>""" <EOL> while <NUM_LIT:1> : <EOL> chunk_type , payload = cls . read_chunk ( sock ) <EOL> yield chunk_type , payload <EOL> if chunk_type == ChunkType . EXIT : <EOL> break <EOL> @ classmethod <EOL> def send_start_reading_input ( cls , sock ) : <EOL> """<STR_LIT>""" <EOL> cls . write_chunk ( sock , ChunkType . START_READING_INPUT ) <EOL> @ classmethod <EOL> def send_stdout ( cls , sock , payload ) : <EOL> """<STR_LIT>""" <EOL> cls . write_chunk ( sock , ChunkType . STDOUT , payload ) <EOL> @ classmethod <EOL> def send_stderr ( cls , sock , payload ) : <EOL> """<STR_LIT>""" <EOL> cls . write_chunk ( sock , ChunkType . STDERR , payload ) <EOL> @ classmethod <EOL> def send_exit ( cls , sock , payload = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> cls . write_chunk ( sock , ChunkType . EXIT , payload ) <EOL> @ classmethod <EOL> def isatty_to_env ( cls , stdin , stdout , stderr ) : <EOL> """<STR_LIT>""" <EOL> fds = ( stdin , stdout , stderr ) <EOL> return { <EOL> cls . TTY_ENV_TMPL . format ( fd_id ) : bytes ( int ( fd . isatty ( ) ) ) for fd_id , fd in enumerate ( fds ) <EOL> } <EOL> @ classmethod <EOL> def isatty_from_env ( cls , env ) : <EOL> """<STR_LIT>""" <EOL> def str_int_bool ( i ) : <EOL> return i . isdigit ( ) and bool ( int ( i ) ) <EOL> return tuple ( str_int_bool ( env . get ( cls . TTY_ENV_TMPL . format ( fd_id ) , '<STR_LIT:0>' ) ) for fd_id in range ( <NUM_LIT:3> ) ) </s>
