<s> import struct <EOL> import logging <EOL> import threading <EOL> import socket <EOL> import errno <EOL> from six . moves import range <EOL> try : <EOL> from bitarray import bitarray <EOL> except ImportError : <EOL> class bitarray ( list ) : <EOL> def __init__ ( self , len ) : <EOL> self [ : ] = [ False ] * len <EOL> def setall ( self , val ) : <EOL> for i in range ( len ( self ) ) : <EOL> self [ i ] = val <EOL> from . tds import Error , readall , skipall <EOL> logger = logging . getLogger ( __name__ ) <EOL> class _SmpSession ( object ) : <EOL> def __init__ ( self , mgr , session_id ) : <EOL> self . _session_id = session_id <EOL> self . _seq_num_for_send = <NUM_LIT:0> <EOL> self . _high_water_for_send = <NUM_LIT:4> <EOL> self . _seq_num_for_recv = <NUM_LIT:0> <EOL> self . _high_water_for_recv = <NUM_LIT:4> <EOL> self . _last_high_water_for_recv = <NUM_LIT:4> <EOL> self . _mgr = mgr <EOL> self . _recv_queue = [ ] <EOL> self . _send_queue = [ ] <EOL> self . _state = '<STR_LIT>' <EOL> self . _curr_buf_pos = <NUM_LIT:0> <EOL> self . _curr_buf = b'<STR_LIT>' <EOL> def __repr__ ( self ) : <EOL> fmt = "<STR_LIT>" <EOL> return fmt . format ( self . _session_id , self . _state , self . _recv_queue , self . _send_queue , <EOL> self . _seq_num_for_send ) <EOL> def close ( self ) : <EOL> self . _mgr . _close_smp_session ( self ) <EOL> def send ( self , data , final ) : <EOL> self . _mgr . _send_packet ( self , data ) <EOL> def read ( self , size ) : <EOL> if not self . _curr_buf [ self . _curr_buf_pos : ] : <EOL> self . _curr_buf = self . _mgr . _recv_packet ( self ) <EOL> self . _curr_buf_pos = <NUM_LIT:0> <EOL> if not self . _curr_buf : <EOL> return b'<STR_LIT>' <EOL> res = self . _curr_buf [ self . _curr_buf_pos : self . _curr_buf_pos + size ] <EOL> self . _curr_buf_pos += len ( res ) <EOL> return res <EOL> def is_connected ( self ) : <EOL> return self . _state == '<STR_LIT>' <EOL> class SmpManager ( object ) : <EOL> _smid = <NUM_LIT> <EOL> _smp_header = struct . Struct ( '<STR_LIT>' ) <EOL> _SYN = <NUM_LIT> <EOL> _ACK = <NUM_LIT> <EOL> _FIN = <NUM_LIT> <EOL> _DATA = <NUM_LIT> <EOL> def __init__ ( self , transport ) : <EOL> self . _transport = transport <EOL> self . _sessions = { } <EOL> self . _used_ids_ba = bitarray ( <NUM_LIT:2> ** <NUM_LIT:16> ) <EOL> self . _used_ids_ba . setall ( False ) <EOL> self . _lock = threading . RLock ( ) <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" . format ( self . _sessions ) <EOL> def create_session ( self ) : <EOL> try : <EOL> session_id = self . _used_ids_ba . index ( False ) <EOL> except ValueError : <EOL> raise Error ( "<STR_LIT>" ) <EOL> session = _SmpSession ( self , session_id ) <EOL> with self . _lock : <EOL> self . _sessions [ session_id ] = session <EOL> self . _used_ids_ba [ session_id ] = True <EOL> hdr = self . _smp_header . pack ( <EOL> self . _smid , <EOL> self . _SYN , <EOL> session_id , <EOL> self . _smp_header . size , <EOL> <NUM_LIT:0> , <EOL> session . _high_water_for_recv , <EOL> ) <EOL> self . _transport . send ( hdr , True ) <EOL> session . _state = '<STR_LIT>' <EOL> return session <EOL> def _close_smp_session ( self , session ) : <EOL> if session . _state in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> return <EOL> elif session . _state == '<STR_LIT>' : <EOL> with self . _lock : <EOL> if self . _transport . is_connected ( ) : <EOL> hdr = self . _smp_header . pack ( <EOL> self . _smid , <EOL> self . _FIN , <EOL> session . _session_id , <EOL> self . _smp_header . size , <EOL> session . _seq_num_for_send , <EOL> session . _high_water_for_recv , <EOL> ) <EOL> session . _state = '<STR_LIT>' <EOL> try : <EOL> self . _transport . send ( hdr , True ) <EOL> self . _recv_packet ( session ) <EOL> except ( socket . error , OSError ) as ex : <EOL> if ex . errno in ( errno . ECONNRESET , errno . EPIPE ) : <EOL> session . _state = '<STR_LIT>' <EOL> else : <EOL> raise ex <EOL> else : <EOL> session . _state = '<STR_LIT>' <EOL> def _send_queued_packets ( self , session ) : <EOL> with self . _lock : <EOL> while session . _send_queue and session . _seq_num_for_send < session . _high_water_for_send : <EOL> data = session . _send_queue . pop ( <NUM_LIT:0> ) <EOL> self . _send_packet ( session , data ) <EOL> @ staticmethod <EOL> def _add_one_wrap ( val ) : <EOL> return <NUM_LIT:0> if val == <NUM_LIT:2> ** <NUM_LIT:32> - <NUM_LIT:1> else val + <NUM_LIT:1> <EOL> def _send_packet ( self , session , data ) : <EOL> with self . _lock : <EOL> if session . _seq_num_for_send < session . _high_water_for_send : <EOL> l = self . _smp_header . size + len ( data ) <EOL> seq_num = self . _add_one_wrap ( session . _seq_num_for_send ) <EOL> hdr = self . _smp_header . pack ( <EOL> self . _smid , <EOL> self . _DATA , <EOL> session . _session_id , <EOL> l , <EOL> seq_num , <EOL> session . _high_water_for_recv , <EOL> ) <EOL> session . _last_high_water_for_recv = session . _high_water_for_recv <EOL> self . _transport . send ( hdr + data , True ) <EOL> session . _seq_num_for_send = self . _add_one_wrap ( session . _seq_num_for_send ) <EOL> else : <EOL> session . _send_queue . append ( data ) <EOL> self . _read_smp_message ( ) <EOL> def _recv_packet ( self , session ) : <EOL> with self . _lock : <EOL> if session . _state == '<STR_LIT>' : <EOL> return b'<STR_LIT>' <EOL> while not session . _recv_queue : <EOL> self . _read_smp_message ( ) <EOL> if session . _state in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> return b'<STR_LIT>' <EOL> session . _high_water_for_recv = self . _add_one_wrap ( session . _high_water_for_recv ) <EOL> if session . _high_water_for_recv - session . _last_high_water_for_recv >= <NUM_LIT:2> : <EOL> hdr = self . _smp_header . pack ( <EOL> self . _smid , <EOL> self . _ACK , <EOL> session . _session_id , <EOL> self . _smp_header . size , <EOL> session . _seq_num_for_send , <EOL> session . _high_water_for_recv , <EOL> ) <EOL> self . _transport . send ( hdr , True ) <EOL> session . _last_high_water_for_recv = session . _high_water_for_recv <EOL> return session . _recv_queue . pop ( <NUM_LIT:0> ) <EOL> @ classmethod <EOL> def _type_to_str ( cls , t ) : <EOL> if t == cls . _SYN : <EOL> return '<STR_LIT>' <EOL> elif t == cls . _ACK : <EOL> return '<STR_LIT>' <EOL> elif t == cls . _DATA : <EOL> return '<STR_LIT>' <EOL> elif t == cls . _FIN : <EOL> return '<STR_LIT>' <EOL> def _bad_stm ( self , message ) : <EOL> self . close ( ) <EOL> raise Error ( message ) <EOL> def _read_smp_message ( self ) : <EOL> with self . _lock : <EOL> smid , flags , sid , l , seq_num , wnd = self . _smp_header . unpack ( readall ( self . _transport , self . _smp_header . size ) ) <EOL> if smid != self . _smid : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> try : <EOL> session = self . _sessions [ sid ] <EOL> except KeyError : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> if wnd < session . _high_water_for_send : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> if seq_num > session . _high_water_for_recv : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> session . _last_recv_seq_num = seq_num <EOL> if flags == self . _ACK : <EOL> if session . _state in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> if seq_num != session . _seq_num_for_recv : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> session . _high_water_for_send = wnd <EOL> self . _send_queued_packets ( session ) <EOL> elif flags == self . _DATA : <EOL> if session . _state == '<STR_LIT>' : <EOL> if seq_num != self . _add_one_wrap ( session . _seq_num_for_recv ) : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> session . _seq_num_for_recv = seq_num <EOL> data = readall ( self . _transport , l - self . _smp_header . size ) <EOL> session . _recv_queue . append ( data ) <EOL> if wnd > session . _high_water_for_send : <EOL> session . _high_water_for_send = wnd <EOL> self . _send_queued_packets ( session ) <EOL> elif session . _state == '<STR_LIT>' : <EOL> skipall ( self . _transport , l - self . _smp_header . size ) <EOL> else : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> elif flags == self . _FIN : <EOL> if session . _state == '<STR_LIT>' : <EOL> session . _state = '<STR_LIT>' <EOL> elif session . _state == '<STR_LIT>' : <EOL> session . _state = '<STR_LIT>' <EOL> del self . _sessions [ session . _session_id ] <EOL> self . _used_ids_ba [ session . _session_id ] = False <EOL> elif session . _state == '<STR_LIT>' : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> else : <EOL> self . _bad_stm ( '<STR_LIT>' + session . _state ) <EOL> elif flags == self . _SYN : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> else : <EOL> self . _bad_stm ( '<STR_LIT>' ) <EOL> def close ( self ) : <EOL> self . _transport . close ( ) <EOL> def _transport_closed ( self ) : <EOL> for session in self . _sessions . values ( ) : <EOL> session . _state = '<STR_LIT>' </s>
