get_classes.py

from ghidra.program.model.address import Address
from ghidra.program.model.mem import MemoryAccessException
from ghidra.program.flatapi import FlatProgramAPI
from ghidra.util.task import TaskMonitor

import math
import json
from os.path import isfile, join, dirname
import struct as st
import string
import sys

cp = currentProgram
fp = FlatProgramAPI(cp)
mem = cp.memory

word_size = cp.getDefaultPointerSize()
word_mask = (1 << (word_size * 8)) - 1
is_big_endian = cp.getLanguage().isBigEndian()

endy_str = ['<', '>'][is_big_endian]
word_str = ['B', 'H', 'I', 'Q'][int(math.log(word_size, 2))]
word_fmt = endy_str + word_str


def read_string(addr, maxlen=200):
    s = ''
    ptr = addr

    while len(s) < maxlen:
        char = bytes(bytearray([getByte(ptr)]))

        if char == '\x00' or char not in string.printable:
            return s

        s += char
        ptr = ptr.add(1)

    return s

def get_class_info(class_ptr):
    #print(class_ptr)
    #data_ptr = r.cmdj('pxwj 4 @ %d' % (class_ptr+0x10))[0]
    data_ptr = toAddr(st.unpack(word_fmt, getBytes(class_ptr.add(4 * word_size), word_size))[0])

    if data_ptr.offset % 2 != 0:
        data_ptr = data_ptr.subtract(1)

    #print(data_ptr)
    #name_ptr = r.cmdj('pxwj 5 @ %d' % (data_ptr+0x10))[0]
    name_ptr = toAddr(st.unpack(word_fmt, getBytes(data_ptr.add(16 + word_size), word_size))[0])
    #print(name_ptr)
    #classname = r.cmd('ps @ %d' % name_ptr)
    classname = read_string(name_ptr)

    if classname is None:
        return None

    classname = classname.strip('\n')
    print(classname)

    #methodlist_ptr = r.cmdj('pxwj 4 @ %d' % (data_ptr+0x14))[0]
    methodlist_ptr = toAddr(st.unpack(word_fmt, getBytes(data_ptr.add(16 + 2 * word_size), word_size))[0])
    #nmethods = r.cmdj('pxwj 4 @ %d' % (methodlist_ptr+4))[0]
    nmethods = 0

    if methodlist_ptr.offset > 0:
        nmethods = st.unpack(endy_str + 'I', getBytes(methodlist_ptr.add(4), 4))[0]

    #ivars_ptr = r.cmdj('pxwj 4 @ %d' % (data_ptr+0x1c))[0]
    ivars_ptr = toAddr(st.unpack(word_fmt, getBytes(data_ptr.add(16 + 3 * word_size), word_size))[0])
    #nivars = r.cmdj('pxwj 4 @ %d' % (ivars_ptr+4))[0]

    nivars = 0

    if ivars_ptr.offset > 0:
        nivars = st.unpack(endy_str + 'I', getBytes(ivars_ptr.add(4), 4))[0]

    return classname, methodlist_ptr.add(8), nmethods, ivars_ptr.add(8), nivars


def get_classlist(fname):
    classlist_sect = None

    for block in currentProgram.memory.blocks:
        if '__objc_classlist' in block.name:
            classlist_sect = block
            break

    if classlist_sect is None:
        printerr('Couldn\'t find classlist section')
        return

    classlist_start = classlist_sect.start
    classlist_end = classlist_sect.start.add(classlist_sect.size)
    nclasses = classlist_sect.size // word_size

    print('%d classes' % nclasses)

    classes = {}

    for i in range(nclasses):
        class_ptr = toAddr(st.unpack(word_fmt, getBytes(classlist_start.add(i * word_size), word_size))[0])
        class_info = get_class_info(class_ptr)

        if class_info is None:
            continue

        classname, methodlist_ptr, nmethods, ivars_ptr, nivars = class_info
        
        classes[classname] = {}
        classes[classname]['methods'] = {}
        classes[classname]['ivars'] = {}

        for i in range(nmethods):
            selname_ptr, _, impptr = st.unpack(endy_str + word_str * 3, getBytes(methodlist_ptr.add(i * word_size * 3), word_size * 3))
            selname = read_string(toAddr(selname_ptr))
            classes[classname]['methods'][selname.strip('\n')] = impptr

        for i in range(nivars):
            _, ivarname_ptr, ivartype, _, _ = st.unpack(endy_str + word_str * 5, getBytes(ivarlist_ptr.add(i * word_size * 5), word_size * 5))
            ivarname = read_string(toAddr(ivarname_ptr))
            ivartype = read_string(toAddr(ivartype_ptr))
            classes[classname]['ivars'][ivarname.strip('\n')] = ivartype.strip('@"<>\n')

    with open(fname, 'w') as f:
        json.dump(classes, f, indent=True)

    return classes


if __name__ == '__main__':
    fname = 'classes.json'
    if len(sys.argv) > 1:
        fname = sys.argv[1]

    fname = join(dirname(sourceFile.absolutePath), fname)
    get_classlist(fname)