from xnu import *

"""
Recursive ipc importance chain viewing macro. This file incorporates complex python datastructures
interspersed with cvalue based objects from lldb interface. 
"""

class TaskNode(object):
    def __init__(self, task_kobj):
        self.task = task_kobj
        self.importance_refs = []
    
    @staticmethod
    def GetHeaderString():
        return GetTaskSummary.header + " " + GetProcSummary.header + " {: <18s}".format("task_imp_base")

    def __str__(self):
        out_arr = []
        if unsigned(self.task) != 0: 
            out_arr.append(GetTaskSummary(self.task) + " " + GetProcSummary(GetProcFromTask(self.task)) + " {: <#018x}".format(self.task.task_imp_base) )
        else:
            out_arr.append("Unknown task.")
        #out_arr.append("TASK: {: <#018x} {: <s}".format(self.task, GetProcNameForTask(self.task))
        for i in self.importance_refs:
            out_arr.append("\t" + i.GetBackRefChain())
        return "\n".join(out_arr)

    def AddImportanceNode(self, iinode):
        self.importance_refs.append(iinode)

class IIINode(object):
    """docstring for IIINode"""
    def __init__(self, elem, parentNode):
        super().__init__()
        self.elem = elem
        self.children = []
        self.parent = parentNode

    def addChildNode(self, elemNode):
        self.children.append(elemNode)

    def __str__(self):
        if unsigned(self.elem.iii_elem.iie_bits) & xnudefines.IIE_TYPE_MASK:
            return GetIPCImportanceInheritSummary(self.elem)
        else:
            return GetIPCImportantTaskSummary(self.elem)
    
    def GetShortSummary(self):
        to_task = self.GetToTask()
        if unsigned(self.elem.iii_elem.iie_bits) & xnudefines.IIE_TYPE_MASK:
            return "{: <#018x} INH ({:d}){: <s}".format(self.elem, GetProcPIDForTask(to_task), GetProcNameForTask(to_task))
        else:
            return "{: <#018x} IIT ({:d}){: <s}".format(self.elem, GetProcPIDForTask(to_task), GetProcNameForTask(to_task))

    def GetChildSummaries(self, prefix="\t"):
        retval = []
        for i in self.children:
            retval.append(prefix + str(i))
            retval.append(i.GetChildSummaries(prefix+"\t"))
        return "\n".join(retval)

    def GetToTask(self):
        if unsigned(self.elem.iii_elem.iie_bits) & xnudefines.IIE_TYPE_MASK:
            return self.elem.iii_to_task.iit_task
        else:
            return self.elem.iit_task

    def GetParentNode(self):
        return self.parent

    def GetBackRefChain(self):
        out_str = ""
        cur_elem = self.elem
        out_str += self.GetShortSummary()
        from_elem = Cast(cur_elem.iii_from_elem, 'ipc_importance_inherit *')
        # NOTE: We are exploiting the layout of iit and iii to have iie at the begining.
        # so casting one to another is fine as long as we tread carefully.

        while (unsigned(from_elem.iii_elem.iie_bits) & xnudefines.IIE_TYPE_MASK) == xnudefines.IIE_TYPE_INHERIT:
            out_str += " <- {: <#018x} INH ({:d}){: <s}".format(from_elem, GetProcPIDForTask(from_elem.iii_to_task.iit_task), GetProcNameForTask(from_elem.iii_to_task.iit_task))
            from_elem = Cast(from_elem.iii_from_elem, 'ipc_importance_inherit *')

        if (unsigned(from_elem.iii_elem.iie_bits) & xnudefines.IIE_TYPE_MASK) == xnudefines.IIE_TYPE_TASK:
            iit_elem = Cast(from_elem, 'ipc_importance_task *')
            out_str += " <- {: <#018x} IIT ({:d}){: <s}".format(iit_elem, GetProcPIDForTask(iit_elem.iit_task), GetProcNameForTask(iit_elem.iit_task))

        return out_str

        #unused
        cur_elem = self
        while cur_elem.parent:
            out_str += "<-" + cur_elem.GetShortSummary()
            cur_elem = cur_elem.GetParentNode()
        return out_str

def GetIIIListFromIIT(iit, rootnode, seen=set()):
    """ walk the iii queue and find each III element in a list format
    """
    for iii in IterateQueue(iit.iit_inherits, 'struct ipc_importance_inherit *',  'iii_inheritance'):
        iiiNode = IIINode(iii, rootnode)
        if unsigned(iii.iii_elem.iie_bits) & xnudefines.IIE_TYPE_MASK:
            rootnode.addChildNode(iiiNode)
            if (iii_to_task_unsigned := unsigned(iii.iii_to_task)) not in seen:
                seen.add(iii_to_task_unsigned)
                GetIIIListFromIIT(iii.iii_to_task, iiiNode, seen)
            GetTaskNodeByKernelTaskObj(iiiNode.GetToTask()).AddImportanceNode(iiiNode)
    return

AllTasksCollection = {}
def GetTaskNodeByKernelTaskObj(task_kobj):
    global AllTasksCollection
    key = hex(unsigned(task_kobj))
    if key not in AllTasksCollection:
        AllTasksCollection[key] = TaskNode(task_kobj)
    return AllTasksCollection[key]
    


@lldb_command('showallipcimportance')
def ShowInheritanceChains(cmd_args=[], cmd_options={}):
    """ show boost inheritance chains.
        Usage: (lldb) showboostinheritancechains  <task_t>
    """
    print(' ' + GetIPCImportantTaskSummary.header + ' ' + GetIPCImportanceElemSummary.header)
    for task in kern.tasks:
        if unsigned(task.task_imp_base):
            print(" " + GetIPCImportantTaskSummary(task.task_imp_base) + ' ' + GetIPCImportanceElemSummary(addressof(task.task_imp_base.iit_elem)))
            base_node = IIINode(Cast(task.task_imp_base, 'ipc_importance_inherit *'), None)
            GetIIIListFromIIT(task.task_imp_base, base_node)
            print(base_node.GetChildSummaries(prefix="\t\t"))
    
    print("\n\n ======================== TASK REVERSE CHAIN OF IMPORTANCES =========================")
    print(TaskNode.GetHeaderString())
    for k in list(AllTasksCollection.keys()):
        t = AllTasksCollection[k]
        print("\n" + str(t))

