Session object (#285)

* Session Object

* Remove unused bookmarks dictionary and update session data structure to include process name

* change behavior based on feedback

- Saving now always allowed
- moved Version functions into session
- renamed Session to session to follow naming conventions

* Implement symbol resolution for loaded bookmarks

* add bookmark resolver by process region

Also add index to memory_region table

* remove symbol resolve
--> only using region resolving

* Add TODO for a potential symbol resolution

---------

Co-authored-by: Korcan Karaokçu <[email protected]>

Author: Bloodiko

GUI/MainWindow.py

@@ -0,0 +1,290 @@
+import os
+from enum import IntFlag, auto
+
+from PyQt6.QtCore import QObject
+from PyQt6.QtGui import QCloseEvent
+from PyQt6.QtWidgets import QFileDialog, QMessageBox
+
+from GUI.States import states
+from libpince import utils, debugcore
+from tr.tr import TranslationConstants as tr
+
+
+class SessionDataChanged(IntFlag):
+    NONE = auto()
+    ADDRESS_TREE = auto()
+    BOOKMARKS = auto()
+    NOTES = auto()
+    PROCESS_NAME = auto()
+
+
+def migrate_version(content: any) -> dict[str, any]:
+    if not hasattr(content, "version") and type(content) == list:
+        return legacy_to_v1(content)
+
+    return content
+
+
+def is_valid_session_data(content: dict[str, any]) -> bool:
+    keys = ["version", "notes", "bookmarks", "address_tree", "process_name"]
+    for key in keys:
+        if key not in content:
+            return False
+
+    return True
+
+
+def legacy_to_v1(content: list) -> dict[str, any]:
+    print("Migrating legacy session data to version 1")
+    return {"version": 1, "notes": "", "bookmarks": {}, "address_tree": content, "process_name": ""}
+
+
+class Session:
+    def __init__(self) -> None:
+        # Anything labled with pct should be saved to the session file
+        self.pct_notes: str = ""
+        self.pct_bookmarks: dict[int, dict] = {}
+        self.pct_version: int = 1
+        self.pct_address_tree: list = []
+        self.pct_process_name: str = ""
+        self.data_changed = SessionDataChanged.NONE
+        self.file_path: str = os.curdir
+        self.last_file_name: str = ""  # process name or file name
+
+    def save_session(self) -> bool:
+        """
+        Save the current session to a file.
+
+        Args:
+            None
+        Returns:
+            bool: True if the session was saved successfully, False otherwise.
+        """
+
+        file_path, _ = QFileDialog.getSaveFileName(
+            None, tr.SAVE_PCT_FILE, self.file_path + "/" + self.last_file_name, tr.FILE_TYPES_PCT
+        )
+        if not file_path:
+            return False
+
+        # until address tree is model view and properly read from this new session object,
+        # address tree must save its data to the session object via signal
+        if self.data_changed & SessionDataChanged.ADDRESS_TREE:
+            states.session_signals.on_save.emit()
+
+        session = {
+            "version": self.pct_version,
+            "notes": self.pct_notes,
+            "bookmarks": self.pct_bookmarks,
+            "address_tree": self.pct_address_tree,
+            "process_name": self.pct_process_name,
+        }
+
+        file_path = utils.append_file_extension(file_path, "pct")
+        if not utils.save_file(session, file_path):
+            QMessageBox.information(None, tr.ERROR, tr.FILE_SAVE_ERROR)
+            return False
+
+        self.file_path = os.path.dirname(file_path)
+        self.last_file_name = os.path.basename(file_path)
+        self.data_changed = SessionDataChanged.NONE
+        return True
+
+    def check_unsaved_changes(self) -> QMessageBox.StandardButton:
+        if self.data_changed == SessionDataChanged.NONE:
+            return QMessageBox.StandardButton.No
+
+        unsaved_changes_result = QMessageBox.question(
+            None,
+            tr.SAVE_SESSION_QUESTION_TITLE,
+            tr.SAVE_SESSION_QUESTION_PROMPT,
+            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No | QMessageBox.StandardButton.Cancel,
+        )
+        return unsaved_changes_result
+
+    def load_session(self) -> bool:
+        """
+        Load a pct session file. Will check for unsaved changes and prompt
+        the user to save them before loading a new session.
+        If the user chooses to cancel, the function will return False.
+        Will also attempt to migrate the session data to the latest version.
+
+        Args:
+            None
+        Returns:
+            bool: True if the session was loaded successfully, False otherwise.
+
+        """
+
+        unsaved_changes_result = self.check_unsaved_changes()
+        if unsaved_changes_result == QMessageBox.StandardButton.Cancel:
+            return False
+        elif unsaved_changes_result == QMessageBox.StandardButton.Yes:
+            if not self.save_session():
+                return False
+
+        file_path, _ = QFileDialog.getOpenFileName(
+            None, tr.OPEN_PCT_FILE, self.file_path + "/" + self.last_file_name, tr.FILE_TYPES_PCT
+        )
+        if not file_path:
+            return False
+
+        content = utils.load_file(file_path)
+        if content is None:
+            QMessageBox.information(None, tr.ERROR, tr.FILE_LOAD_ERROR.format(file_path))
+            return False
+        content = migrate_version(content)
+        if not is_valid_session_data(content):
+            QMessageBox.information(None, tr.ERROR, tr.FILE_LOAD_ERROR.format(file_path))
+            return False
+
+        self.pct_version = content["version"]
+        self.pct_notes = content["notes"]
+
+        # Load bookmarks with symbol resolution
+        self.pct_bookmarks = content["bookmarks"]
+        self.recalculate_bookmarks()
+
+        self.pct_address_tree = content["address_tree"]
+
+        self.file_path = os.path.dirname(file_path)
+        self.last_file_name = os.path.basename(file_path)
+
+        states.session_signals.on_load.emit()
+        self.data_changed = SessionDataChanged.NONE
+        return True
+
+    def pre_exit(self, close_event: QCloseEvent) -> None:
+        """
+        Event handler for the close event of the application.
+        If there are unsaved changes, prompt the user to save them.
+        Accepts or ignores the close event based on the user's choice.
+
+        Args:
+            close_event (QCloseEvent): The close event to be handled.
+        Returns:
+            None
+        """
+
+        if self.data_changed == SessionDataChanged.NONE:
+            close_event.accept()
+            return
+
+        pre_exit_unsaved_changes_result = self.check_unsaved_changes()
+        if pre_exit_unsaved_changes_result == QMessageBox.StandardButton.Yes:
+            self.save_session()
+            close_event.accept()
+
+        elif pre_exit_unsaved_changes_result == QMessageBox.StandardButton.Cancel:
+            close_event.ignore()
+        else:
+            close_event.accept()
+
+    def recalculate_bookmarks(self):
+        """
+        Recalculate all bookmarks to update their addresses based on their symbols or region info.
+        This is useful when the process has changed and the addresses may have shifted.
+        Will not mark session as changed.
+
+        Args:
+            None
+        Returns:
+            None
+        """
+
+        region_dict = utils.get_region_dict(debugcore.currentpid)
+
+        new_bookmarks: dict[int, dict] = {}
+        for addr, value in self.pct_bookmarks.items():
+            comment = value["comment"]
+            symbol = value["symbol"]
+            address_region_details = value["address_region_details"]
+            new_addr = addr
+
+            # We used to resolve the symbol first but parentheses in symbols causes the functions to be called
+            # For instance, main() won't be resolved to main but rather the function main() will be called
+            # Because of this behavior, we rely on resolving via regions
+            # Resolving via regions is more reliable unless the binary has been changed
+            # TODO: Resolving via symbols could be re-addressed if this behavior were fixed
+
+            # resolve via region details
+            region_name, offset, region_index = address_region_details.values()
+            region = region_dict.get(region_name, None)
+            if region is not None:
+                new_addr = int(region[region_index], 16) + int(offset, 16)
+            else:
+                print("[WARN] BookmarkResolver: Could not find region with name:", region_name)
+                continue
+
+            new_bookmarks[new_addr] = {
+                "symbol": symbol,
+                "comment": comment,
+                "address_region_details": address_region_details,
+            }
+
+        self.pct_bookmarks = new_bookmarks
+
+
+class SessionManager:
+    session = Session()
+
+    @staticmethod
+    def get_session() -> Session:
+        return SessionManager.session
+
+    @staticmethod
+    def reset_session() -> None:
+        session = SessionManager.get_session()
+        # User has one last chance to save the session before resetting
+        # result is ignored, because the session is going to be reset anyway
+        session.check_unsaved_changes()
+        SessionManager.session = Session()
+        states.session_signals.new_session.emit()
+        # Reset the session data changed flag
+        session.data_changed = SessionDataChanged.NONE
+
+    @staticmethod
+    def save_session() -> None:
+        SessionManager.get_session().save_session()
+
+    @staticmethod
+    def load_session() -> None:
+        SessionManager.get_session().load_session()
+
+    @staticmethod
+    def on_process_changed() -> None:
+        if debugcore.currentpid == -1:
+            return
+
+        if states.exiting:
+            return
+
+        process_name = utils.get_process_name(debugcore.currentpid)
+        session = SessionManager.get_session()
+
+        if session.pct_process_name == process_name:
+            # process is the same as last one, probably process restarted / reattached
+            session.recalculate_bookmarks()
+            return
+
+        if session.pct_process_name == "":
+            # silently set the process name and file name if necessary
+            session.pct_process_name = process_name
+            if session.last_file_name == "":
+                session.last_file_name = utils.append_file_extension(process_name, "pct")
+            return
+
+        if session.pct_process_name != process_name:
+            # Ask if the user wants to keep the session
+            keep_session_result = QMessageBox.question(
+                None,
+                tr.SESSION_PROCESS_CHANGED_TITLE,
+                tr.SESSION_PROCESS_CHANGED_PROMPT,
+                QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
+            )
+            if keep_session_result == QMessageBox.StandardButton.Yes:
+                session.pct_process_name = process_name
+            else:
+                SessionManager.reset_session()
+                session.pct_process_name = process_name
+                session.last_file_name = utils.append_file_extension(process_name, "pct")

GUI/MainWindow.ui

@@ -3,6 +3,7 @@
 
 process_signals = guitypedefs.ProcessSignals()
 setting_signals = guitypedefs.SettingSignals()
+session_signals = guitypedefs.SessionSignals()
 
 status_thread = guitypedefs.CheckInferiorStatus()
 status_thread.start()
@@ -19,10 +20,6 @@
 # Currently only used in address_table_loop
 exp_cache: dict[str, str] = {}
 
-# Used by tableWidget_Disassemble in MemoryViewer
-# Format -> {bookmark_address:comment}
-bookmarks: dict[int, str] = {}
-
 # Set to True when app is about to exit
 exiting = False
 

GUI/Session/session.py

@@ -25,6 +25,12 @@ class ProcessSignals(QObject):
     exit = pyqtSignal()
 
 
+class SessionSignals(QObject):
+    new_session = pyqtSignal()
+    on_save = pyqtSignal()
+    on_load = pyqtSignal()
+
+
 class Worker(QRunnable):
     def __init__(self, fn, *args, **kwargs):
         super().__init__()

GUI/States/states.py

@@ -1,11 +1,13 @@
-from PyQt6.QtWidgets import QWidget, QMessageBox, QMenu, QListWidgetItem
 from PyQt6.QtCore import Qt, pyqtSignal
-from PyQt6.QtGui import QShortcut, QKeySequence
+from PyQt6.QtGui import QKeySequence, QShortcut
+from PyQt6.QtWidgets import QListWidgetItem, QMenu, QMessageBox, QWidget
+
+from GUI.Session.session import SessionDataChanged, SessionManager
 from GUI.States import states
 from GUI.Utils import guiutils, utilwidgets
 from GUI.Widgets.Bookmark.Form.BookmarkWidget import Ui_Form
-from tr.tr import TranslationConstants as tr
 from libpince import debugcore, utils
+from tr.tr import TranslationConstants as tr
 
 
 # This widget is too intertwined with MemoryViewer, it will be need to be reworked if it gets used in anywhere else
@@ -26,12 +28,14 @@ def __init__(self, parent):
         self.shortcut_delete.activated.connect(self.delete_record)
         self.shortcut_refresh = QShortcut(QKeySequence("R"), self)
         self.shortcut_refresh.activated.connect(self.refresh_table)
+        self.session = SessionManager.get_session()
         self.refresh_table()
+        states.session_signals.new_session.connect(self.on_new_session)
         guiutils.center_to_parent(self)
 
     def refresh_table(self):
         self.listWidget.clear()
-        address_list = [hex(address) for address in states.bookmarks.keys()]
+        address_list = [hex(address) for address in self.session.pct_bookmarks.keys()]
         if debugcore.currentpid == -1:
             self.listWidget.addItems(address_list)
         else:
@@ -45,7 +49,7 @@ def change_display(self, row):
             self.lineEdit_Info.clear()
         else:
             self.lineEdit_Info.setText(debugcore.get_address_info(current_address))
-        self.lineEdit_Comment.setText(states.bookmarks[int(current_address, 16)])
+        self.lineEdit_Comment.setText(self.session.pct_bookmarks[int(current_address, 16)]["comment"])
 
     def listWidget_item_double_clicked(self, item: QListWidgetItem):
         self.double_clicked.emit(utils.extract_address(item.text()))
@@ -60,16 +64,18 @@ def exec_add_entry_dialog(self):
                 return
             self.bookmarked.emit(int(address, 16))
             self.refresh_table()
+            self.session.data_changed |= SessionDataChanged.BOOKMARKS
 
     def exec_change_comment_dialog(self, current_address):
         self.comment_changed.emit(current_address)
         self.refresh_table()
+        self.session.data_changed |= SessionDataChanged.BOOKMARKS
 
     def listWidget_context_menu_event(self, event):
         current_item = guiutils.get_current_item(self.listWidget)
         if current_item:
             current_address = int(utils.extract_address(current_item.text()), 16)
-            if current_address not in states.bookmarks:
+            if current_address not in self.session.pct_bookmarks:
                 QMessageBox.information(self, tr.ERROR, tr.INVALID_ENTRY)
                 self.refresh_table()
                 return
@@ -104,3 +110,8 @@ def delete_record(self):
         current_address = int(utils.extract_address(current_item.text()), 16)
         self.deleted.emit(current_address)
         self.refresh_table()
+        self.session.data_changed |= SessionDataChanged.BOOKMARKS
+
+    def on_new_session(self):
+        self.session = SessionManager.get_session()
+        self.refresh_table()