mirror of
https://github.com/DragonMinded/bemaniutils.git
synced 2026-10-12 09:18:17 +03:00
Beginning of an AFP ByteCode decompiler, starting with a massive code reorg and a control flow graph analyzer.
This commit is contained in:
Jennifer Taylor
@@ -0,0 +1,282 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user