Cool_Tools/bemaniutils
1
0
Fork 0
mirror of synced 2024-11-12 01:00:46 +01:00
Code Releases Activity

Move PE operations into common library, fix 64-bit struct pointer dereferencing.

Browse Source
This commit is contained in:
Jennifer Taylor 2021-08-15 00:41:37 +00:00
parent 40614c106f
commit a6da39e469
5 changed files with 54 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
bemani/common/__init__.py
View File

@ -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",
]

22
bemani/common/pe.py Normal file
View File

@ -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'

23
bemani/utils/psmap.py
View File

@ -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)

42
bemani/utils/read.py
View File

@ -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 = []

35
bemani/utils/struct.py
View File

@ -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!")