<s> import socket <EOL> import random <EOL> import time <EOL> from Queue import Queue , Empty as QueueEmpty <EOL> from rpyc . core . stream import Stream , TunneledSocketStream , ClosedFile <EOL> COOKIE_LENGTH = <NUM_LIT:8> <EOL> class ReconnectingTunnelStream ( Stream ) : <EOL> RETRIES = <NUM_LIT:5> <EOL> def __init__ ( self , remote_machine , destination_port , retries = RETRIES ) : <EOL> self . remote_machine = remote_machine <EOL> self . destination_port = destination_port <EOL> self . retries = retries <EOL> self . cookie = "<STR_LIT>" . join ( chr ( random . randint ( <NUM_LIT:0> , <NUM_LIT:255> ) ) for _ in range ( COOKIE_LENGTH ) ) <EOL> self . stream = None <EOL> def close ( self ) : <EOL> if self . stream is not None and not self . closed : <EOL> self . stream . close ( ) <EOL> self . stream = ClosedFile <EOL> @ property <EOL> def closed ( self ) : <EOL> return self . stream is ClosedFile <EOL> def fileno ( self ) : <EOL> return self . _safeio ( lambda stream : stream . fileno ( ) ) <EOL> def _reconnect ( self ) : <EOL> s = socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) <EOL> s . bind ( ( "<STR_LIT:localhost>" , <NUM_LIT:0> ) ) <EOL> local_port = s . getsockname ( ) [ <NUM_LIT:1> ] <EOL> s . close ( ) <EOL> tun = self . remote_machine . tunnel ( local_port , self . destination_port ) <EOL> stream = TunneledSocketStream . connect ( "<STR_LIT:localhost>" , local_port ) <EOL> stream . write ( self . cookie ) <EOL> stream . tun = tun <EOL> return stream <EOL> def _safeio ( self , callback ) : <EOL> if self . closed : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> for i in range ( self . retries ) : <EOL> try : <EOL> if not self . stream : <EOL> self . stream = self . _reconnect ( ) <EOL> return callback ( self . stream ) <EOL> except ( EOFError , IOError , OSError , socket . error ) : <EOL> if i >= self . retries - <NUM_LIT:1> : <EOL> raise <EOL> if self . stream : <EOL> self . stream . close ( ) <EOL> self . stream = None <EOL> time . sleep ( <NUM_LIT:0.5> ) <EOL> def write ( self , data ) : <EOL> return self . _safeio ( lambda stream : stream . write ( data ) ) <EOL> def read ( self , count ) : <EOL> return self . _safeio ( lambda stream : stream . read ( count ) ) <EOL> class MultiplexingListener ( object ) : <EOL> REACCEPT_TIMEOUT = <NUM_LIT:10> <EOL> RETRIES = <NUM_LIT:5> <EOL> def __init__ ( self , reaccept_timeout = REACCEPT_TIMEOUT , retries = RETRIES ) : <EOL> self . listener = socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) <EOL> self . reaccept_timeout = reaccept_timeout <EOL> self . retries = retries <EOL> self . client_map = { } <EOL> def close ( self ) : <EOL> self . listener . close ( ) <EOL> self . listener = ClosedFile <EOL> def fileno ( self ) : <EOL> return self . listener . getfileno ( ) <EOL> def getsockname ( self ) : <EOL> return self . listener . getsockname ( ) <EOL> def listen ( self , backlog ) : <EOL> self . listener . listen ( backlog ) <EOL> def bind ( self , addrinfo ) : <EOL> self . listener . bind ( addrinfo ) <EOL> def settimeout ( self , timeout ) : <EOL> self . listener . settimeout ( timeout ) <EOL> def setsockopt ( self , level , option , value ) : <EOL> return self . listener . setsockopt ( level , option , value ) <EOL> def shutdown ( self , mode ) : <EOL> self . listener . shutdown ( mode ) <EOL> def accept ( self ) : <EOL> while True : <EOL> sock , addrinfo = self . listener . accept ( ) <EOL> cookie = sock . recv ( COOKIE_LENGTH ) <EOL> if cookie not in self . client_map : <EOL> self . client_map [ cookie ] = Queue ( <NUM_LIT:1> ) <EOL> self . client_map [ cookie ] . put ( sock ) <EOL> resock = ReconnectingSocket ( self , cookie , self . retries ) <EOL> return resock , addrinfo <EOL> else : <EOL> self . client_map [ cookie ] . put ( sock ) <EOL> def reaccept ( self , cookie ) : <EOL> try : <EOL> return self . client_map [ cookie ] . get ( self . reaccept_timeout ) <EOL> except QueueEmpty : <EOL> raise EOFError ( "<STR_LIT>" ) <EOL> class ReconnectingSocket ( object ) : <EOL> def __init__ ( self , listener , cookie , retries ) : <EOL> self . listener = listener <EOL> self . cookie = cookie <EOL> self . sock = None <EOL> self . retries = retries <EOL> self . blocking_mode = None <EOL> def close ( self ) : <EOL> if self . sock : <EOL> self . sock . close ( ) <EOL> self . sock = ClosedFile <EOL> def fileno ( self ) : <EOL> return self . _safeio ( lambda sock : sock . fileno ( ) ) <EOL> def getsockname ( self ) : <EOL> return self . _safeio ( lambda sock : sock . getsockname ( ) ) <EOL> def getpeername ( self ) : <EOL> return self . _safeio ( lambda sock : sock . getpeername ( ) ) <EOL> def shutdown ( self , mode ) : <EOL> if self . sock : <EOL> self . sock . shutdown ( mode ) <EOL> def setblocking ( self , mode ) : <EOL> self . blocking_mode = mode <EOL> if self . sock : <EOL> self . sock . setblocking ( mode ) <EOL> def _safeio ( self , callback ) : <EOL> for i in range ( self . retries ) : <EOL> if self . sock is None : <EOL> self . sock = self . listener . reaccept ( self . cookie ) <EOL> if self . blocking_mode is not None : <EOL> self . sock . setblocking ( self . blocking_mode ) <EOL> try : <EOL> return callback ( self . sock ) <EOL> except ( EOFError , IOError , OSError , socket . error ) : <EOL> if i >= self . retries - <NUM_LIT:1> : <EOL> raise <EOL> if self . sock : <EOL> self . sock . close ( ) <EOL> self . sock = None <EOL> def recv ( self , count ) : <EOL> return self . _safeio ( lambda sock : sock . recv ( count ) ) <EOL> def send ( self , data ) : <EOL> return self . _safeio ( lambda sock : sock . send ( data ) ) </s>
