#/usr/bin/python
from pydbg import *
from pydbg.defines import *
import os.path
import pefile
import re
import struct
import sys

def GetSection(pe, section_name):
  for section in pe.sections:
    # Section names are \0-padded on the right
    if section.Name.rstrip("\x00") == section_name:
      return section
  
  raise Exception("No %s section" % section_name)

# Find RenderProcessHost::run_renderer_in_process_ by
# BrowserRenderProcessHost::BrowserRenderProcessHost
# (chrome/browser/renderer_host/browser_render_process_host.cc:146)
def FindRripByBrph(pe):
  section = GetSection(pe, ".text")
  data = section.get_data(section.VirtualAddress)
  
  # chrome.dll 2.0.172.33 10,031,088 bytes
  # .text:01DB921E 80 3D 2E 60 4C 02 00    cmp     RenderProcessHost__run_renderer_in_process_, 0
  # .text:01DB9225 74 14                   jz      short loc_1DB923B
  # .text:01DB9227 83 05 0C 5D 4B 02 03    add     dword_24B5D0C, 3
  RE_BRPH = re.compile(r"\x80\x3d(....)\x00..\x83\x05....\x03", re.S)
  matches = list(RE_BRPH.finditer(data))
  if len(matches) != 1:
    raise Exception("Wrong number of instruction pattern matches")
  
  return struct.unpack("<L", matches[0].group(1))[0]

# Find kMaxRenderersByRamTier
# (chrome/browser/renderer_host/render_process_host.cc:25)
def FindMrbrt(pe):
  section = GetSection(pe, ".rdata")
  data = section.get_data(section.VirtualAddress)
  
  # chrome.dll 2.0.172.33 10,031,088 bytes
  # .rdata:023FD460 03 00 00 00 06 00 00 00+GetMaxRendererProcessCount__kMaxRenderersByRamTier dd    3,   6,   9, 0Ch, 0Eh, 12h, 14h, 16h, 18h, 1Ah, 1Dh
  kMaxRenderersByRamTier = [3, 6, 9, 12, 14, 18, 20, 22, 24, 26, 29, 32, 35, 38, 40]
  mrbrt_data = struct.pack("<" + "I" * len(kMaxRenderersByRamTier),
    *kMaxRenderersByRamTier)
  
  # FIXME: handle relocations
  # FIXME: ensure it's the only match
  return (pe.OPTIONAL_HEADER.ImageBase + section.VirtualAddress +
          data.index(mrbrt_data)) # raises if not found

# Find RenderProcessHost::run_renderer_in_process_ by
# GetMaxRendererProcessCount()
# (chrome/browser/renderer_host/render_process_host.cc:15)
def FindRripByGmrpc(pe):
  mrbrt_va = FindMrbrt(pe)
  
  section = GetSection(pe, ".text")
  data = section.get_data(section.VirtualAddress)
  
  # chrome.dll 2.0.172.33 10,031,088 bytes
  # .text:01CDEF06                         GetMaxRendererProcessCount proc near    ; CODE XREF: sub_1D7D325+Cp
  # .text:01CDEF06 80 3D 2E 60 4C 02 00    cmp     RenderProcessHost__run_renderer_in_process_, 0
  # .text:01CDEF0D 56                      push    esi
  # .text:01CDEF0E 8B 35 EC 3F 4B 02       mov     esi, dword_24B3FEC
  # ...
  # .text:01CDEF3A 8B 04 85 60 D4 3F 02    mov     eax, ds:GetMaxRendererProcessCount__kMaxRenderersByRamTier[eax*4]
  ref_mrbrt = data.index(struct.pack("<L", mrbrt_va)) # raises if not found
  
  # Account for 50% function size increase in subsequent versions
  projected_function_size = int(1.5 * (0x01CDEF3A - 0x01CDEF06))
  
  func_data = data[ref_mrbrt - projected_function_size:ref_mrbrt]
  RE_GMRPC = re.compile(r"\x80\x3d(....)\x00", re.S)
  matches = list(RE_GMRPC.finditer(func_data))
  if len(matches) != 1:
    raise Exception("Wrong number of instruction pattern matches")
  
  match = matches[0]
  return struct.unpack("<L", match.group(1))[0]

def MultiFind(name, finders):
  def MultiFinder(pe):
    value = None
    for finder in finders:
      try:
        candidate = finder(pe)
      except Exception, e:
        print "%s: %s" % (finder.__name__, str(e))
      else:
        print "%s: 0x%08x" % (finder.__name__, candidate)
        
        if value is None:
          value = candidate
        elif candidate != value:
          raise Exception("Candidate value 0x%08x disagrees with "
                          "first value 0x%08x" % (candidate, value))
    
    if value is None:
      raise Exception("No candidates")
    
    return value
  return MultiFinder

FindRrip = MultiFind("FindRrip", [FindRripByBrph, FindRripByGmrpc])

def FindChromeDll(path):
  for drive in "CDEF":
    fullpath = drive + ":" + path
    if os.path.exists(fullpath):
      return fullpath
  
  raise Exception("chrome.dll not found")

def handle_load_dll(dbg):
  last_dll = dbg.get_system_dll(-1)
  print "loading:%s into:%08x size:%d (from %s)" % (last_dll.name, last_dll.base, last_dll.size, last_dll.path)
  
  if last_dll.name == "chrome.dll":
    pe = pefile.PE(FindChromeDll(last_dll.path))

    rrip = FindRrip(pe) - pe.OPTIONAL_HEADER.ImageBase + last_dll.base
    print ("RenderProcessHost::run_renderer_in_process_ is (re)located "
           "at 0x%08x" % rrip)
    
    # It all comes down to this.
    dbg.write(rrip, "\x01", 1)
    
    # bai.
    dbg.detach()
  
  return DBG_CONTINUE

def usage():
  print r'Usage:'
  print r'  chromehack.py "%USERPROFILE%\...\chrome.exe" [--single-process]'

def main(args):
  if struct.calcsize("<L") != 4:
    raise Exception("This only works on Win32")
  
  if len(args) < 1:
    usage()
    return 1

  chrome_exe = args[0]
  chrome_args = []
  single_process = False
  for arg in args[1:]:
    if arg == "--single-process":
      single_process = True
      # swallow switch
    else:
      chrome_args.append(arg)

  dbg = pydbg()
  if single_process:
    dbg.set_callback(LOAD_DLL_DEBUG_EVENT, handle_load_dll)
  # FIXME: handle spaces in arguments
  dbg.load(chrome_exe, " ".join(chrome_args))
  if single_process:
    dbg.run()
  else:
    dbg.detach()
  
  # FIXME return Chrome's exit code?
  return 0

if __name__ == "__main__":
  sys.exit(main(sys.argv[1:]))