import os from typing import Any, Dict, List, Tuple, cast from .types import AP2Action, JumpAction, IfAction from .util import VerboseOutput class ByteCode: # A list of bytecodes to execute. def __init__(self, actions: List[AP2Action], end_offset: int) -> None: self.actions = actions self.end_offset = end_offset def as_dict(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: if kwargs.get('decompile_bytecode', False): decompiler = ByteCodeDecompiler(self) code = decompiler.decompile(verbose=True) return { 'code': code, } else: return { 'actions': [a.as_dict(*args, **kwargs) for a in self.actions], 'end_offset': self.end_offset, } def __repr__(self) -> str: entries: List[str] = [] for action in self.actions: entries.extend([f" {s}" for s in str(action).split(os.linesep)]) return f"ByteCode({os.linesep}{os.linesep.join(entries)}{os.linesep} {self.end_offset}: END{os.linesep})" class ControlFlow: def __init__(self, beginning: int, end: int, next_flow: List[int]) -> None: self.beginning = beginning self.end = end self.next_flow = next_flow def contains(self, offset: int) -> bool: return (self.beginning <= offset) and (offset < self.end) def is_first(self, offset: int) -> bool: return self.beginning == offset def is_last(self, offset: int) -> bool: return self.end == (offset + 1) def split(self, offset: int, link: bool = False) -> Tuple["ControlFlow", "ControlFlow"]: if not self.contains(offset): raise Exception(f"This ControlFlow does not contain offset {offset}") # First, make the second half that the first half will point to. second = ControlFlow( offset, self.end, self.next_flow, ) # Now, make the first half that we can point to. first = ControlFlow( self.beginning, offset, [second.beginning] if link else [], ) return (first, second) def __repr__(self) -> str: return f"ControlFlow(beginning={self.beginning}, end={self.end}, next={(', '.join(str(n) for n in self.next_flow)) or 'N/A'}" class ByteCodeChunk: def __init__(self, actions: List[AP2Action], next_chunk: List[int]) -> None: self.actions = actions self.next_chunk = next_chunk @property def offset(self) -> int: return self.actions[0].offset def __repr__(self) -> str: entries: List[str] = [] for action in self.actions: entries.extend([f" {s}" for s in str(action).split(os.linesep)]) return f"ByteCodeChunk({os.linesep}{os.linesep.join(entries)}{os.linesep} Next Offsets: {', '.join(str(n) for n in self.next_chunk) or 'None'}{os.linesep})" class ByteCodeDecompiler(VerboseOutput): def __init__(self, bytecode: ByteCode) -> None: super().__init__() self.bytecode = bytecode def __graph_control_flow(self) -> List[ByteCodeChunk]: # Start by assuming that the whole bytecode never directs flow. This is, confusingly, # indexed by AP2Action offset, not by actual bytecode offset, so we can avoid the # prickly problem of opcodes that take more than one byte in the data. flows: Dict[int, ControlFlow] = {} end = len(self.bytecode.actions) beginning = 0 # The end of the program. flows[end] = ControlFlow(end, end + 1, []) # The rest of the program. flows[beginning] = ControlFlow(beginning, end, [end]) # Function that helps us find a flow by position. def find(opcodeno: int) -> int: for start, cf in flows.items(): if cf.contains(opcodeno): return start raise Exception(f"Offset {opcodeno} somehow not in our control flow graph!") # Now, walk the entire bytecode, and every control flow point split the graph at that point. for i, action in enumerate(self.bytecode.actions): current_action = i next_action = i + 1 if action.opcode in [AP2Action.THROW, AP2Action.RETURN]: # This should end execution, so we should cap off the current execution # and send it to the end. current_action_flow = find(current_action) next_action_flow = find(next_action) if current_action_flow == next_action_flow: # We need to split this on the next_action boundary. first, second = flows[current_action_flow].split(next_action) first.next_flow = [end] self.vprint(f"{action} action split {flows[current_action_flow]} into {first}, {second}") flows[current_action_flow] = first flows[next_action] = second else: # This already was split in two, presumably by something # earlier in the chain jumping to the opcode after this. # We need to unlink the current flow from the second and # link it to the end. flows[current_action_flow].next_flow = [end] self.vprint(f"{action} action repointed {flows[current_action_flow]} to end") elif action.opcode == AP2Action.JUMP: # Unconditional control flow redirection after this, we should split the # section if necessary and point this section at the new offset. # First, we need to find the jump point and make sure that its the start # of a section. action = cast(JumpAction, action) for j, dest in enumerate(self.bytecode.actions): if dest.offset == action.jump_offset: dest_action = j break else: raise Exception(f"{action} jumps to an opcode that doesn't exist!") # If the destination action flow already starts with the jump offset, # then we're good, we just need to point our current split at this new # offset. If it doesn't start with the jump offset, then we need to split # that flow so we can point to the opcode directly. dest_action_flow = find(dest_action) if not flows[dest_action_flow].is_first(dest_action): first, second = flows[dest_action_flow].split(dest_action, link=True) self.vprint(f"{action} action required split of {flows[dest_action_flow]} into {first, second}") flows[dest_action_flow] = first flows[dest_action] = second # Now, the second is what we want to point at in the next section. dest_action_flow = dest_action # Now, we must split the current flow at the point of this jump. current_action_flow = find(current_action) next_action_flow = find(next_action) if current_action_flow == next_action_flow: # We need to split this on the next_action boundary. first, second = flows[current_action_flow].split(next_action) first.next_flow = [dest_action_flow] self.vprint(f"{action} action split {flows[current_action_flow]} into {first}, {second}") flows[current_action_flow] = first flows[next_action] = second else: # This already was split in two, presumably by something # earlier in the chain jumping to the opcode after this. # We need to unlink the current flow from the second and # link it to the end. flows[current_action_flow].next_flow = [dest_action_flow] self.vprint(f"{action} action repointed {flows[current_action_flow]} to new chunk") elif action.opcode in [AP2Action.IF, AP2Action.IF2]: # Conditional control flow redirection after this, we should split the # section if necessary and point this section at the new offset as well # as the second half of the split section. # First, we need to find the jump point and make sure that its the start # of a section. action = cast(IfAction, action) for j, dest in enumerate(self.bytecode.actions): if dest.offset == action.jump_if_true_offset: dest_action = j break else: raise Exception(f"{action} conditional jumps to an opcode that doesn't exist!") # If the destination action flow already starts with the jump offset, # then we're good, we just need to point our current split at this new # offset. If it doesn't start with the jump offset, then we need to split # that flow so we can point to the opcode directly. dest_action_flow = find(dest_action) if not flows[dest_action_flow].is_first(dest_action): first, second = flows[dest_action_flow].split(dest_action, link=True) self.vprint(f"{action} action required split of {flows[dest_action_flow]} into {first, second}") flows[dest_action_flow] = first flows[dest_action] = second # Now, the second is what we want to point at in the next section. dest_action_flow = dest_action # Now, we must split the current flow at the point of this jump. current_action_flow = find(current_action) next_action_flow = find(next_action) if current_action_flow == next_action_flow: # We need to split this on the next_action boundary. first, second = flows[current_action_flow].split(next_action) first.next_flow = [next_action, dest_action_flow] self.vprint(f"{action} action split {flows[current_action_flow]} into {first}, {second}") flows[current_action_flow] = first flows[next_action] = second else: # This already was split in two, presumably by something # earlier in the chain jumping to the opcode after this. # We need to unlink the current flow from the second and # link it to the end. flows[current_action_flow].next_flow = [next_action, dest_action_flow] self.vprint(f"{action} action repointed {flows[current_action_flow]} to new chunk") # Finally, return chunks of contiguous execution. chunks: List[ByteCodeChunk] = [] for start, flow in flows.items(): if start == end: # We don't want to render out the end of the graph, it was only there to make # the above algorithm easier. continue if len(flow.next_flow) == 1 and flow.next_flow[0] == end: # This flow is a termination state. chunks.append(ByteCodeChunk(self.bytecode.actions[flow.beginning:flow.end], [])) else: next_chunks: List[int] = [] for ano in flow.next_flow: if ano == end: raise Exception("Logic error!") next_chunks.append(self.bytecode.actions[ano].offset) chunks.append(ByteCodeChunk(self.bytecode.actions[flow.beginning:flow.end], next_chunks)) return sorted(chunks, key=lambda c: c.offset) def decompile(self, verbose: bool = False) -> str: with self.debugging(verbose): return self.__decompile() def __decompile(self) -> str: # First, we need to construct a control flow graph. chunks = self.__graph_control_flow() self.vprint(chunks) return "TODO"