diff --git a/bemani/common/__init__.py b/bemani/common/__init__.py index 0f8f9ac..85570c0 100644 --- a/bemani/common/__init__.py +++ b/bemani/common/__init__.py @@ -7,6 +7,7 @@ from bemani.common.id import ID from bemani.common.aes import AESCipher from bemani.common.time import Time from bemani.common.parallel import Parallel +from bemani.common.pe import PEFile __all__ = [ @@ -24,4 +25,5 @@ __all__ = [ "Time", "Parallel", "intish", + "PEFile", ] diff --git a/bemani/common/pe.py b/bemani/common/pe.py new file mode 100644 index 0000000..9022db2 --- /dev/null +++ b/bemani/common/pe.py @@ -0,0 +1,22 @@ +import pefile # type: ignore + + +class PEFile: + def __init__(self, data: bytes) -> None: + self.__pe = pefile.PE(data=data, fast_load=True) + + def virtual_to_physical(self, offset: int) -> int: + for section in self.__pe.sections: + start = section.VirtualAddress + self.__pe.OPTIONAL_HEADER.ImageBase + end = start + section.SizeOfRawData + + if offset >= start and offset < end: + return (offset - start) + section.PointerToRawData + + raise Exception(f"Couldn't find physical offset for virtual offset 0x{offset:08x}") + + def is_virtual(self, offset: int) -> bool: + return offset >= self.__pe.OPTIONAL_HEADER.ImageBase + + def is_64bit(self) -> bool: + return hex(self.__pe.FILE_HEADER.Machine) == '0x8664' diff --git a/bemani/utils/psmap.py b/bemani/utils/psmap.py index b159dc7..11226d3 100644 --- a/bemani/utils/psmap.py +++ b/bemani/utils/psmap.py @@ -1,33 +1,24 @@ import argparse -import pefile # type: ignore import struct from typing import List +from bemani.common import PEFile from bemani.protocol import Node from bemani.utils.responsegen import generate_lines def parse_psmap(data: bytes, offset: str, rootname: str) -> Node: - pe = pefile.PE(data=data, fast_load=True) + pe = PEFile(data=data) root = Node.void(rootname) base = int(offset, 16) - def virtual_to_physical(offset: int) -> int: - for section in pe.sections: - start = section.VirtualAddress + pe.OPTIONAL_HEADER.ImageBase - end = start + section.SizeOfRawData - - if offset >= start and offset < end: - return (offset - start) + section.PointerToRawData - raise Exception(f'Couldn\'t find raw offset for virtual offset 0x{offset:08x}') - - if base >= pe.OPTIONAL_HEADER.ImageBase: + if pe.is_virtual(base): # Assume this is virtual - base = virtual_to_physical(base) + base = pe.virtual_to_physical(base) def read_string(offset: int) -> str: # First, translate load offset in memory to disk offset - offset = virtual_to_physical(offset) + offset = pe.virtual_to_physical(offset) # Now, grab bytes until we're null-terminated bytestring = [] @@ -43,7 +34,7 @@ def parse_psmap(data: bytes, offset: str, rootname: str) -> Node: saved_loc: List[int] = [] while True: - if hex(pe.FILE_HEADER.Machine) == '0x8664': # 64 bit + if pe.is_64bit(): # 64 bit chunk = data[base:(base + 24)] base = base + 24 @@ -83,7 +74,7 @@ def parse_psmap(data: bytes, offset: str, rootname: str) -> Node: if defaultptr != 0: saved_loc.append(base) - base = virtual_to_physical(defaultptr) + base = pe.virtual_to_physical(defaultptr) else: saved_loc.append(None) diff --git a/bemani/utils/read.py b/bemani/utils/read.py index 4571afd..d8624dc 100644 --- a/bemani/utils/read.py +++ b/bemani/utils/read.py @@ -7,7 +7,6 @@ import io import jaconv # type: ignore import json import os -import pefile # type: ignore import struct import yaml import xml.etree.ElementTree as ET @@ -18,7 +17,7 @@ from sqlalchemy.sql import text # type: ignore from sqlalchemy.exc import IntegrityError # type: ignore from typing import Any, Dict, List, Optional, Tuple -from bemani.common import GameConstants, VersionConstants, DBConstants, Time +from bemani.common import GameConstants, VersionConstants, DBConstants, PEFile, Time from bemani.format import ARC, IFS, IIDXChart, IIDXMusicDB from bemani.data import Server, Song from bemani.data.interfaces import APIProviderInterface @@ -394,16 +393,7 @@ class ImportPopn(ImportBase): data = myfile.read() myfile.close() - pe = pefile.PE(data=data, fast_load=True) - - def virtual_to_physical(offset: int) -> int: - for section in pe.sections: - start = section.VirtualAddress + pe.OPTIONAL_HEADER.ImageBase - end = start + section.SizeOfRawData - - if offset >= start and offset < end: - return (offset - start) + section.PointerToRawData - raise Exception(f'Couldn\'t find raw offset for virtual offset 0x{offset:08x}') + pe = PEFile(data) if self.version == VersionConstants.POPN_MUSIC_TUNE_STREET: # Based on K39:J:A:A:2010122200 @@ -984,7 +974,7 @@ class ImportPopn(ImportBase): def read_string(offset: int) -> str: # First, translate load offset in memory to disk offset - offset = virtual_to_physical(offset) + offset = pe.virtual_to_physical(offset) # Now, grab bytes until we're null-terminated bytestring = [] @@ -1650,19 +1640,10 @@ class ImportIIDX(ImportBase): import_qpros = True # by default, try to import qpros try: - pe = pefile.PE(data=binarydata, fast_load=True) + pe = PEFile(binarydata) except BaseException: import_qpros = False # if it failed then we're reading a music db file, not the executable - def virtual_to_physical(offset: int) -> int: - for section in pe.sections: - start = section.VirtualAddress + pe.OPTIONAL_HEADER.ImageBase - end = start + section.SizeOfRawData - - if offset >= start and offset < end: - return (offset - start) + section.PointerToRawData - raise Exception(f'Couldn\'t find raw offset for virtual offset 0x{offset:08x}') - songs: List[Dict[str, Any]] = [] if not import_qpros: musicdb = IIDXMusicDB(binarydata) @@ -1859,7 +1840,7 @@ class ImportIIDX(ImportBase): def read_string(offset: int) -> str: # First, translate load offset in memory to disk offset - offset = virtual_to_physical(offset) + offset = pe.virtual_to_physical(offset) # Now, grab bytes until we're null-terminated bytestring = [] @@ -2759,16 +2740,7 @@ class ImportSDVX(ImportBase): data = myfile.read() myfile.close() - pe = pefile.PE(data=data, fast_load=True) - - def virtual_to_physical(offset: int) -> int: - for section in pe.sections: - start = section.VirtualAddress + pe.OPTIONAL_HEADER.ImageBase - end = start + section.SizeOfRawData - - if offset >= start and offset < end: - return (offset - start) + section.PointerToRawData - raise Exception(f'Couldn\'t find raw offset for virtual offset 0x{offset:08x}') + pe = PEFile(data) if self.version == VersionConstants.SDVX_BOOTH: offset = 0xFFF28 @@ -2779,7 +2751,7 @@ class ImportSDVX(ImportBase): def read_string(spot: int) -> str: # First, translate load offset in memory to disk offset - spot = virtual_to_physical(spot) + spot = pe.virtual_to_physical(spot) # Now, grab bytes until we're null-terminated bytestring = [] diff --git a/bemani/utils/struct.py b/bemani/utils/struct.py index efe0293..f8399a2 100644 --- a/bemani/utils/struct.py +++ b/bemani/utils/struct.py @@ -1,9 +1,10 @@ import argparse -import pefile # type: ignore import struct import sys from typing import Optional, Tuple, List, Any +from bemani.common import PEFile + class LineNumber: def __init__(self, offset: int, hex: bool) -> None: @@ -20,7 +21,7 @@ class LineNumber: class StructPrinter: def __init__(self, data: bytes) -> None: self.data = data - self.pe = pefile.PE(data=data, fast_load=True) + self.pe = PEFile(data) def parse_format_spec(self, fmt: str) -> Tuple[str, List[Any]]: prefix: str = "" @@ -106,15 +107,6 @@ class StructPrinter: return prefix, specs - def virtual_to_physical(self, offset: int) -> int: - for section in self.pe.sections: - start = section.VirtualAddress + self.pe.OPTIONAL_HEADER.ImageBase - end = start + section.SizeOfRawData - - if offset >= start and offset < end: - return (offset - start) + section.PointerToRawData - raise Exception(f'Couldn\'t find raw offset for virtual offset 0x{offset:08x}') - def parse_struct(self, startaddr: str, endaddr: str, countstr: str, fmt: str) -> List[Any]: start: int = int(startaddr, 16) end: Optional[int] = int(endaddr, 16) if endaddr is not None else None @@ -125,13 +117,13 @@ class StructPrinter: if end is not None and count is not None: raise Exception("Can't handle providing two ends!") - if start >= self.pe.OPTIONAL_HEADER.ImageBase: + if self.pe.is_virtual(start): # Assume this is virtual - start = self.virtual_to_physical(start) + start = self.pe.virtual_to_physical(start) - if end is not None and end >= self.pe.OPTIONAL_HEADER.ImageBase: + if end is not None and self.pe.is_virtual(end): # Assume this is virtual - end = self.virtual_to_physical(end) + end = self.pe.virtual_to_physical(end) # Parse out any dereference instructions. prefix, specs = self.parse_format_spec(fmt) @@ -191,9 +183,14 @@ class StructPrinter: line.append(struct.unpack(prefix + spec, chunk)[0]) offset += size else: - chunk = self.data[offset:(offset + 4)] - pointer = struct.unpack(prefix + "I", chunk)[0] - offset += 4 + if self.pe.is_64bit(): + chunk = self.data[offset:(offset + 8)] + pointer = struct.unpack(prefix + "Q", chunk)[0] + offset += 8 + else: + chunk = self.data[offset:(offset + 4)] + pointer = struct.unpack(prefix + "I", chunk)[0] + offset += 4 # Resolve the physical address of this pointer, trick the substructure into # parsing only one iteration. @@ -201,7 +198,7 @@ class StructPrinter: # Null pointer line.append(None) else: - pointer = self.virtual_to_physical(pointer) + pointer = self.pe.virtual_to_physical(pointer) subparse = self.__parse_struct(pointer, pointer + 1, None, prefix, spec) if len(subparse) != 1: raise Exception("Logic error!")