<s> from scapy . all import TCP , IP , Padding <EOL> import random <EOL> import time <EOL> class BadPacketError ( Exception ) : <EOL> pass <EOL> def get_payload ( packet ) : <EOL> while not isinstance ( packet , TCP ) : <EOL> packet = packet . payload <EOL> return packet . payload <EOL> class TCPSocket ( object ) : <EOL> def __init__ ( self , listener , verbose = <NUM_LIT:0> ) : <EOL> self . state = "<STR_LIT>" <EOL> self . verbose = verbose <EOL> self . src_ip = listener . ip_address <EOL> self . recv_buffer = "<STR_LIT>" <EOL> self . listener = listener <EOL> self . seq = self . _generate_seq ( ) <EOL> self . last_ack_sent = <NUM_LIT:0> <EOL> @ staticmethod <EOL> def _has_load ( packet ) : <EOL> payload = get_payload ( packet ) <EOL> if isinstance ( payload , Padding ) : <EOL> return False <EOL> return bool ( payload ) <EOL> def _set_dest ( self , host , port ) : <EOL> self . dest_port = port <EOL> self . dest_ip = host <EOL> self . ip_header = IP ( dst = self . dest_ip , src = self . src_ip ) <EOL> def connect ( self , host , port ) : <EOL> self . _set_dest ( host , port ) <EOL> self . src_port = self . listener . get_port ( ) <EOL> self . listener . open ( self . src_ip , self . src_port , self ) <EOL> self . _send_syn ( ) <EOL> def bind ( self , host , port ) : <EOL> self . src_ip = host <EOL> self . src_port = port <EOL> self . listener . open ( self . src_ip , self . src_port , self ) <EOL> self . state = "<STR_LIT>" <EOL> @ staticmethod <EOL> def _generate_seq ( ) : <EOL> return random . randint ( <NUM_LIT:0> , <NUM_LIT> ) <EOL> def _send ( self , flags = "<STR_LIT>" , load = None ) : <EOL> """<STR_LIT>""" <EOL> packet = TCP ( dport = self . dest_port , <EOL> sport = self . src_port , <EOL> seq = self . seq , <EOL> ack = self . last_ack_sent , <EOL> flags = flags ) <EOL> full_packet = self . ip_header / packet <EOL> if load : <EOL> full_packet = full_packet / load <EOL> self . listener . send ( full_packet ) <EOL> if load is not None : <EOL> self . seq += len ( load ) <EOL> def _send_syn ( self ) : <EOL> self . state = "<STR_LIT>" <EOL> self . _send ( flags = "<STR_LIT:S>" ) <EOL> def _send_ack ( self , flags = "<STR_LIT>" , load = None ) : <EOL> """<STR_LIT>""" <EOL> self . _send ( flags = flags + "<STR_LIT:A>" , load = load ) <EOL> def close ( self ) : <EOL> if self . state == "<STR_LIT>" : <EOL> return <EOL> self . state = "<STR_LIT>" <EOL> self . _send_ack ( flags = "<STR_LIT:F>" ) <EOL> @ staticmethod <EOL> def next_seq ( packet ) : <EOL> tcp_flags = packet . sprintf ( "<STR_LIT>" ) <EOL> if TCPSocket . _has_load ( packet ) : <EOL> return packet . seq + len ( packet . load ) <EOL> elif '<STR_LIT:S>' in tcp_flags or '<STR_LIT:F>' in tcp_flags : <EOL> return packet . seq + <NUM_LIT:1> <EOL> else : <EOL> return packet . seq <EOL> def _close ( self ) : <EOL> self . state = "<STR_LIT>" <EOL> self . listener . close ( self . src_ip , self . src_port ) <EOL> def handle ( self , packet ) : <EOL> if self . last_ack_sent and self . last_ack_sent != packet . seq : <EOL> return <EOL> self . last_ack_sent = max ( self . next_seq ( packet ) , self . last_ack_sent ) <EOL> recv_flags = packet . sprintf ( "<STR_LIT>" ) <EOL> if self . _has_load ( packet ) : <EOL> self . recv_buffer += packet . load <EOL> self . _send_ack ( ) <EOL> elif "<STR_LIT:R>" in recv_flags : <EOL> self . _close ( ) <EOL> elif "<STR_LIT:S>" in recv_flags : <EOL> if self . state == "<STR_LIT>" : <EOL> self . state = "<STR_LIT>" <EOL> self . _set_dest ( packet . payload . src , packet . sport ) <EOL> self . _send_ack ( flags = "<STR_LIT:S>" ) <EOL> elif self . state == "<STR_LIT>" : <EOL> self . seq += <NUM_LIT:1> <EOL> self . state = "<STR_LIT>" <EOL> self . _send_ack ( ) <EOL> elif "<STR_LIT:F>" in recv_flags : <EOL> if self . state == "<STR_LIT>" : <EOL> self . seq += <NUM_LIT:1> <EOL> self . state = "<STR_LIT>" <EOL> self . _send_ack ( flags = "<STR_LIT:F>" ) <EOL> elif self . state == "<STR_LIT>" : <EOL> self . seq += <NUM_LIT:1> <EOL> self . _send_ack ( ) <EOL> self . _close ( ) <EOL> elif "<STR_LIT:A>" in recv_flags : <EOL> if self . state == "<STR_LIT>" : <EOL> self . state = "<STR_LIT>" <EOL> elif self . state == "<STR_LIT>" : <EOL> self . _close ( ) <EOL> else : <EOL> raise BadPacketError ( "<STR_LIT>" ) <EOL> def send ( self , payload ) : <EOL> while self . state != "<STR_LIT>" : <EOL> time . sleep ( <NUM_LIT> ) <EOL> self . _send_ack ( load = payload , flags = "<STR_LIT:P>" ) <EOL> def recv ( self , size , timeout = None ) : <EOL> start_time = time . time ( ) <EOL> while len ( self . recv_buffer ) < size : <EOL> time . sleep ( <NUM_LIT> ) <EOL> if self . state in [ "<STR_LIT>" , "<STR_LIT>" ] : <EOL> break <EOL> if timeout < ( time . time ( ) - start_time ) : <EOL> break <EOL> recv = self . recv_buffer [ : size ] <EOL> self . recv_buffer = self . recv_buffer [ size : ] <EOL> return recv </s>
