diff --git a/osaca/api/kerncraft_interface.py b/osaca/api/kerncraft_interface.py index 54360c2..28a1894 100755 --- a/osaca/api/kerncraft_interface.py +++ b/osaca/api/kerncraft_interface.py @@ -1,6 +1,8 @@ #!/usr/bin/env python3 import collections +import sys +from io import StringIO from osaca.frontend import Frontend from osaca.parser import ParserAArch64v81, ParserX86ATT @@ -8,6 +10,19 @@ from osaca.semantics import (INSTR_FLAGS, KernelDG, MachineModel, SemanticsAppender, reduce_to_section) +# Stolen from https://stackoverflow.com/a/16571630 +class Capturing(list): + def __enter__(self): + self._stdout = sys.stdout + sys.stdout = self._stringio = StringIO() + return self + + def __exit__(self, *args): + self.extend(self._stringio.getvalue().splitlines()) + del self._stringio # free up some memory + sys.stdout = self._stdout + + class KerncraftAPI(object): def __init__(self, arch, code): self.machine_model = MachineModel(arch=arch) @@ -25,7 +40,9 @@ class KerncraftAPI(object): def create_output(self, verbose=False): kernel_graph = KernelDG(self.kernel, self.parser, self.machine_model) frontend = Frontend(arch=self.machine_model.get_arch()) - frontend.print_full_analysis(self.kernel, kernel_graph, verbose=verbose) + with Capturing() as output: + frontend.print_full_analysis(self.kernel, kernel_graph, verbose=verbose) + return '\n'.join(output) def get_unmatched_instruction_ratio(self): unmatched_counter = 0 diff --git a/osaca/parser/parser_x86att.py b/osaca/parser/parser_x86att.py index 045eb21..149d8de 100755 --- a/osaca/parser/parser_x86att.py +++ b/osaca/parser/parser_x86att.py @@ -212,9 +212,8 @@ class ParserX86ATT(BaseParser): { self.INSTRUCTION_ID: result['mnemonic'], self.OPERANDS_ID: operands, - self.COMMENT_ID: ' '.join(result[self.COMMENT_ID]) - if self.COMMENT_ID in result - else None, + self.COMMENT_ID: + ' '.join(result[self.COMMENT_ID]) if self.COMMENT_ID in result else None, } ) return return_dict