#Global / AlliedSignal GNS-X CDU - Inline string processor
#@author Phil Pemberton
#@category 00-Projects.GNS-X
#@keybinding 
#@menupath 
#@toolbar 

from ghidra.program.model.listing import CodeUnit
from ghidra.program.model.symbol  import FlowType

###############

# helper function to get a Ghidra Address type
def getAddress(offset):
    return currentProgram.getAddressFactory().getDefaultAddressSpace().getAddress(offset)

###############

# setup
fm = currentProgram.getFunctionManager()
mem = currentProgram.getMemory()

# Find methods which call inline print functions
xrefs = []
funcs = fm.getFunctions(True)
for func in funcs:
  if func.getName().startswith('DisplayScreenAfterCall'):
    print("Found inline print function {} @ 0x{}".format(func.getName(), func.getEntryPoint()))
    entry_point = func.getEntryPoint()
    references = getReferencesTo(entry_point)
    for xref in references:
      xrefs.append(xref)

# Work on each xref...

fixlist = []
for xref in xrefs:
  if xref.getReferenceType() != FlowType.UNCONDITIONAL_CALL:
    print("*** Warning, requires manual fixup: " + str(xref))
    createBookmark(xref.getFromAddress(), "InlineStringError", "Manual fixup required, xref source is not a CALL")
    continue

  # Start from the xref address. Seek past the CALL instruction.
  # Convert byte to data. (charactor position byte)
  # Scan bytes until 0xFF (terminator)
  # Convert text to string

  # Get the xref's from address (the call)
  addr = xref.getFromAddress()

  # Get an Instruction to the call instruction
  callIns = getInstructionAt(addr)   # the call address

  # DEBUG
  print("xref - {} => {}".format(str(xref), callIns.getMnemonicString()))
  # DEBUG

  # First byte after the call is the Offset Address. Mark it as data.
  offsetByteAddr = addr.add(callIns.getLength())
  clearListing(offsetByteAddr)
  offsetByteData = createByte(offsetByteAddr)
  offsetByteData.setComment(CodeUnit.EOL_COMMENT, "Byte offset from start of VRAM")

  # Advance past the offset byte
  stringStartAddr = offsetByteAddr.next()

  # Scan the input string
  stringEndAddr = stringStartAddr
  sbyte = getByte(stringEndAddr) & 0xFF    # AND 0xFF is to force the byte to be unsigned
  text = ""
  while sbyte != 0xFF:
    text += chr(sbyte)
    stringEndAddr = stringEndAddr.next()
    sbyte = getByte(stringEndAddr) & 0xFF    # AND 0xFF is to force the byte to be unsigned

  # calculate string length
  stringLen = stringEndAddr.getOffset() - stringStartAddr.getOffset()

  print("  String starts at: {}, ends at {}, text '{}' ({} characters)".format(stringStartAddr, stringEndAddr, text, stringLen))

  # String start and end address got - end address points at the terminator.

  # Wipe out anything already there and replace it with a string
  clearListing(stringStartAddr, stringEndAddr)
  createAsciiString(stringStartAddr, stringLen)

  # Turn the terminator into a byte
  strTermData = createByte(stringEndAddr)
  strTermData.setComment(CodeUnit.EOL_COMMENT, "Terminator")

  # Set a fallthrough on the call instruction
  callIns.setFallThrough(stringEndAddr.next())

  # Set a bookmark on the end of the string (in case manual cleanup is needed)
  createBookmark(stringEndAddr.next(), "InlineString", "End of inline string, disassembly cleanup may be required")

  # Start disassembly at end of string
  disassemble(stringEndAddr.next())