Added error message for sqlalchemy (RhinoSecurityLabs#189)

* Added error message for sqlalchemy * Fix indentation and writing style consistency Co-authored-by: BFriedland <[email protected]>
aaronrea · Jul 21, 2020 · b345370 · b345370
1 parent 1dd2053
commit b345370
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/pacu.py b/pacu.py
@@ -24,6 +24,7 @@
 
     from core.models import AWSKey, PacuSession
     from setup_database import setup_database_if_not_present
+    from sqlalchemy import exc
     from utils import get_database_connection, set_sigint_handler
 except ModuleNotFoundError as error:
     exception_type, exception_value, tb = sys.exc_info()
@@ -813,7 +814,6 @@ def export_keys(self, command):
         else:
             return
 
-
     ###### Some module notes
     # For any argument that needs a value and a region for that value, use the form
     # value@region
@@ -1577,6 +1577,19 @@ def run_gui(self):
                             local_data=local_data,
                             global_data=global_data
                         )
+
+                # Catch sqlalchemy error
+                elif exception_type == exc.OperationalError:
+                    traceback_text = '\nTraceback (most recent call last):\n{}{}: {}\n\n'.format(''.join(traceback.format_tb(tb)), str(exception_type), str(exception_value))
+                    session, global_data, local_data = self.get_data_from_traceback(tb)
+                    self.log_error(
+                        traceback_text,
+                        exception_info='{}: {}\n\nPacu database error. This could be caused by a recent update in Pacu\'s database\'s structure. If your Pacu has been updated recently, try removing your old db.sqlite3 database file.'.format(exception_type, exception_value),
+                        session=session,
+                        local_data=local_data,
+                        global_data=global_data
+                    )
+
                 else:
                     traceback_text = '\nTraceback (most recent call last):\n{}{}: {}\n\n'.format(''.join(traceback.format_tb(tb)), str(exception_type), str(exception_value))
                     session, global_data, local_data = self.get_data_from_traceback(tb)