diff --git a/Dockerfile b/Dockerfile index 8820baf40..38f4b455d 100644 --- a/Dockerfile +++ b/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 diff --git a/decompiler/util/closeable_named_temporary_file.py b/decompiler/util/closeable_named_temporary_file.py new file mode 100644 index 000000000..90491e4ad --- /dev/null +++ b/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) diff --git a/decompiler/util/decoration.py b/decompiler/util/decoration.py index 2453a4bcc..44a2aa281 100644 --- a/decompiler/util/decoration.py +++ b/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 diff --git a/decompiler/util/to_dot_converter.py b/decompiler/util/to_dot_converter.py index af169e573..b1d459d71 100644 --- a/decompiler/util/to_dot_converter.py +++ b/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.""" diff --git a/tests/util/test_closeable_named_temporary_file.py b/tests/util/test_closeable_named_temporary_file.py new file mode 100644 index 000000000..3d1e7611a --- /dev/null +++ b/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() diff --git a/tests/util/test_decoration.py b/tests/util/test_decoration.py index 152d1fcb6..e055eeeb3 100644 --- a/tests/util/test_decoration.py +++ b/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")