diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/examples/pseudo_python.py | 137 | ||||
| -rw-r--r-- | python/function.py | 53 | ||||
| -rw-r--r-- | python/highlevelil.py | 30 | ||||
| -rw-r--r-- | python/languagerepresentation.py | 56 |
4 files changed, 228 insertions, 48 deletions
diff --git a/python/examples/pseudo_python.py b/python/examples/pseudo_python.py index 0d303d64..8c1a237e 100644 --- a/python/examples/pseudo_python.py +++ b/python/examples/pseudo_python.py @@ -35,10 +35,15 @@ class PseudoPythonFunction(LanguageRepresentationFunction): def perform_init_token_emitter(self, emitter): # Python never allows braces, even if the user has set the option to show them emitter.brace_requirement = BraceRequirement.BracesNotAllowed + emitter.has_collapsable_regions = True def perform_begin_lines(self, instr: HighLevelILInstruction, tokens: HighLevelILTokenEmitter): # Ensure that the function body is indented relative to the function declaration tokens.increase_indent() + tokens.prepend_blank_collapse_indicator() + + def perform_end_lines(self, instr: HighLevelILInstruction, tokens: HighLevelILTokenEmitter): + tokens.prepend_blank_collapse_indicator() def perform_get_expr_text( self, instr: HighLevelILInstruction, tokens: HighLevelILTokenEmitter, settings: Optional[DisassemblySettings], @@ -46,6 +51,8 @@ class PseudoPythonFunction(LanguageRepresentationFunction): statement: bool = False ): with tokens.expr(instr): + if instr.as_ast: + tokens.prepend_instr_collapse_indicator(self.function, instr) if instr.operation == HighLevelILOperation.HLIL_BLOCK: need_separator = False body = instr.body @@ -71,37 +78,54 @@ class PseudoPythonFunction(LanguageRepresentationFunction): elif instr.operation == HighLevelILOperation.HLIL_IF: tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "if ")) self.perform_get_expr_text(instr.condition, tokens, settings) - tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, ":\n")) + tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, ":")) if instr.as_ast: - # Only display the if body when printing the full AST. When printing basic blocks in graph view, - # the body of the if and the else part are rendered as other nodes in the graph. - tokens.begin_scope(ScopeType.BlockScopeType) - self.perform_get_expr_text(instr.true, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.BlockScopeType) + if self.function.is_instruction_collapsed(instr): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + # Python can't put multiple `elif`s on the same line + tokens.new_line() + else: + # Only display the if body when printing the full AST. When printing basic blocks in graph view, + # the body of the if and the else part are rendered as other nodes in the graph. + tokens.begin_scope(ScopeType.BlockScopeType) + self.perform_get_expr_text(instr.true, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.BlockScopeType) # Else statements need to be handled as chains, since "else if" in Python should be rendered # as "elif" statements. if_chain = instr.false + idx = 1 while if_chain is not None and if_chain.operation not in [HighLevelILOperation.HLIL_NOP, HighLevelILOperation.HLIL_UNREACHABLE]: if if_chain.operation == HighLevelILOperation.HLIL_IF: + tokens.prepend_instr_collapse_indicator(self.function, instr, idx) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "elif ")) self.perform_get_expr_text(if_chain.condition, tokens, settings) tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ":")) - tokens.begin_scope(ScopeType.BlockScopeType) - self.perform_get_expr_text(if_chain.true, tokens, settings, + if self.function.is_instruction_collapsed(instr, idx): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + # Python can't put multiple `elif`s on the same line + tokens.new_line() + else: + tokens.begin_scope(ScopeType.BlockScopeType) + self.perform_get_expr_text(if_chain.true, tokens, settings, OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.BlockScopeType) + tokens.end_scope(ScopeType.BlockScopeType) if_chain = if_chain.false else: + tokens.prepend_instr_collapse_indicator(self.function, instr, idx) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "else")) tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ":")) - tokens.begin_scope(ScopeType.BlockScopeType) - self.perform_get_expr_text(if_chain, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.BlockScopeType) + if self.function.is_instruction_collapsed(instr, idx): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.BlockScopeType) + self.perform_get_expr_text(if_chain, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.BlockScopeType) break + idx += 1 elif instr.operation == HighLevelILOperation.HLIL_FOR: tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "for ")) @@ -153,27 +177,36 @@ class PseudoPythonFunction(LanguageRepresentationFunction): self.perform_get_expr_text(instr.update, tokens, settings) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, ":")) if instr.as_ast: - tokens.begin_scope(ScopeType.BlockScopeType) - self.perform_get_expr_text(instr.body, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.BlockScopeType) + if self.function.is_instruction_collapsed(instr): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.BlockScopeType) + self.perform_get_expr_text(instr.body, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.BlockScopeType) elif instr.operation == HighLevelILOperation.HLIL_WHILE: tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "while ")) self.perform_get_expr_text(instr.condition, tokens, settings) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, ":")) if instr.as_ast: - tokens.begin_scope(ScopeType.BlockScopeType) - self.perform_get_expr_text(instr.body, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.BlockScopeType) + if self.function.is_instruction_collapsed(instr): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.BlockScopeType) + self.perform_get_expr_text(instr.body, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.BlockScopeType) elif instr.operation == HighLevelILOperation.HLIL_DO_WHILE: if instr.as_ast: tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "do")) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, ":")) - tokens.begin_scope(ScopeType.BlockScopeType) - self.perform_get_expr_text(instr.body, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.BlockScopeType) + if self.function.is_instruction_collapsed(instr): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.BlockScopeType) + self.perform_get_expr_text(instr.body, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.BlockScopeType) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "while ")) self.perform_get_expr_text(instr.condition, tokens, settings) else: @@ -184,25 +217,32 @@ class PseudoPythonFunction(LanguageRepresentationFunction): tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "match ")) self.perform_get_expr_text(instr.condition, tokens, settings) tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, ":")) - tokens.begin_scope(ScopeType.SwitchScopeType) - if instr.as_ast: - # Output each case - for case in instr.cases: - self.perform_get_expr_text(case, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.new_line() + if self.function.is_instruction_collapsed(instr): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.SwitchScopeType) - # Check for default case - if instr.default is not None and instr.default.operation not in [HighLevelILOperation.HLIL_NOP, - HighLevelILOperation.HLIL_UNREACHABLE]: - tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "default")) - tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ":")) - tokens.begin_scope(ScopeType.CaseScopeType) - self.perform_get_expr_text(instr.default, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.CaseScopeType) - tokens.end_scope(ScopeType.SwitchScopeType) + # Output each case + for case in instr.cases: + self.perform_get_expr_text(case, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.new_line() + + # Check for default case + if instr.default is not None and instr.default.operation not in [HighLevelILOperation.HLIL_NOP, + HighLevelILOperation.HLIL_UNREACHABLE]: + tokens.prepend_instr_collapse_indicator(self.function, instr, 1) + tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "default")) + tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ":")) + if self.function.is_instruction_collapsed(instr, 1): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.CaseScopeType) + self.perform_get_expr_text(instr.default, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.CaseScopeType) + tokens.end_scope(ScopeType.SwitchScopeType) elif instr.operation == HighLevelILOperation.HLIL_CASE: tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "case ")) for (i, value) in enumerate(instr.values): @@ -210,10 +250,13 @@ class PseudoPythonFunction(LanguageRepresentationFunction): tokens.append(InstructionTextToken(InstructionTextTokenType.OperationToken, " | ")) self.perform_get_expr_text(value, tokens, settings) tokens.append(InstructionTextToken(InstructionTextTokenType.TextToken, ":")) - tokens.begin_scope(ScopeType.CaseScopeType) - self.perform_get_expr_text(instr.body, tokens, settings, - OperatorPrecedence.TopLevelOperatorPrecedence, True) - tokens.end_scope(ScopeType.CaseScopeType) + if self.function.is_instruction_collapsed(instr): + tokens.append(InstructionTextToken(InstructionTextTokenType.CollapsedInformationToken, " ...")) + else: + tokens.begin_scope(ScopeType.CaseScopeType) + self.perform_get_expr_text(instr.body, tokens, settings, + OperatorPrecedence.TopLevelOperatorPrecedence, True) + tokens.end_scope(ScopeType.CaseScopeType) elif instr.operation == HighLevelILOperation.HLIL_BREAK: tokens.append(InstructionTextToken(InstructionTextTokenType.KeywordToken, "break")) elif instr.operation == HighLevelILOperation.HLIL_CONTINUE: diff --git a/python/function.py b/python/function.py index ec2fbab8..b101a44f 100644 --- a/python/function.py +++ b/python/function.py @@ -3337,6 +3337,59 @@ class Function: bc.confidence = core.max_confidence core.BNSetUserFunctionInlinedDuringAnalysis(self.handle, bc) + def toggle_region(self, hash): + """ + Toggle the collapsed state of a region during rendering, by hash value + :param hash: Hash value of region + """ + core.BNFunctionToggleRegion(self.handle, hash) + + def collapse_region(self, hash): + """ + Collapse a region during rendering + :param hash: Hash value of region + """ + core.BNFunctionCollapseRegion(self.handle, hash) + + def expand_region(self, hash): + """ + Un-collapse a region during rendering + :param hash: Hash value of region + """ + core.BNFunctionExpandRegion(self.handle, hash) + + def expand_all(self): + """ + Expand all regions in the function + """ + core.BNFunctionExpandAll(self.handle) + + @property + def is_collapsed(self): + """If the entire function is collapsed during rendering.""" + return self.is_region_collapsed(self.start) + + def is_instruction_collapsed( + self, + instr: 'highlevelil.HighLevelILInstruction', + discriminator: int = 0 + ) -> bool: + """ + Determine if a given HLIL instruction (with discriminator) is collapsed during rendering. + :param instr: Instruction which might be collapsed + :param discriminator: Unique discriminator id for the region + :return: True if the instruction should be rendered as collapsed + """ + return self.is_region_collapsed(instr.get_instruction_hash(discriminator)) + + def is_region_collapsed(self, hash) -> bool: + """ + Determine if a given region is collapsed during rendering. + :param hash: Hash value of region + :return: True if the region should be rendered as collapsed + """ + return core.BNFunctionIsRegionCollapsed(self.handle, hash) + class AdvancedFunctionAnalysisDataRequestor: def __init__(self, func: Optional['Function'] = None): diff --git a/python/highlevelil.py b/python/highlevelil.py index 8503bbe4..dc16e333 100644 --- a/python/highlevelil.py +++ b/python/highlevelil.py @@ -943,6 +943,36 @@ class HighLevelILInstruction(BaseILInstruction): finally: core.BNFreeDisassemblyTextLines(lines, count.value) + @property + def can_collapse(self) -> bool: + """If this instruction can be collapsed in rendered lines""" + return self.operation in [ + HighLevelILOperation.HLIL_IF, + HighLevelILOperation.HLIL_WHILE, + HighLevelILOperation.HLIL_WHILE_SSA, + HighLevelILOperation.HLIL_DO_WHILE, + HighLevelILOperation.HLIL_DO_WHILE_SSA, + HighLevelILOperation.HLIL_FOR, + HighLevelILOperation.HLIL_FOR_SSA, + HighLevelILOperation.HLIL_SWITCH, + HighLevelILOperation.HLIL_CASE + ] + + def get_instruction_hash(self, discriminator: int) -> int: + """ + Hash of instruction matching the C++ HighLevelILInstruction::GetInstructionHash, + used for collapsed region matching. + :param discriminator: Extra value to include in the hash to differentiate regions + """ + + def rotl(value, shift): + return ((value << shift) & 0xffffffffffffffff) | (value >> (64 - shift)) + + hash = self.operation.value + hash ^= rotl(self.address, 23) + hash ^= rotl(discriminator, 47) + return hash + @dataclass(frozen=True, repr=False, eq=False) class HighLevelILUnaryBase(HighLevelILInstruction, UnaryOperation): diff --git a/python/languagerepresentation.py b/python/languagerepresentation.py index 51c14da0..44c34bf2 100644 --- a/python/languagerepresentation.py +++ b/python/languagerepresentation.py @@ -35,7 +35,7 @@ from . import variable from . import types from .log import log_error from .enums import BraceRequirement, HighlightStandardColor, InstructionTextTokenType, OperatorPrecedence, ScopeType, \ - SymbolDisplayType, SymbolDisplayResult + SymbolDisplayType, SymbolDisplayResult, InstructionTextTokenContext class HighLevelILTokenEmitter: @@ -92,6 +92,60 @@ class HighLevelILTokenEmitter: """Forces there to be no indentation for the next line.""" core.BNHighLevelILTokenEmitterNoIndentForThisLine(self.handle) + def prepend_blank_collapse_indicator(self): + """ + Insert, at the beginning of the line, a collapse indicator token. + The indicator will be a blank space and not have any functionality. + """ + core.BNHighLevelILTokenPrependCollapseBlankIndicator(self.handle); + + def prepend_instr_collapse_indicator( + self, + function_: 'function.Function', + instr: 'highlevelil.HighLevelILInstruction', + discriminator: int = 0 + ): + """ + Insert, at the beginning of the line, a collapse indicator token. + + The indicator will allow the user to collapse the region specified by + (instr, discriminator) on the function. + Implementations can use :py:func:`Function.is_instruction_collapsed` and + :py:func:`Function.is_region_collapsed` to account for collapsed regions in rendering. + + :param function_: Function whose instructions are being emitted + :param instr: Instruction being emitted which can be collapsed + :param discriminator: Unique discriminator id for the region + """ + if not self.has_collapsable_regions: + return + + # Insert the collapse indicator at the beginning of the line if one isn't already there or the + # one that is there is empty + context = InstructionTextTokenContext.ContentCollapsiblePadding + if instr.can_collapse: + if function_ and function_.is_instruction_collapsed(instr, discriminator): + context = InstructionTextTokenContext.ContentCollapsedContext + else: + context = InstructionTextTokenContext.ContentExpandedContext + self.prepend_region_collapse_indicator(context, instr.get_instruction_hash(discriminator)) + + def prepend_region_collapse_indicator( + self, + context: InstructionTextTokenContext, + hash: int + ): + core.BNHighLevelILTokenPrependCollapseIndicator(self.handle, context, hash) + + @property + def has_collapsable_regions(self) -> bool: + """If the emitter can emit regions which can be collapsed""" + return core.BNHighLevelILTokenEmitterHasCollapsableRegions(self.handle) + + @has_collapsable_regions.setter + def has_collapsable_regions(self, value: bool): + core.BNHighLevelILTokenEmitterSetHasCollapsableRegions(self.handle, value) + class ZeroConfidenceContext: """ ``class ZeroConfidenceContext`` is a context manager that optionally forces tokens to be of zero confidence |