# LiMEaide
# Copyright (c) 2011-2018 Daryl Bennett
# Author:
# Daryl Bennett - kd8bny@gmail.com
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at
# your option) any later version.
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
import logging
import selectors
import socket
import struct
import sys
import time
import threading
from termcolor import colored
from queue import Queue
class TCP_CLIENT(threading.Thread):
"""Non-blocking Client to transfer files from server"""
def __init__(self, qresult, ip, port, output):
super(TCP_CLIENT, self).__init__()
self.logger = logging.getLogger(__name__)
self.qresult = qresult
self.ip = ip
self.port = port
self.output = output
self.result = {'success': True, 'terminal': False}
self.byte_count = 0
def __transfer_status__(self, bytes_len):
"""Provide status of the files being transfered.
:param Length of bytes received """
self.byte_count += bytes_len
mbytes_so_far = int(self.byte_count / float(1 << 20))
print(colored(
"Transfer of {0} is at {1:d} MiB".format(
self.output, mbytes_so_far), 'cyan'), end='\r', flush=True)
def __write_out__(self, data):
"""Write data to file by appending
:param data to append"""
try:
with open(self.output, 'ab') as f:
f.write(data)
except Exception as e:
self.logger.error("Unable to save output: {}".format(e))
self.result['success'] = False
self.result['terminal'] = True
sys.exit(colored("Unable to save output", 'red'))
def __handle_client__(self, key, mask, sel):
"""Connect to client using a selector
:return Should client retry to connect
"""
retry = True
sock = key.fileobj
if mask & selectors.EVENT_READ:
try:
recv_data = sock.recv(1024)
if recv_data:
self.__transfer_status__(len(recv_data))
self.__write_out__(recv_data)
else:
self.logger.info("File saved")
print("\n")
retry = False
except socket.error as e:
self.logger.error(e)
self.result['success'] = False
retry = False
return retry
def run(self):
retry = True
self.logger.info("Connecting to Socket")
sel = selectors.DefaultSelector()
conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
conn.setsockopt(
socket.SOL_SOCKET, socket.SO_LINGER, struct.pack('ii', 1, 0))
conn.setblocking(False)
conn.connect_ex((self.ip, self.port))
sel.register(conn, selectors.EVENT_READ, data=None)
while retry:
events = sel.select()
for key, mask in events:
retry = self.__handle_client__(key, mask, sel)
sel.unregister(conn)
if self.result['success']:
conn.shutdown(socket.SHUT_RDWR)
self.qresult.put(self.result)
self.logger.info("Socket Closed")
class CONNECTION_MANAGER(threading.Thread):
"""Handle all client connections and attempts"""
def __init__(self, q, e):
super(CONNECTION_MANAGER, self).__init__()
self.logger = logging.getLogger(__name__)
self.event_kill = e
self.queue = q
self.qstatus = Queue()
self.retry_count = 0
def __start_client__(self, conn):
"""Start a new client thread"""
client = TCP_CLIENT(self.qstatus, conn[0], conn[1], conn[2])
client.start()
client.join()
def run(self):
while not self.event_kill.is_set():
time.sleep(3)
if self.queue.empty():
continue
else:
connection = self.queue.get()
self.__start_client__(connection)
status = self.qstatus.get()
if not status['success']:
if status['terminal']:
self.logger.warning(
"Failed connection, closing down")
self.event_kill.set()
elif self.retry_count > 4:
self.logger.warning(
"Retry count exceeded, closing down")
self.event_kill.set()
else:
self.retry_count += 1
self.logger.warning(
"connection failed, adding back to queue")
self.queue.put(connection)
self.logger.info("Connection manager is finished")
