Fix errors 7

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}")
 

app_security/app_vulnerability_scanner.py

@@ -2,54 +2,63 @@
 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)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
 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():

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'))

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

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()