diff --git a/ai/ai_simulations.py b/ai/ai_simulations.py index da01ce6..3888128 100644 --- a/ai/ai_simulations.py +++ b/ai/ai_simulations.py @@ -11,13 +11,19 @@ def __init__(self): ] def simulate_attack(self): + if not self.scenarios: + print("Error: No scenarios available for simulation.") + return + try: if not self.scenarios: raise IndexError("No scenarios available.") scenario = random.choice(self.scenarios) print(f"[SIMULATION] Executing simulated attack: {scenario}") + except IndexError as e: print(f"Error during simulation: {e}") + except Exception as e: print(f"Error during simulation: {e}") diff --git a/app_security/app_vulnerability_scanner.py b/app_security/app_vulnerability_scanner.py index 62e15f4..7008e66 100644 --- a/app_security/app_vulnerability_scanner.py +++ b/app_security/app_vulnerability_scanner.py @@ -2,6 +2,7 @@ from database.models import DocumentAnalysis from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker +import time DATABASE_URL = "sqlite:///document_analysis.db" engine = create_engine(DATABASE_URL) @@ -9,47 +10,55 @@ def scan_application(app_url): print(f"Scanning application for vulnerabilities: {app_url}") - try: - session = SessionLocal() + retries = 3 + for attempt in range(retries): try: - response = requests.get(app_url) - response.raise_for_status() - vulnerabilities = {"vulnerabilities_found": 2, "critical_issues": ["SQL Injection", "XSS"]} - - # Save scan results to the database - scan_result = DocumentAnalysis( - source=app_url, - title="Vulnerability Scan", - links=str(vulnerabilities["critical_issues"]), - error=None - ) - session.add(scan_result) - session.commit() - return vulnerabilities - except requests.exceptions.HTTPError as http_err: - print(f"HTTP error occurred: {http_err}") - scan_result = DocumentAnalysis( - source=app_url, - title="Vulnerability Scan", - links=None, - error=str(http_err) - ) - session.add(scan_result) - session.commit() - except Exception as err: - print(f"Other error occurred: {err}") - scan_result = DocumentAnalysis( - source=app_url, - title="Vulnerability Scan", - links=None, - error=str(err) - ) - session.add(scan_result) - session.commit() - finally: - session.close() - except Exception as db_err: - print(f"Database connection error: {db_err}") + session = SessionLocal() + try: + response = requests.get(app_url) + response.raise_for_status() + vulnerabilities = {"vulnerabilities_found": 2, "critical_issues": ["SQL Injection", "XSS"]} + + # Save scan results to the database + scan_result = DocumentAnalysis( + source=app_url, + title="Vulnerability Scan", + links=str(vulnerabilities["critical_issues"]), + error=None + ) + session.add(scan_result) + session.commit() + return vulnerabilities + except requests.exceptions.HTTPError as http_err: + print(f"HTTP error occurred: {http_err}") + scan_result = DocumentAnalysis( + source=app_url, + title="Vulnerability Scan", + links=None, + error=str(http_err) + ) + session.add(scan_result) + session.commit() + except Exception as err: + print(f"Other error occurred: {err}") + scan_result = DocumentAnalysis( + source=app_url, + title="Vulnerability Scan", + links=None, + error=str(err) + ) + session.add(scan_result) + session.commit() + finally: + session.close() + except Exception as db_err: + print(f"Database connection error: {db_err}") + if attempt < retries - 1: + print("Retrying database connection...") + time.sleep(2) + else: + print("Failed to connect to the database after multiple attempts.") + return {"vulnerabilities_found": 0, "critical_issues": []} return {"vulnerabilities_found": 0, "critical_issues": []} def verify_database_connection(): diff --git a/core/email_server/EmailServer.py b/core/email_server/EmailServer.py index 622f5e4..fed656c 100644 --- a/core/email_server/EmailServer.py +++ b/core/email_server/EmailServer.py @@ -103,8 +103,13 @@ def Save_Email_To_Recipient(client_socket, data, msg, requests, subject, sender, filename = filename.split("/")[-1] # Save the image file - with open(os.path.join(recipient_directory, filename), "wb") as f: - f.write(part.get_payload(decode=True)) + try: + with open(os.path.join(recipient_directory, filename), "wb") as f: + f.write(part.get_payload(decode=True)) + except Exception as e: + logging.error(f"Error saving email attachment: {e}") + client_socket.sendall("Error saving email attachment".encode('utf-8')) + return print(f"From: {sender}") print(f"To: {recipient}") @@ -121,18 +126,28 @@ def Save_Email_To_Recipient(client_socket, data, msg, requests, subject, sender, if not os.path.isfile(f"{recipient_directory}/{recipient}_received_emails.csv") or ( os.stat(f"{recipient_directory}/{recipient}_received_emails.csv").st_size == 0): # If the file doesn't exist, then create the file and save the email to the file df = pd.DataFrame(email_data, columns=MyColumns) - df.to_csv(f"{recipient_directory}/{recipient}_received_emails.csv", mode='w', header=True, index=False) # Save the email to the recipient's inbox - df.to_csv(f"{recipient_directory}/{recipient}_received_emailsHistory.csv", mode='w', header=True, index=False) # Save the email to the recipient's inbox history + try: + df.to_csv(f"{recipient_directory}/{recipient}_received_emails.csv", mode='w', header=True, index=False) # Save the email to the recipient's inbox + df.to_csv(f"{recipient_directory}/{recipient}_received_emailsHistory.csv", mode='w', header=True, index=False) # Save the email to the recipient's inbox history + except Exception as e: + logging.error(f"Error saving email to CSV: {e}") + client_socket.sendall("Error saving email to CSV".encode('utf-8')) + return else: # If the file already exists, then append the email to the file - df = pd.read_csv(f"{recipient_directory}/{recipient}_received_emails.csv") # Read the csv file of the recipient - new_row_df = pd.DataFrame(email_data, columns=df.columns) - df = pd.concat([df, new_row_df], ignore_index=True) - df.to_csv(f"{recipient_directory}/{recipient}_received_emails.csv", mode='w', header=True, index=False) - df = pd.read_csv(f"{recipient_directory}/{recipient}_received_emailsHistory.csv") - df = pd.concat([df, new_row_df], ignore_index=True) - df.to_csv(f"{recipient_directory}/{recipient}_received_emailsHistory.csv", mode='w', header=True, index=False) + try: + df = pd.read_csv(f"{recipient_directory}/{recipient}_received_emails.csv") # Read the csv file of the recipient + new_row_df = pd.DataFrame(email_data, columns=df.columns) + df = pd.concat([df, new_row_df], ignore_index=True) + df.to_csv(f"{recipient_directory}/{recipient}_received_emails.csv", mode='w', header=True, index=False) + df = pd.read_csv(f"{recipient_directory}/{recipient}_received_emailsHistory.csv") + df = pd.concat([df, new_row_df], ignore_index=True) + df.to_csv(f"{recipient_directory}/{recipient}_received_emailsHistory.csv", mode='w', header=True, index=False) + except Exception as e: + logging.error(f"Error appending email to CSV: {e}") + client_socket.sendall("Error appending email to CSV".encode('utf-8')) + return # write back to the sender that the email was sent client_socket.sendall("Email Sent".encode('utf-8')) diff --git a/core/end_user/EndUserClient.py b/core/end_user/EndUserClient.py index fe1ae25..06ab641 100644 --- a/core/end_user/EndUserClient.py +++ b/core/end_user/EndUserClient.py @@ -55,78 +55,86 @@ def receive_complete_data(client_socket): # this function is used to receive the def parse_email_data(data): # this function gets the data from the inbox and parse it to the email data - msg = email.message_from_bytes(data) + try: + msg = email.message_from_bytes(data) - Command, subject, sender, recipient = msg['Command'], msg["Subject"], msg["From"], msg["To"] - recipient_directory = f"{saveMail_directory}/{recipient}" - os.makedirs(recipient_directory, exist_ok=True) + Command, subject, sender, recipient = msg['Command'], msg["Subject"], msg["From"], msg["To"] + recipient_directory = f"{saveMail_directory}/{recipient}" + os.makedirs(recipient_directory, exist_ok=True) - if msg.is_multipart(): - for part in msg.get_payload(): - if part.get_content_type() == "text/plain": - body = part.get_payload() - else: - print(msg.get_payload()) - for part in msg.walk(): - if part.get_content_maintype() == "multipart": - continue - if part.get("Content-Disposition") is None: - continue - - filename = part.get_filename() - #filename = filename.split("\\")[-1] - filename = filename.split("/")[-1] - - # Save the image file - with open(os.path.join(recipient_directory, filename), "wb") as f: - f.write(part.get_payload(decode=True)) - print(f'\n Opened and parsed new email from {sender} to {recipient} with subject {subject}') - print(f'Email body: {body}') - print(f'Email attachment: {filename}') - - filepath = str(f"{recipient_directory}/{filename}") - try: #We faced some network errors resulting in images being sent partially black. To address this issue, we implemented a try-except block to handle such occurrences. Now, if an image fails to send correctly, a default image is sent for that experiment. - with open(filepath) as f: # TEST IF THE FILE IS A VALID IMAGE - img = MIMEImage(f.read()) - except: # network error - if default_image=='': - print('Network Error: No default image is set') - return + if msg.is_multipart(): + for part in msg.get_payload(): + if part.get_content_type() == "text/plain": + body = part.get_payload() else: - filepath = default_image + print(msg.get_payload()) + for part in msg.walk(): + if part.get_content_maintype() == "multipart": + continue + if part.get("Content-Disposition") is None: + continue + + filename = part.get_filename() + #filename = filename.split("\\")[-1] + filename = filename.split("/")[-1] + + # Save the image file + with open(os.path.join(recipient_directory, filename), "wb") as f: + f.write(part.get_payload(decode=True)) + print(f'\n Opened and parsed new email from {sender} to {recipient} with subject {subject}') + print(f'Email body: {body}') + print(f'Email attachment: {filename}') + + filepath = str(f"{recipient_directory}/{filename}") + try: #We faced some network errors resulting in images being sent partially black. To address this issue, we implemented a try-except block to handle such occurrences. Now, if an image fails to send correctly, a default image is sent for that experiment. + with open(filepath) as f: # TEST IF THE FILE IS A VALID IMAGE + img = MIMEImage(f.read()) + except: # network error + if default_image=='': + print('Network Error: No default image is set') + return + else: + filepath = default_image - return (sender, recipient, subject, body, filepath) + return (sender, recipient, subject, body, filepath) + except Exception as e: + print(f"Error parsing email data: {e}") + return None def send_Email(Command, sender, recipient, subject, body, attachment_path, SERVER_HOST, SERVER_PORT, AdditionalQuery=['']): # this function sends a new email to the email server - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket: - client_socket.connect((SERVER_HOST, SERVER_PORT)) - - # Create the message - msg = MIMEMultipart() - msg["Command"] = Command - msg["Subject"] = subject - msg["From"] = sender - msg["To"] = recipient - - if AdditionalQuery != '': - for i in range(len(AdditionalQuery)): - msg["AdditionalQuery" + str(i)] = AdditionalQuery[i] - msg["AdditionalQueryNum"] = str(len(AdditionalQuery)) - msg.attach(MIMEText(body, "plain")) - - filename = attachment_path - with open(filename, "rb") as f: - img = MIMEImage(f.read()) - img.add_header("Content-Disposition", "attachment", filename=filename) - msg.attach(img) - message = msg.as_string().encode('utf-8') - - client_socket.sendall(message) # send the message to the server - response = receive_complete_data(client_socket) # get the response from the server - - return response.decode('utf-8') + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket: + client_socket.connect((SERVER_HOST, SERVER_PORT)) + + # Create the message + msg = MIMEMultipart() + msg["Command"] = Command + msg["Subject"] = subject + msg["From"] = sender + msg["To"] = recipient + + if AdditionalQuery != '': + for i in range(len(AdditionalQuery)): + msg["AdditionalQuery" + str(i)] = AdditionalQuery[i] + msg["AdditionalQueryNum"] = str(len(AdditionalQuery)) + msg.attach(MIMEText(body, "plain")) + + filename = attachment_path + with open(filename, "rb") as f: + img = MIMEImage(f.read()) + img.add_header("Content-Disposition", "attachment", filename=filename) + msg.attach(img) + message = msg.as_string().encode('utf-8') + + client_socket.sendall(message) # send the message to the server + response = receive_complete_data(client_socket) # get the response from the server + + return response.decode('utf-8') + except Exception as e: + print(f"Error sending email: {e}") + return "Error sending email" def show_email_popup(email_data): # this function shows a popup with the email data @@ -188,10 +196,11 @@ def check_email_inbox(): # this function checks the inbox for new emails from t client_socket.close() try: email_data = parse_email_data(inbox_data) - show_email_popup(email_data) - Handle_New_Inbox_Email(email_data) - except: - pass + if email_data: + show_email_popup(email_data) + Handle_New_Inbox_Email(email_data) + except Exception as e: + print(f"Error handling new inbox email: {e}") def read_emails_from_file(): # this function reads 5 emails from the Email csv file and returns them as a list diff --git a/dashboard/dashboard.py b/dashboard/dashboard.py index 17ec297..3b98502 100644 --- a/dashboard/dashboard.py +++ b/dashboard/dashboard.py @@ -270,6 +270,7 @@ def add_tool_tips(): session.commit() except Exception as e: logging.error(f"Error saving dashboard data to database: {e}") + session.rollback() finally: session.close()