def construct_data_pkg(data: bytes, local_host: str, local_port: int,
dest_host: str, dest_port: int, seq_number: int,
ack_number: int, window_size: int) -> 'TCPPackage':
flags = ut.PacketFlags(False, False, False, False, False, False, False,
False)
info = ut.PacketInfo(local_host, dest_host, local_port, dest_port,
seq_number, ack_number, 0, flags, 0, window_size)
return TCPPackage(data, info, True)
def send(self, data:bytes)->int:
init_len = len(data)
if not len(data):
flags = ut.PacketFlags(False, False, False, False, False, False, False, False)
data = self._build_pkg(b"", flags)
self.sender.schedule(data)
while not self.sender.finished:
time.sleep(0.05)
return init_len
def _get_packet(self) -> ut.PacketInfo:
flags = ut.PacketFlags(self.ack_flag, self.rst_flag, self.syn_flag,
self.fin_flag, self.ece_flag, self.cwr_flag,
self.psh_flag, self.urg_flag)
packet = ut.PacketInfo(self.source_host, self.dest_host,
self.source_port, self.dest_port,
self.seq_number, self.ack_number, self.checksum,
flags, self.urg_ptr, self.window_size)
return packet
def _handle_syn(self, pkg:TCPPackage):
if pkg.ack_flag: # SYNACK
self._send_ack()
if self.state == self.SYN_SENT:
self.state = Conn.CONNECTED
elif self.state == Conn.CLOSED: # SYN
synack_pkg = self._build_pkg(b"",ut.PacketFlags(True, False, True, False, False, False, False, False))
self.sender.schedule(synack_pkg)
self.state = Conn.SYNACK_SENT
self.__synack_seq_number = synack_pkg.seq_number
def close(self):
if self.state == Conn.CLOSED:
raise ConnException(exc.CON_ALREADY_CLOSED_MSG)
if self.is_server:
self._close_connection(0)
return
flags = ut.PacketFlags(False, False, False, True, False, False, False, False)
fin_pkg = self._build_pkg(b"", flags)
self.state = Conn.FIN_SENT
self.sender.schedule(fin_pkg)
def _flags_info(self):
flags = ut.PacketFlags(self.ack_flag, self.rst_flag, self.syn_flag,
self.fin_flag, self.ece_flag, self.cwr_flag,
self.psh_flag, self.urg_flag)
str_flags = ""
for value, name in zip(flags, flags._fields):
if value:
str_flags += name.upper() + " "
else:
if str_flags:
str_flags = str_flags[:len(str_flags) -
1] # Remove last space if any
str_flags = "[" + str_flags + "]"
return str_flags
def _send_no_conn_reply(self, package:TCPPackage):
"""
Send a rst package to client
"""
with ut.get_raw_socket() as s:
# Change package information
# Set rst flag to notify that no conn exist with the package address
flags = ut.PacketFlags(False, True, False, False, False, False, False, False)
info = ut.PacketInfo(package.dest_host, package.source_host, package.dest_port,
package.source_port, package.ack_number, package.seq_number,0,
flags, 0, package.window_size)
package = TCPPackage(b"", info, True)
s.sendto(package.to_bytes(),(package.dest_host, package.dest_port))
log.info(f"RST Package sent {package._endpoint_info()}")
def _handle_fin(self, pkg:TCPPackage):
if pkg.ack_flag: # FINACK Pkg
if self.state == Conn.FIN_SENT:
self.state = Conn.FIN_WAIT
# Start waiting
if not self.__fin_executor:
self.__fin_executor = ThreadPoolExecutor(1)
self.__fin_executor.submit(self._close_connection, self.FIN_WAITING)
if self.state in [Conn.FIN_WAIT, Conn.FINACK_SENT]: # FINACK_SENT because both endpoints closed at the same time
self._send_ack()
elif self.state in [Conn.CONNECTED, Conn.FIN_SENT]: # FIN pkg
if self.state == Conn.FIN_SENT: # Both endpoints closed at the same time
pkg = pkg.copy()
pkg.ack_number += 1 # ACK the FIN PKG
self.sender.rcv_pkg(pkg)
while not self.sender.finished:
time.sleep(0.05)
self.state = Conn.FINACK_SENT
flags = ut.PacketFlags(True, False, False, True, False, False, False, False)
finack_pkg = self._build_pkg(b"", flags)
self.__finack_seq_number = finack_pkg.seq_number
self.sender.schedule(finack_pkg)
def _send_ack(self):
flags = ut.PacketFlags(True, False, False, False, False, False, False, False)
pkg = self._build_pkg(b"", flags)
self.sender.schedule(pkg, False)
return pkg
def init_connection(self):
syn_pkg = self._build_pkg(b"",ut.PacketFlags(False, False, True, False, False, False, False, False))
self.state = Conn.SYN_SENT
self.sender.schedule(syn_pkg)