#
# ios-12-kernelcache-tagged-pointers.py
# Brandon Azad
#
# An idapython script that shows how to work with the new tagged pointers in the iOS 12
# kernelcache.
#

import idc
import idaapi
import idautils

import random
from collections import defaultdict, Counter

def untag_pointer(tp):
    return tp | 0xffff000000000000

def tagged_pointer_tag(tp):
    return (tp >> 48) & 0xffff

def is_tagged_pointer_format(pointer):
    return (pointer & 0x0000FFFFf0000000) == 0x0000FFF000000000

def is_valid_tagged_pointer(pointer):
    return is_tagged_pointer_format(pointer) and idaapi.getseg(untag_pointer(pointer))

def tagged_pointer_tag_to_span(tag):
    return (tag >> 1) & ~0x3

def tagged_pointer_span(tp):
    return tagged_pointer_tag_to_span(tagged_pointer_tag(tp))

def tagged_pointers():
    if tagged_pointers.eas_and_tps is None:
        eas_and_tps = []
        for seg in idautils.Segments():
            start = idc.SegStart(seg)
            end = idc.SegEnd(seg)
            for ea in range(start, end, 4):
                value = idc.Qword(ea)
                if is_valid_tagged_pointer(value):
                    eas_and_tps.append((ea, value))
        tagged_pointers.eas_and_tps = eas_and_tps
    return tagged_pointers

tagged_pointers.eas_and_tps = None

def print_tagged_pointer_counts_per_section():
    seg_to_tag_count = Counter()
    for ea, tp in tagged_pointers().eas_and_tps:
        if tagged_pointer_tag(tp) == 0xffff:
            continue
        seg = idc.SegStart(ea)
        seg_to_tag_count[seg] += 1
    segs = list(seg_to_tag_count.keys())
    segs.sort()
    for seg in segs:
        print '{:32s} {:8d}'.format(idc.SegName(seg), seg_to_tag_count[seg])

def print_untagged_pointer_counts_per_section():
    seg_to_tag_count = Counter()
    for ea, tp in tagged_pointers().eas_and_tps:
        if tagged_pointer_tag(tp) != 0xffff:
            continue
        seg = idc.SegStart(ea)
        seg_to_tag_count[seg] += 1
    segs = list(seg_to_tag_count.keys())
    segs.sort()
    for seg in segs:
        print '{:32s} {:8d}'.format(idc.SegName(seg), seg_to_tag_count[seg])

def print_tagged_pointer_counts_by_tag_per_section():
    seg_to_tag_counts = defaultdict(Counter)
    for ea, tp in tagged_pointers().eas_and_tps:
        seg = idc.SegStart(ea)
        tag = tagged_pointer_tag(tp)
        seg_to_tag_counts[seg][tag] += 1
    segs = list(seg_to_tag_counts.keys())
    segs.sort()
    for seg in segs:
        tag_counts = seg_to_tag_counts[seg]
        tags = list(tag_counts.keys())
        tags.sort()
        print '{:32s} {}'.format(idc.SegName(seg),
                ', '.join('{:04x} ({})'.format(tag, tag_counts[tag]) for tag in tags))

def print_references_for_tagged_pointers():
    tag_to_eas_and_tps = defaultdict(list)
    for ea, tp in tagged_pointers().eas_and_tps:
        tag_to_eas_and_tps[tagged_pointer_tag(tp)].append((ea, tp))
    tags = list(tag_to_eas_and_tps.keys())
    tags.sort()
    for tag in tags:
        eas_and_tps = tag_to_eas_and_tps[tag]
        ea, tp = random.choice(eas_and_tps)
        kp = untag_pointer(tp)
        print '{:04x}  {:8d}    {:016x} {:32s} ->  {:016x} {:s}'.format(
                tag, len(eas_and_tps), ea, idc.SegName(ea), kp, idc.SegName(kp))

def find_tagged_pointers_to_address(address):
    tag_to_eas = defaultdict(list)
    for ea, tp in tagged_pointers().eas_and_tps:
        if untag_pointer(tp) == address:
            tag_to_eas[tagged_pointer_tag(tp)].append(ea)
    return dict(tag_to_eas)

def print_tagged_pointers_to_address(address):
    tag_to_eas = find_tagged_pointers_to_address(address)
    tags = list(tag_to_eas.keys())
    tags.sort()
    for tag in tags:
        eas = tag_to_eas[tag]
        print '{:04x}  {:8d}    {}'.format(
                tag, len(eas), ', '.join('{:016x}'.format(ea) for ea in eas))

def check_pointer_tags_as_links():
    last_seg, last_ea, last_tp = None, None, None
    correct, warnings, errors, skipped = 0, 0, 0, 0
    for ea, tp in tagged_pointers().eas_and_tps:
        if tagged_pointer_tag(tp) == 0xffff:
            skipped += 1
            continue
        seg = idc.SegStart(ea)
        if seg == last_seg:
            span = tagged_pointer_span(last_tp)
            if span == 0:
                real_span = ea - last_ea
                print 'WARNING: ZERO SPAN LINK: ea={:016x}, tp={:016x}, last_ea={:016x}, last_tp={:016x}, span={:04x}'.format(
                        ea, tp, last_ea, last_tp, real_span)
                warnings += 1
            elif ea != last_ea + span:
                print 'ERROR: BAD LINK: ea={:016x}, tp={:016x}, last_ea={:016x}, last_tp={:016x}'.format(
                        ea, tp, last_ea, last_tp)
                errors += 1
                break
            else:
                correct += 1
        elif last_ea:
            span = tagged_pointer_span(last_tp)
            if span == 0:
                correct += 1
            else:
                print 'WARNING: NON-ZERO SPAN LINK AT END OF SEGMENT: ea={:016x}, tp={:016x}, last_ea={:016x}, last_tp={:016x}, seg={:016x}'.format(
                        ea, tp, last_ea, last_tp, seg)
                warnings += 1
        last_seg, last_ea, last_tp = seg, ea, tp
    print '{} correct, {} warnings, {} errors, {} skipped'.format(
            correct, warnings, errors, skipped)