[Decoration] Exporting of rendered plots fails on windows (#370)

* Create draft PR for #369

* Fix export_plot on windows

* black

* Move temporary file creation out of method

* Add missed flush

* add tests for png export

* add graphviz to dockerfile requirements

* Use CloseableNamedTemporaryFile.py

* Remove unnecessary use of temporary file

* Add missing file deletion for 'closeable_temporary_file'

* Rename CloseableNamedTemporaryFile.py

* test CloseableNamedTemporaryFile

* close CloseableNamedTemporaryFile before removing

* black

---------

Co-authored-by: rihi <[email protected]>
Co-authored-by: rihi <[email protected]>
Co-authored-by: Manuel Blatt <[email protected]>
Co-authored-by: Manuel Blatt <[email protected]>

Dockerfile

@@ -22,6 +22,7 @@ RUN apt -y update && apt -y upgrade && apt install -y --no-install-recommends \
     virtualenv \
     unzip \
 	astyle \
+    graphviz \
     # plotting ascii graphs for debug purposes
     libgraph-easy-perl \
     z3

decompiler/util/closeable_named_temporary_file.py

@@ -0,0 +1,29 @@
+import os
+from contextlib import contextmanager
+from tempfile import NamedTemporaryFile
+
+
+@contextmanager
+def CloseableNamedTemporaryFile(**kwargs):
+    """
+    Context manager wrapper for NamedTemporaryFile, which allows
+    closing the file handle without deleting the underling file.
+    Deletion is delegated to the context manager closing.
+
+    Note: With Python 3.12, a new parameter 'delete_on_close' is introduced
+    for NamedTemporaryFile which accomplishes the same thing. Consequently,
+    this api should be replaced when the code is updated to 3.12.
+    """
+
+    kwargs["delete"] = False
+    with NamedTemporaryFile(**kwargs) as file:
+        try:
+            yield file
+        finally:
+            # Close the file to be sure that it can be removed.
+            # It's ok if the file was already closed because NamedTemporaryFile's close method is idempotent.
+            file.close()
+            # If file was already deleted outside of this contextmanager, this will crash
+            # (just like the original NamedTemporaryFile).
+            # On NT, this might also crash if another handle to this file is still open
+            os.remove(file.name)

decompiler/util/decoration.py

@@ -1,14 +1,12 @@
 """Module handling plotting and pretty printing."""
 from __future__ import annotations
 
-import os
 import textwrap
 from logging import warning
 from re import compile
 from subprocess import CompletedProcess, run
 from sys import stdout
-from tempfile import NamedTemporaryFile
-from typing import Dict, Optional, TextIO
+from typing import Dict, TextIO
 
 import z3
 from binaryninja import BranchType, EdgePenStyle, EdgeStyle, FlowGraph, FlowGraphNode, HighlightStandardColor, ThemeColor, show_graph_report
@@ -30,11 +28,13 @@
 from decompiler.structures.ast.syntaxtree import AbstractSyntaxTree
 from decompiler.structures.graphs.cfg import BasicBlock, BasicBlockEdge, BasicBlockEdgeCondition, ControlFlowGraph
 from decompiler.structures.pseudo.operations import Condition
+from decompiler.util.closeable_named_temporary_file import CloseableNamedTemporaryFile
 from decompiler.util.to_dot_converter import ToDotConverter
 from networkx import DiGraph
-from pygments import format, lex
+from pygments import highlight
 from pygments.formatters.html import HtmlFormatter
-from pygments.lexers.c_like import CLexer
+from pygments.formatters.terminal import TerminalFormatter
+from pygments.lexers.c_cpp import CppLexer
 
 try:
     run(["graph-easy", "-v"], capture_output=True)
@@ -59,22 +59,21 @@ def graph(self) -> DiGraph:
         """Return the graph being decorated."""
         return self._graph
 
-    def _write_dot(self, handle: Optional[TextIO] = None):
-        """Write the graph to the given handle or NamedTemporaryFile."""
-        if not handle:
-            handle = NamedTemporaryFile(mode="w+")
-        ToDotConverter.write(self._graph, handle)
-        handle.flush()
-        handle.seek(0)
-        return handle
+    def _write_dot(self, handle: TextIO):
+        """Write the graph to the given handle."""
+        handle.write(ToDotConverter.write(self._graph))
 
     def export_ascii(self) -> str:
         """Export the current graph into an ascii representation."""
         if not GRAPH_EASY_INSTALLED:
             warning(f"Invoking graph-easy although it seems like it is not installed on the system.")
-        with self._write_dot() as handle:
-            result: CompletedProcess = run(["graph-easy", "--as=ascii", handle.name], capture_output=True)
-        return result.stdout.decode("utf-8")
+
+        with CloseableNamedTemporaryFile(mode="w", encoding="utf-8") as file:
+            self._write_dot(file)
+            file.close()
+
+            result: CompletedProcess = run(["graph-easy", "--as=ascii", file.name], capture_output=True)
+            return result.stdout.decode("utf-8")
 
     def export_dot(self, path: str):
         """Export the graph into a dotfile at the given location."""
@@ -88,10 +87,14 @@ def export_plot(self, path: str, type="png"):
         path -- Path to the plot to be created.
         type -- a string describing the output type (commonly pdf, png)
         """
-        with self._write_dot() as handle:
-            result = run(["dot", f"-T{type}", f"-o{path}", f"{handle.name}"], capture_output=True)
-        if result.returncode:
-            raise ValueError(f"Could not plot graph! ({result.stderr.decode('utf-8')}")
+
+        with CloseableNamedTemporaryFile(mode="w", encoding="utf-8") as file:
+            self._write_dot(file)
+            file.close()
+
+            result = run(["dot", f"-T{type}", f"-o{path}", f"{file.name}"], capture_output=True)
+            if result.returncode:
+                raise ValueError(f"Could not plot graph! ({result.stderr.decode('utf-8')}")
 
 
 class DecoratedCFG(DecoratedGraph):
@@ -301,22 +304,6 @@ def _format_node_content(label: str, max_width: int = 60):
 class DecoratedCode:
     """Class representing C code ready for pretty printing."""
 
-    class TempFile:
-        """Context manager to write content to NamedTemporaryFile and release for windows, returns file name"""
-
-        def __init__(self, content: str):
-            self.tmpf = NamedTemporaryFile(mode="w", delete=False)
-            self.tmpf.write(content)
-            self.name = self.tmpf.name
-            self.tmpf.flush()
-            self.tmpf.close()
-
-        def __enter__(self) -> str:
-            return self.name
-
-        def __exit__(self, exc_type, exc_val, exc_tb):
-            os.unlink(self.name)
-
     def __init__(self, code: str, style="paraiso-dark"):
         """Generate an object handling code decoration."""
         self._text = code
@@ -358,20 +345,22 @@ def reformat(self):
         """Call astyle on command line to reformat the code."""
         if not ASTYLE_INSTALLED:
             warning(f"Invoking astyle although it seems like it is not installed on the system.")
-        with self.TempFile(self._text) as filename:
-            run(["astyle", "-z2", "-n", filename], check=True, capture_output=True)
-            with open(filename, "r") as output:
+
+        with CloseableNamedTemporaryFile(mode="w", encoding="utf-8") as file:
+            file.write(self._text)
+            file.close()
+
+            run(["astyle", "-z2", "-n", file.name], check=True, capture_output=True)
+
+            with open(file.name, "r") as output:
                 self._text = output.read()
 
     def export_ascii(self) -> str:
-        with self.TempFile(self._text) as filename:
-            result: CompletedProcess = run(["pygmentize", "-l", "cpp", f"-O style={self._style}", filename], capture_output=True)
-        return result.stdout.decode("ascii")
+        return highlight(self._text, CppLexer(), TerminalFormatter(style=self._style))
 
     def export_html(self) -> str:
         """Export an html representation of the current code."""
-        tokens = lex(self._text, CLexer())
-        html = format(tokens, HtmlFormatter(full=True, style=self._style))
+        html = highlight(self._text, CppLexer(), HtmlFormatter(full=True, style=self._style))
         return self._filter_css_comments(html)
 
     @staticmethod

decompiler/util/to_dot_converter.py

@@ -1,5 +1,4 @@
 """Module handling conversion to dot-format."""
-from typing import TextIO
 
 from networkx import DiGraph
 
@@ -16,10 +15,10 @@ def __init__(self, graph: DiGraph):
         self._graph = graph
 
     @classmethod
-    def write(cls, graph: DiGraph, handle: TextIO):
+    def write(cls, graph: DiGraph) -> str:
         """Write dot-format of given graph into handle."""
         converter = cls(graph)
-        handle.write(converter._create_dot())
+        return converter._create_dot()
 
     def _create_dot(self) -> str:
         """Create dot-file content."""

tests/util/test_closeable_named_temporary_file.py

@@ -0,0 +1,27 @@
+import os
+
+from decompiler.util.closeable_named_temporary_file import CloseableNamedTemporaryFile
+
+
+class TestCloseableNamedTemporaryFile:
+    def test_usage_after_closing(self):
+        with CloseableNamedTemporaryFile(mode="w") as file:
+            file.write("test")
+            file.close()
+            with open(file.name, "r") as reopened_file:
+                assert reopened_file.read() == "test"
+
+    def test_deletion_with_close(self):
+        with CloseableNamedTemporaryFile(mode="w") as file:
+            file.close()
+        assert not os.path.exists(file.name)
+
+    def test_deletion_without_close(self):
+        with CloseableNamedTemporaryFile(mode="w") as file:
+            pass
+        assert not os.path.exists(file.name)
+
+    def test_close_after_delete(self):
+        with CloseableNamedTemporaryFile(mode="w") as file:
+            pass
+        file.close()

tests/util/test_decoration.py

@@ -1,3 +1,6 @@
+import os
+from io import StringIO
+
 import pytest
 from binaryninja import HighlightStandardColor
 from decompiler.structures.ast.ast_nodes import SeqNode
@@ -214,6 +217,22 @@ def test_convert_to_dot_with_string(self, graph_with_string):
 }"""
         )
 
+    def test_png_export(self, simple_graph, png_magic):
+        decorated = DecoratedCFG.from_cfg(simple_graph)
+        export_path = "_test_cfg.png"
+        # remove potential left over file
+        try:
+            os.remove(export_path)
+        except OSError:
+            pass
+        # export plot
+        decorated.export_plot(export_path)
+        # check if exported plot has correct value (and implicitly if it exists)
+        with open(export_path, "rb") as file:
+            assert file.read(len(png_magic)) == png_magic
+        # remove file
+        os.remove(export_path)
+
 
 class TestDecoratedAST:
     @pytest.fixture
@@ -442,9 +461,11 @@ def test_ascii_switch(self, ast_switch):
     @pytest.mark.usefixtures("ast_condition")
     def test_dotviz_output(self, ast_condition):
         decorated = DecoratedAST.from_ast(ast_condition)
-        handle = decorated._write_dot()
-        data = handle.read()
-        handle.close()
+
+        with StringIO() as stringIo:
+            decorated._write_dot(stringIo)
+            data = stringIo.getvalue()
+
         assert all(
             [
                 x in data
@@ -608,6 +629,22 @@ def test_convert_to_dot_switch(self, ast_switch):
 }"""
         )
 
+    def test_png_export(self, ast_for_loop, png_magic):
+        decorated = DecoratedAST.from_ast(ast_for_loop)
+        export_path = "_test_ast.png"
+        # remove potential left over file
+        try:
+            os.remove(export_path)
+        except OSError:
+            pass
+        # export plot
+        decorated.export_plot(export_path)
+        # check if exported plot has correct value (and implicitly if it exists)
+        with open(export_path, "rb") as file:
+            assert file.read(len(png_magic)) == png_magic
+        # remove file
+        os.remove(export_path)
+
 
 class TestDecoratedCode:
     @pytest.fixture
@@ -640,3 +677,8 @@ def test_does_html_output(self, simple_code):
         html_code = decorated.export_html()
         assert len(html_code) > 100
         assert r"\*" not in html_code
+
+
+@pytest.fixture
+def png_magic():
+    return bytes.fromhex("89 50 4e 47 0d 0a 1a 0a")