Replace RTTI information with text from an input string file

rtti.py

import sys, os

def main():
  if len(sys.argv) < 3:
    print('Usage: py rtti.py <path> <txt>')
    return 1
    
  path = sys.argv[1]
  if not os.path.exists(path):
    print('Input file not found')
    return 1
    
  txt = sys.argv[2]
  if not os.path.exists(txt):
    print('String file not found')
    return 1
    
  contents = bytearray(open(path, 'rb').read())
  if len(contents) < 2 or contents[:2] != b'MZ':
    print('Invalid PE file')
    return 1
    
  replacements = open(txt, 'rb').read().replace(b'\x0D\x0A', b'\x20')
  if not replacements:
    print('Invalid string file')
    return 1
  
  queries = [ '.?AV', '.?AU' ]

  for i in range(len(contents) - 4):
    b = contents[i:i+4]
    if b != queries[0].encode() and b != queries[1].encode():
      continue
  
    length = contents.find(b'\x00', i) - i
    text = contents[i:contents.find(b'\x00', i)].decode()
    print(f'[{hex(i)}] {text}')
    
    for j in range(length):
      contents[j+i] = replacements[j % len(replacements)]
  
  with open(path, 'wb') as out:
    out.write(contents)
  
  return 0
    
if __name__ == '__main__':
  exit(main())