import queue import socket import selectors import time from datetime import datetime from threading import Thread from time import sleep import ClientTempMail import Crypto # Included the unix timestamp function from the DBConn file so we dont have to use the entire class just for this def get_unix_timestamp(seconds_offset=0) -> int: """ Function to get the current unix timestamp to sim the NOW() sql method, seconds offset will pull a past of future timestamp """ return int(time.time()) + seconds_offset class ClientThread(Thread): """Class for dealing with the threaded actions of the client app""" def __init__(self, current_socket: socket, address: tuple): Thread.__init__(self) self.__address = address self.__socket: socket.socket = current_socket self.__selector = selectors.DefaultSelector() self.__buffer_in = queue.Queue() self.__buffer_out = queue.Queue() # flag for the current network status, -1 for awaiting request, 0 for awaiting response, 1 for ready self.current_status = -1 self.state = "INIT" self.__init_timeout = get_unix_timestamp() self.last_response_content = [] # Login and User variables self.__session_ref = None self.login_username = None self.__user_full_name = None self.__user_email = None self.__user_role = None self.initialised = False # Mail class to hold the mail data structure that mimics the server version self.client_mail_temp: ClientTempMail.ClientTempMail or None = None self.__queued_messages = [] # Handshake and Encryption Variables self.__pending_handshake = False self.__hand_shook = False self.__crypto_manager_receive: Crypto.Receiver or None = None self.__crypto_manager_transmit: Crypto.Transmitter or None = None # Active and Heartbeat Vars self.active = True self.__heartbeat_timer = None self.__heartbeat_tick = -120 self.__selector.register(self.__socket, selectors.EVENT_READ | selectors.EVENT_WRITE, data=None) def is_connected(self) -> bool: """Method for checking if the connection is active, and sending the noop command for the heartbeat""" if not self.__socket: return False # If the connection is initialised we can send commands if self.initialised: try: # Send heartbeat NOOP command, this is a background command as to not ruin the flow of the program # since actual SMTP servers work in a specific order self.queue_command("NOOP") except Exception: return False return True def get_users_role(self) -> str: """Method for getting the current users server role (USER or ADMIN)""" return self.__user_role def run(self): """Main threads looped run method that listens and writes our network data""" pending_handshake_counter = 0 try: while self.active: waiting_commands = self.__selector.select(timeout=1) for key, mask in waiting_commands: try: if mask & selectors.EVENT_READ: self.__read() elif mask & selectors.EVENT_WRITE and not self.__buffer_out.empty(): self.__write() # Process any queued messages when an actual message is not being processed elif mask & selectors.EVENT_WRITE and len(self.__queued_messages) > 0: self.__write_queue() except Exception as ex: print(f"Exception hit: {repr(ex)}") self.__close() # Init time out makes sure that a "valid" connection it makes will time out if the server # does not respond the way we want it to if self.state == "INIT": if self.__init_timeout < get_unix_timestamp(-40): print("No INIT processed, the connection is most likely at fault.") self.active = False # Pending handshake so we can send the AUTH response then set the handshook flag after it has been sent if self.__pending_handshake: if pending_handshake_counter == 1: self.__pending_handshake = False self.__hand_shook = True pending_handshake_counter = 0 else: pending_handshake_counter += 1 # If logged in we can do our Heartbeat NOOP request to faux keep our connection alive if self.__session_ref: if self.__heartbeat_timer < get_unix_timestamp(self.__heartbeat_tick): if not self.is_connected(): print("Server did not respond") self.__close() if not self.__selector.get_map(): break finally: self.__selector.close() def __read(self): """Private method for reading the network data coming into the client""" try: data_in = self.__socket.recv(1024) except BlockingIOError: pass else: if data_in: if self.__hand_shook: # We decode the bytes of the message we receive, this should be a hex string data_in_decoded_string = data_in.decode() # If needed we split the response into chunks so that we can decrypt the multiple blocks that # were sent chunk_size = self.__crypto_manager_receive.get_expected_block_size() data_in_chunks = [self.__crypto_manager_receive.decrypt_string( data_in_decoded_string[i:i + chunk_size] ) for i in range(0, len(data_in_decoded_string), chunk_size)] # Join all of the decrypted chunks back together full_data_in = "".join(data_in_chunks) self.__buffer_in.put(full_data_in) else: # Data is treated normally when we have not shaken hands self.__buffer_in.put(data_in.decode()) else: raise RuntimeError("Connection Closed") self.process_response() def __write_queue(self): """Private method for writing a queue of data with a simulated 'latency' to stop the messages being joined""" # Some SMTP response require the server to send multiple responses per request while len(self.__queued_messages) > 0: # Sleep to simulate some kind of latency to stop the queue to hopefully stop # the response from being added to the previous. sleep(1) try: self.process_silent_command(self.__queued_messages.pop(0)) response = self.__buffer_out.get_nowait() except Exception: response = None if response: # This would sometimes interrupt the users input because a queued write would happen during it # print(f"Sending Queued response to Client: {self.__address}") try: self.__do_write(response) except BlockingIOError as ex: print(f"IO Error: {repr(ex)}") pass def __write(self): """Private method for doing a single write to the socket""" try: response = self.__buffer_out.get_nowait() except Exception: response = None if response: print(f"Sending data: {repr(response)} to {self.__address}") try: self.__do_write(response) except BlockingIOError as ex: print(f"Could not send data: {repr(ex)}") pass def __do_write(self, response: bytes) -> None: """Shared Private method to do the actual write to the socket""" # Check if we have shook hands if self.__hand_shook: # Get the max chunk size for the key we generated in case we need to split the request down chunk_size = self.__crypto_manager_transmit.get_usable_byte_length() response_chunks = [self.__crypto_manager_transmit.encrypt_string(response[i:i + chunk_size]) for i in range(0, len(response), chunk_size)] full_response = b"".join(response_chunks) self.__socket.send(full_response) else: self.__socket.send(response) # Set the heartbeat timer since we just sent a command and wont need to send a noop again self.__heartbeat_timer = get_unix_timestamp() def process_command(self, command: str) -> None: """Method to Process the command and add it into the buffer, we also set the client to read mode""" self.current_status = 0 self.__buffer_out.put(command.encode()) def process_silent_command(self, command: str) -> None: """Method to silently add a command to the buffer, this does not cause the client to enter read mode""" self.__buffer_out.put(command.encode()) def queue_command(self, command: str) -> None: """Method to add a command to the queued message array for queued writing""" self.__queued_messages.append(command) def process_response(self): """Method to process the response from the server and to set the client to the required state""" current_data = self.__buffer_in.get() # Split the response string by the first space, this should give us a response code and response data response_parts = current_data.split(" ", 1) if len(response_parts) >= 2: response_code = int(response_parts[0]) # If the response is not from the heartbeat print it if response_code != 295: print(f"Response: {repr(current_data)}") if len(response_parts) > 1: response_data = response_parts[1] else: response_data = None # Check the current state and respond accordingly if self.state == "INIT": if response_code == 220: self.state = "HELO" # Initialised, now send the HELO command self.process_command("HELO " + self.__address[0]) self.current_status = 0 self.initialised = True else: print("No response from the server...") self.__close() elif self.state == "HELO": if response_code == 250: # HELO responded, send the AUTH command self.state = "AUTH" self.current_status = 0 elif self.state == "AUTH": if not self.__hand_shook: if response_code == 570: # set up the RSA encryption variables for encrypting the connection self.__crypto_manager_transmit = Crypto.Transmitter(response_data) self.__crypto_manager_receive = Crypto.Receiver() self.__crypto_manager_receive.generate_keys() # Send our public key back to the server self.process_command("AUTH " + self.__crypto_manager_receive.get_public_key_pair()) # Set the client to a pending handshake self.__pending_handshake = True else: if response_code == 250: # Auth key was sent successfully, start the confirmation handshake self.process_command("HSHK " + response_data) self.state = "HSHK" elif self.state == "HSHK": if response_code == 250: # The validation handshake was current, start the login state and set the write self.state = "LOGI" self.current_status = -1 elif self.state == "LOGI": if response_code == 250: # Login was successful, wait for the secondary role response so we can tell # if the user is and ADMIN or USER self.current_status = 0 elif response_code == 265: self.__session_ref = response_data self.current_status = 0 # Send the VRFY command to the get logged in users name and email address self.state = "VRFY" self.process_command("VRFY " + self.login_username) elif response_code == 421: pass else: # Bad login, try again self.last_response_content.clear() self.last_response_content.append(response_data) self.current_status = -1 elif self.state == "VRFY": if response_code == 261: self.__user_role = response_data elif response_code == 250: # Process the email address that gets sent back with the VRFY response if '@' in response_data: if '<' in response_data: start_bracket_index = response_data.find('<') end_bracket_index = response_data.rfind('>') self.__user_full_name = response_data[0:start_bracket_index-1].strip() self.__user_email = response_data[start_bracket_index+1:end_bracket_index].strip() else: # self.__user_full_name = "N/A" self.__user_email = response_data.strip() # Send the user to the Home menu (HMEN) as they are not logged in self.state = "HMEN" self.current_status = -1 else: # VRFY failed and we send the user back to the login self.reset_session() self.state = "LOGI" self.current_status = -1 elif response_code == 500: # VRFY failed and we send the user back to the login self.reset_session() self.state = "LOGI" self.current_status = -1 elif self.state == "MAIL": # Start an email input, we start by asking for the receipt addresses (RCPT) if response_code == 250: self.state = "RCPT" self.current_status = -1 else: self.state = "HMEN" self.current_status = -1 elif self.state == "RCPT": self.last_response_content.clear() if response_code == 250: # The address was local and added self.last_response_content.append("Address Added.") elif response_code == 251: # The address was valid but not local and added self.last_response_content.append("Address Added, but not local.") elif response_code == 503: # The server said we are out of sequence, reset the transaction and start over self.last_response_content.append("There was an issue with the order of commands sent, RSET sent.") self.state = "RSET" self.process_command("RSET") else: # The address was rejected, we remove it from the clients end last_address = self.client_mail_temp.to_address_list.pop(-1) self.last_response_content.append(f"Address {last_address} not added for reason: " + response_data) if self.state != "RSET": self.state = "RCPT" self.current_status = -1 elif self.state == "SUBJ": # The receipts were fine and we move onto the subject, subject does not have an official "state" # we process it fully from the client side and send it along with all of the starting DATA requests if response_code == 354: string_date = datetime.now() date_formatted = string_date.strftime("%d %b %y %H:%M:%S") # send the email headers that get used when "forwarding" to another server, for completeness # DATA starts with 4 lines, Date, From, Subject and To as a comma seperated list self.queue_command("Date: " + date_formatted) self.queue_command("From: " + self.__user_full_name + " <" + self.__user_email + ">") self.queue_command("Subject: " + self.client_mail_temp.subject) self.queue_command("To: " + ", ".join(self.client_mail_temp.to_address_list)) # Set the state to DATA and start accepting lines from the client until the >SEND command is used self.state = "DATA" self.current_status = -1 else: # The server said we are out of sequence, reset the transaction and start over self.last_response_content.append("There was an issue with the order of commands sent, RSET sent.") self.state = "RSET" self.process_command("RSET") elif self.state == "DATA": if response_code == 250: # The data was sent and the server responded with an OK response, return to the home menu self.last_response_content.clear() self.last_response_content.append("Email sent.") self.client_mail_temp = None self.state = "HMEN" self.current_status = -1 elif self.state == "RSET": if response_code == 250: # The server responded OK to our request to reset the current transaction self.client_mail_temp = None self.state = "HMEN" self.current_status = -1 elif self.state == "VIEW": if response_code == 211: # Server responded to the VIEW command self.current_status = 0 self.last_response_content.clear() elif response_code == 214: # Process the VIEW response until we get the end of DATA string if response_data == "\n.\n": self.current_status = -1 else: self.last_response_content.append(response_data) elif response_code == 504: # There was an error getting a mail, usually if they try to enter # a number that doesnt belong to them self.last_response_content.clear() self.last_response_content.append("There was an issue loading this mail") self.current_status = -1 else: # There was an unspecified error so we send the user back to the home menu self.last_response_content.clear() self.state = "HMEN" self.current_status = -1 elif self.state == "ADMN": if response_code == 250: # Server has left ADMIN mode if response_data == "OK QUIT": self.last_response_content.clear() self.state = "HMEN" self.current_status = -1 # Server has processed the command given elif response_data == "OK DONE": self.last_response_content.clear() self.last_response_content.append("Action Completed.") self.current_status = -1 else: self.last_response_content.clear() self.current_status = -1 # Server sent a multipart ADMIN response elif response_code == 211: self.current_status = 0 self.last_response_content.clear() elif response_code == 214: if response_data == "\n.\n": self.current_status = -1 else: self.last_response_content.append(response_data) # Server sent an ADMIN error response elif response_code == 500: self.last_response_content.clear() self.last_response_content.append("Action Failed.") self.current_status = -1 # Permission was denied by the server else: self.last_response_content.clear() self.last_response_content.append("Permission Denied.") self.state = "HMEN" self.current_status = -1 elif self.state == "LOUT": # Server logged the client out if response_code == 250: # Logged out self.reset_session() self.state = "LOGI" else: # Couldn't log out? self.state = "HMEN" self.current_status = -1 elif self.state == "HELP": # Server responded to the HELP command and is sending back the multipart response if response_code == 211: self.current_status = 0 self.last_response_content.clear() elif response_code == 214: if response_data == "\n.\n": self.current_status = -1 else: self.last_response_content.append(response_data) # Server responded with a command not implemented error becuase the help section does not exist elif response_code == 504: self.last_response_content.clear() self.last_response_content.append("There was no help on file for your request.") self.current_status = -1 # Unspecified error from the server, just sending the user to the home menu else: self.last_response_content.clear() self.state = "HMEN" self.current_status = -1 # Status 421 and 221 is only used when the connection has been terminated and we just exit at that point if response_code == 421: print("The connection was closed by the server with the following message:") for line in self.last_response_content: print(line) else: print(f"Reason: {response_data}") self.__close() self.current_status = -2 if response_code == 221: self.state = "QUIT" self.__close() self.current_status = -2 else: # The response did not have a valid response code print("Response too short.") def __close(self): """Private method to close the connection to the server and end the thread""" print("Attempting to end the Connection...") try: self.__selector.unregister(self.__socket) self.__socket.close() print(f"Connection ended with {repr(self.__address)}") except OSError as ex: print(f"Could not close the connection: {repr(ex)}") pass finally: self.__socket = None self.active = False self.current_status = -1 def get_session_ref(self) -> str: """Method to return the current session reference from the server""" return self.__session_ref def get_user_name(self) -> str: """Method to return the current users full name""" return self.__user_full_name def get_user_email(self) -> str: """Method to return the current users email address""" return self.__user_email def reset_session(self) -> None: """Method to restart a clients session without closing the connection, used for the logout command""" self.__session_ref = None self.__user_email = None self.__user_full_name = None self.__user_role = None self.login_username = None