150 lines
5.1 KiB
Python
150 lines
5.1 KiB
Python
import sys
|
|
|
|
from typing import Any, List, Optional, Tuple
|
|
|
|
|
|
def align(val: int) -> int:
|
|
return (val + 3) & 0xFFFFFFFFC
|
|
|
|
|
|
def pad(data: bytes, length: int) -> bytes:
|
|
if len(data) == length:
|
|
return data
|
|
elif len(data) > length:
|
|
raise Exception("Logic error, padding request in data already written!")
|
|
return data + (b"\0" * (length - len(data)))
|
|
|
|
|
|
def descramble_text(text: bytes, obfuscated: bool) -> str:
|
|
if text:
|
|
if obfuscated and (text[0] - 0x20) > 0x7F:
|
|
# Gotta do a weird demangling where we swap the
|
|
# top bit.
|
|
return bytes(((x + 0x80) & 0xFF) for x in text).decode('ascii')
|
|
else:
|
|
return text.decode('ascii')
|
|
else:
|
|
return ""
|
|
|
|
|
|
def scramble_text(text: str, obfuscated: bool) -> bytes:
|
|
if obfuscated:
|
|
return bytes(((x + 0x80) & 0xFF) for x in text.encode('ascii')) + b'\0'
|
|
else:
|
|
return text.encode('ascii') + b'\0'
|
|
|
|
|
|
class TrackedCoverageManager:
|
|
def __init__(self, covered_class: "TrackedCoverage", verbose: bool) -> None:
|
|
self.covered_class = covered_class
|
|
self.verbose = verbose
|
|
|
|
def __enter__(self) -> "TrackedCoverageManager":
|
|
if self.verbose:
|
|
self.covered_class._tracking = True
|
|
return self
|
|
|
|
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
self.covered_class._tracking = False
|
|
|
|
|
|
class TrackedCoverage:
|
|
def __init__(self) -> None:
|
|
self.coverage: List[bool] = []
|
|
self._tracking: bool = False
|
|
|
|
def covered(self, size: int, verbose: bool) -> TrackedCoverageManager:
|
|
if verbose:
|
|
self.coverage = [False] * size
|
|
return TrackedCoverageManager(self, verbose)
|
|
|
|
def add_coverage(self, offset: int, length: int, unique: bool = True) -> None:
|
|
if not self._tracking:
|
|
# Save some CPU cycles if we aren't verbose.
|
|
return
|
|
for i in range(offset, offset + length):
|
|
if self.coverage[i] and unique:
|
|
raise Exception(f"Already covered {hex(offset)}!")
|
|
self.coverage[i] = True
|
|
|
|
def print_coverage(self, req_start: Optional[int] = None, req_end: Optional[int] = None) -> None:
|
|
for start, offset in self.get_uncovered_chunks(req_start, req_end):
|
|
print(f"Uncovered: {hex(start)} - {hex(offset)} ({offset-start} bytes)", file=sys.stderr)
|
|
|
|
def get_uncovered_chunks(self, req_start: Optional[int] = None, req_end: Optional[int] = None, adjust_offsets: bool = False) -> List[Tuple[int, int]]:
|
|
# First offset that is not coverd in a run.
|
|
start = None
|
|
chunks: List[Tuple[int, int]] = []
|
|
|
|
for offset, covered in enumerate(self.coverage):
|
|
if covered:
|
|
if start is not None:
|
|
chunks.append((start, offset))
|
|
start = None
|
|
else:
|
|
if start is None:
|
|
start = offset
|
|
if start is not None:
|
|
# Print final range
|
|
offset = len(self.coverage)
|
|
chunks.append((start, offset))
|
|
|
|
if req_start is None and req_end is None:
|
|
return chunks
|
|
|
|
filtered_chunks: List[Tuple[int, int]] = []
|
|
for start, end in chunks:
|
|
if start >= end:
|
|
raise Exception("Logic error!")
|
|
|
|
if req_start is not None:
|
|
if end <= req_start:
|
|
# Don't care this is wholly before our start filter.
|
|
continue
|
|
if start < req_start and end > req_start:
|
|
# This overlaps our start filter, so update the start to be
|
|
# our start filter.
|
|
start = req_start
|
|
if req_end is not None:
|
|
if start >= req_end:
|
|
# Don't care, this is wholly after our end filter.
|
|
continue
|
|
if start < req_end and end > req_end:
|
|
# This overlaps our end filter, so update the end to be
|
|
# our end filter.
|
|
end = req_end
|
|
|
|
if adjust_offsets:
|
|
filtered_chunks.append((start - req_start if req_start else 0, end - req_start if req_start else 0))
|
|
else:
|
|
filtered_chunks.append((start, end))
|
|
return filtered_chunks
|
|
|
|
|
|
class VerboseOutputManager:
|
|
def __init__(self, covered_class: "VerboseOutput", verbose: bool) -> None:
|
|
self.covered_class = covered_class
|
|
self.verbose = verbose
|
|
|
|
def __enter__(self) -> "VerboseOutputManager":
|
|
if self.verbose:
|
|
self.covered_class.verbose = True
|
|
else:
|
|
self.covered_class.verbose = False
|
|
return self
|
|
|
|
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
self.covered_class.verbose = False
|
|
|
|
|
|
class VerboseOutput:
|
|
def __init__(self) -> None:
|
|
self.verbose: bool = False
|
|
|
|
def debugging(self, verbose: bool) -> VerboseOutputManager:
|
|
return VerboseOutputManager(self, verbose)
|
|
|
|
def vprint(self, *args: Any, **kwargs: Any) -> None:
|
|
if self.verbose:
|
|
print(*args, **kwargs, file=sys.stderr)
|