apol: Add graphical output for Infoflow and DTA.

Signed-off-by: Chris PeBenito <pebenito@ieee.org>

setoolsgui/widgets/queryupdater.py

@@ -6,7 +6,8 @@
 import logging
 import typing
 
-from PyQt6 import QtCore
+import networkx as nx
+from PyQt6 import QtCore, QtGui, QtWidgets
 import setools
 
 from . import models
@@ -42,17 +43,19 @@ class QueryResultsUpdater(QtCore.QObject, typing.Generic[Q]):
     finished = QtCore.pyqtSignal(int)
     raw_line = QtCore.pyqtSignal(str)
 
-    def __init__(self, query: Q,
-                 model: models.SEToolsTableModel | None = None,
+    def __init__(self, query: Q, /, *,
+                 table_model: models.SEToolsTableModel | None = None,
+                 graphics_buffer: QtWidgets.QLabel | None = None,
                  render: RenderFunction = lambda _, x: str(x),
                  result_limit: int = 0) -> None:
 
         super().__init__()
         self.log: typing.Final = logging.getLogger(query.__module__)
         self.query: typing.Final[Q] = query
-        self.model = model
+        self.table_model = table_model
         self.render = render
         self.result_limit = result_limit
+        self.graphics_buffer = graphics_buffer
 
     def update(self) -> None:
         """Run the query and update results."""
@@ -83,8 +86,19 @@ class QueryResultsUpdater(QtCore.QObject, typing.Generic[Q]):
 
             self.log.info(f"Generated {counter} total results.")
 
-            if self.model:
-                self.model.item_list = results
+            if self.table_model:
+                self.log.info("Updating results in table model.")
+                self.table_model.item_list = results
+
+            if self.graphics_buffer:
+                assert isinstance(self.query, setools.query.DirectedGraphAnalysis)
+                self.log.info("Generating graphical results.")
+                self.graphics_buffer.clear()
+                pgv = nx.nx_agraph.to_agraph(self.query.graphical_results())
+                pic = QtGui.QPixmap()
+                pic.loadFromData(pgv.draw(prog="dot", format="png"), "PNG")
+                self.graphics_buffer.setPixmap(pic)
+
             self.finished.emit(counter)
 
         except Exception as e:

setoolsgui/widgets/tab.py

@@ -8,7 +8,7 @@ import typing
 from PyQt6 import QtCore, QtGui, QtWidgets
 import setools
 
-from . import criteria, exception, models, views
+from . import criteria, exception, models, util, views
 from .models.typing import QObjectType
 from .queryupdater import QueryResultsUpdater
 
@@ -462,7 +462,7 @@ class TableResultTabWidget(BaseAnalysisTabWidget):
         """Set the table results model for this tab and set up the processing thread for it."""
         self.sort_proxy.setSourceModel(model)
 
-        self.worker = QueryResultsUpdater(self.query, model)
+        self.worker = QueryResultsUpdater(self.query, table_model=model)
         self.worker.moveToThread(self.processing_thread)
         self.worker.raw_line.connect(self.raw_results.appendPlainText)
         self.worker.finished.connect(self.query_completed)
@@ -583,17 +583,27 @@ class DirectedGraphResultTab(BaseAnalysisTabWidget, typing.Generic[DGA]):
         sizePolicy.setHeightForWidth(self.results.sizePolicy().hasHeightForWidth())
 
         #
-        # Create placeholder future graphical tab
+        # Create graphical tab
         #
-        self.graphical_results = QtWidgets.QWidget(self.results)
+        self.graphical_scroll = QtWidgets.QScrollArea(self.results)
+        self.graphical_scroll.setSizeAdjustPolicy(
+            QtWidgets.QAbstractScrollArea.SizeAdjustPolicy.AdjustToContents)
+        self.graphical_scroll.setWidgetResizable(True)
+        self.results.addTab(self.graphical_scroll, "Graphical Results")
+
+        image_size_policy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Policy.MinimumExpanding,
+                                                  QtWidgets.QSizePolicy.Policy.MinimumExpanding)
+        image_size_policy.setHorizontalStretch(1)
+        image_size_policy.setVerticalStretch(1)
+        image_size_policy.setHeightForWidth(True)
+
+        self.graphical_results = QtWidgets.QLabel(self.graphical_scroll)
         self.graphical_results.setObjectName("graphical_results")
-        self.graphical_results.setSizePolicy(sizePolicy)
-        self.results.addTab(self.graphical_results, "Graphical Results")
-        self.results.setTabEnabled(DirectedGraphResultTab.ResultTab.Graph, False)
-        self.results.setTabWhatsThis(DirectedGraphResultTab.ResultTab.Graph,
-                                     "Future graphical results feature.")
-        self.results.setTabToolTip(DirectedGraphResultTab.ResultTab.Graph,
-                                   "Future graphical results feature.")
+        self.graphical_results.setSizePolicy(image_size_policy)
+        self.graphical_results.setContextMenuPolicy(QtCore.Qt.ContextMenuPolicy.CustomContextMenu)
+        self.graphical_results.customContextMenuRequested.connect(
+            self._graphical_results_context_menu)
+        self.graphical_scroll.setWidget(self.graphical_results)
 
         #
         # Create tree browser tab
@@ -627,7 +637,7 @@ class DirectedGraphResultTab(BaseAnalysisTabWidget, typing.Generic[DGA]):
                                      "<b>This tab has plain text results of the query.</b>")
 
         # set initial tab
-        self.results.setCurrentIndex(DirectedGraphResultTab.ResultTab.Tree)
+        self.results.setCurrentIndex(DirectedGraphResultTab.ResultTab.Graph)
 
         # set up processing thread
         self.processing_thread = QtCore.QThread(self.analysis_widget)
@@ -641,7 +651,7 @@ class DirectedGraphResultTab(BaseAnalysisTabWidget, typing.Generic[DGA]):
         self.busy.reset()
 
         # set up results worker
-        self.worker = QueryResultsUpdater[DGA](self.query)
+        self.worker = QueryResultsUpdater[DGA](self.query, graphics_buffer=self.graphical_results)
         self.worker.moveToThread(self.processing_thread)
         self.worker.raw_line.connect(self.raw_results.appendPlainText)
         self.worker.finished.connect(self.query_completed)
@@ -662,7 +672,33 @@ class DirectedGraphResultTab(BaseAnalysisTabWidget, typing.Generic[DGA]):
     @tree_results_model.setter
     def tree_results_model(self, model: models.SEToolsTableModel) -> None:
         self.tree_results.setModel(model)
-        self.worker.model = model
+        self.worker.table_model = model
+
+    def _graphical_results_context_menu(self, pos: QtCore.QPoint) -> None:
+        """Generate context menu for graphical results widget."""
+        save_action = QtGui.QAction("Save As...", self.graphical_results)
+        save_action.triggered.connect(self._save_graphical_results)
+
+        menu = QtWidgets.QMenu(self.graphical_results)
+        menu.setAttribute(QtCore.Qt.WidgetAttribute.WA_DeleteOnClose)
+        menu.addActions((save_action,))
+        menu.exec(self.graphical_results.mapToGlobal(pos))
+
+    def _save_graphical_results(self) -> None:
+        """Save the graphical results to a file."""
+        with util.QMessageOnException("Error",
+                                      "<b>Failed to save graphical results.</b>",
+                                      log=self.log,
+                                      parent=self):
+
+            filename, _ = QtWidgets.QFileDialog.getSaveFileName(
+                self, "Save graphical results", "", "PNG files (*.png);;All files (*)")
+
+            if filename:
+                if not self.graphical_results.pixmap().save(filename, format="PNG"):
+                    # The save method does not raise an exception, so unfortunately
+                    # there is no additional info to share with the user.
+                    raise RuntimeError(f"Failed to save graphical results to {filename}.")
 
     #
     # Start/end of processing

setoolsgui/widgets/util.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: LGPL-2.1-only
+
+import logging
+import traceback
+import types
+import typing
+
+from PyQt6 import QtWidgets
+
+__all__: typing.Final[tuple[str, ...]] = ("QMessageOnException",)
+
+
+class QMessageOnException:
+
+    """Context manager to display a message box on exception."""
+
+    def __init__(self, title: str, message: str, /, *,
+                 suppress: bool = True,
+                 log: logging.Logger | None = None,
+                 icon: QtWidgets.QMessageBox.Icon = QtWidgets.QMessageBox.Icon.Critical,
+                 parent: QtWidgets.QWidget | None = None) -> None:
+
+        self.title: typing.Final[str] = title
+        self.message: typing.Final[str] = message
+        self.suppress: typing.Final[bool] = suppress
+        self.parent: typing.Final[QtWidgets.QWidget | None] = parent
+        self.log: typing.Final[logging.Logger] = log if log else logging.getLogger(__name__)
+        self.icon: typing.Final[QtWidgets.QMessageBox.Icon] = icon
+
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(self,
+                 exc_type: type[BaseException] | None,
+                 exc_value: BaseException | None,
+                 tb: types.TracebackType | None) -> bool:
+
+        if exc_type:
+            self.log.critical(self.message)
+            self.log.debug("Backtrace", exc_info=True)
+
+            msg = QtWidgets.QMessageBox(self.icon,
+                                        self.title,
+                                        self.message,
+                                        parent=self.parent)
+
+            msg.setInformativeText(str(exc_value))
+            msg.setDetailedText("\n".join(traceback.format_tb(tb)))
+            msg.exec()
+
+            return self.suppress
+
+        return False