""" IDAPython plugin to display all system calls of AMD/Intel 32/64bit
    Linux userland binaries.

    Work in progress...

    The system call ABI from the following link are supported.
    http://esec-lab.sogeti.com/posts/2011/07/05/linux-syscall-abi.html

    by n0p
"""

import idaapi
import ida_idaapi
import ida_bytes
import ida_gdl
import ida_idp
import ida_search
import ida_segment
import idautils
import idc

import time

try:
    from miasm2.core.bin_stream_ida import bin_stream_ida
    from miasm2.analysis.depgraph import DependencyGraph

    from utils import guess_machine
except:
    guess_machine = None

from SystemCalls_constants import *


class SystemCallView(idaapi.Choose2):

    def __init__(self, systemCalls):

        self.systemCalls = systemCalls

        idaapi.Choose2.__init__(self, 'System call', [['Address', 13],
                                                      ['Type', 10],
                                                      ['Number', 10],
                                                      ['Name', 20],
                                                      ['Pointer Size', 12]])

        self.items = list()

        self.nop_items = list()

        self.icon = 150

        self.cmd_nop = None

        self.initialized = False

    def __fillView(self):

        start = time.time()

        self.systemCalls.searchSystemCalls()

        end = time.time()
        print ('[*] It took %d seconds to discover the system calls.' %
               (end - start))

        self.items = list()

        start = time.time()

        if len(self.systemCalls.x86) != 0:
            for faddr in self.systemCalls.x86.iterkeys():
                calls = self.systemCalls.getSystemCallNumber(
                    self.systemCalls.x86[faddr], x86SystemCalls)
                for call in calls:
                    try:
                        self.items.append(
                            ['0x%08X' % call[0], systemCallTypes[call[1]],
                             '0x%03X' % int(call[2]),
                             x86SystemCalls[int(call[2])], '32bit'])
                    except:
                        # No hex system call number found.
                        self.items.append(
                            ['0x%08X' % call[0], systemCallTypes[
                                call[1]], str(call[2]), '', '32bit'])

        if len(self.systemCalls.x86_64) != 0:
            for faddr in self.systemCalls.x86_64.iterkeys():
                calls = self.systemCalls.getSystemCallNumber(
                    self.systemCalls.x86_64[faddr], x86_64SystemCalls)
                for call in calls:
                    try:
                        self.items.append(
                            ['0x%08X' % call[0], systemCallTypes[call[1]],
                             '0x%03X' % int(call[2]),
                             x86_64SystemCalls[int(call[2])], '64bit'])
                    except:
                        # No hex system call number found.
                        self.items.append(
                            ['0x%08X' % call[0], systemCallTypes[
                                call[1]], str(call[2]), '', '64bit'])

        end = time.time()
        print ('[*] It took %d seconds to analyze the system calls.' %
               (end - start))

        self.items.sort(key=lambda tup: tup[0])

    def OnClose(self):
        pass

    def OnCommand(self, n, cmd_id):
        if cmd_id == self.cmd_nop:
            start_ea = int(self.items[n][0], 16)
            end_ea = start_ea + idc.get_item_size(start_ea)

            self.nop_items.append(self.items[n][0])

            ida_bytes.patch_bytes(start_ea, (end_ea - start_ea) * '\x90')

    def OnGetIcon(self, n):
        if not len(self.items) > 0:
            return -1

        if self.items[n][3] == '':
            # No system call number found => display red icon.
            return 59
        else:
            # Display green icon.
            return 61

    def OnGetLine(self, n):
        return self.items[n]

    def OnGetLineAttr(self, n):
        if len(self.nop_items) > 0:
            if self.items[n][0] in self.nop_items:
                return [0xB6B6B4, 0]

    def OnGetSize(self):
        return len(self.items)

    def OnRefresh(self, n):
        self.__fillView()

    def OnSelectLine(self, n):
        idaapi.jumpto(int(self.items[n][0], 16))

    def show(self):
        if not self.initialized:
            self.initialized = True
            self.__fillView()

        if self.Show() < 0:
            return False

        if self.cmd_nop is None:
            self.cmd_nop = self.AddCommand(
                'NOP system call', flags=idaapi.CHOOSER_POPUP_MENU, icon=50)

        return True


class SystemCall():

    class Func():
        def __init__(self):

            # System calls inside the function.
            self.calls = set()

            # Return value of idaapi.get_func(addr)
            self.f = None

    class Call():
        def __init__(self, addr, sctype):

            self.addr = addr
            self.sctype = sctype

        def __hash__(self):
            return self.addr

    def __init__(self):

        # Init miasm stuff.
        if guess_machine is not None:
            self.machine = guess_machine()
            self.mn = self.machine.mn
            self.dis_engine = self.machine.dis_engine
            self.ira = self.machine.ira

            self.mdis = self.dis_engine(
                bin_stream_ida(), dont_dis_nulstart_bloc=True)
            self.ir_arch = self.ira(self.mdis.symbol_pool)

            # Populate symbols with ida names
            for ad, name in idautils.Names():
                if name is None:
                    continue
                self.mdis.symbol_pool.add_label(name, ad)

            self.elements = set([self.mn.regs.RAX])

        self.x86 = dict()
        self.x86_64 = dict()

        self.systemCallView = SystemCallView(self)

    def __getSystemCallNumberByComment(self, addr, scstrings):
        cmt = idc.Comment(addr)

        if cmt and cmt.startswith('LINUX - '):
            try:
                return scstrings.index(cmt.replace('LINUX - ', ''))
            except:
                return None

    def getSystemCallNumber(self, func, scstrings):
        """ Get the value of rax/eax at the time of the system call.
        """

        sol = list()

        # Get the analysis results from IDA, by reading IDA's comments at
        # system calls.
        calls = set()

        for call in func.calls:
            number = self.__getSystemCallNumberByComment(call.addr, scstrings)

            if number:
                sol.append([call.addr, call.sctype, str(number)])
                calls.add(call)

        func.calls -= calls

        # Just proceed with depgraph if IDA detected a function and miasm had
        # been imported.
        if not func.f or guess_machine is None:
            for call in func.calls:
                sol.append([call.addr, call.sctype, ''])
            return sol

        blocs = self.mdis.dis_multibloc(func.f.startEA)

        # Generate IR
        try:
            for bloc in blocs:
                self.ir_arch.add_bloc(bloc)
        except:
            for call in func.calls:
                sol.append([call.addr, call.sctype, ''])
            return sol

        for call in func.calls:
            addr = call.addr

            # Check if addr is in a basic block without an entry.
            if len(self.ir_arch.getby_offset(addr)) == 0:
                fc = qflow_chart_t('', func.f, ida_idaapi.BADADDR,
                                   ida_idaapi.BADADDR, ida_gdl.FC_PREDS)

                try:
                    # Iterate through all basic blocks.
                    for i in xrange(0, fc.size()):
                        if fc[i].startEA <= addr and addr < fc[i].endEA:
                            # Basic block without entry found.
                            blocs = self.mdis.dis_multibloc(fc[i].startEA)

                            # Generate IR
                            for bloc in blocs:
                                self.ir_arch.add_bloc(bloc)
                except:
                    sol.append([addr, call.sctype, ''])
                    continue

            cur_bloc = list(self.ir_arch.getby_offset(addr))[0]
            cur_label = cur_bloc.label

            for line_nb, l in enumerate(cur_bloc):
                if l.instr.offset == addr:
                    break

            # Get dependency graphs
            dg = DependencyGraph(self.ir_arch, follow_call=False)
            graphs = dg.get(cur_label, self.elements, line_nb, set(
                [self.ir_arch.symbol_pool.getby_offset(func.f.startEA)]))

            while 1:
                try:
                    graph = graphs.next()
                except StopIteration:
                    break

                if not graph.has_loop:
                    sol.append([addr, call.sctype, graph.emul().values()[0]])

        return sol

    def __addCall(self, addr, sctype, arch):
        f = idaapi.get_func(addr)

        if not f:
            faddr = 0
        else:
            faddr = f.startEA

        if faddr in arch:
            arch[faddr].calls.add(self.Call(addr, sctype))
        else:
            arch[faddr] = self.Func()
            arch[faddr].calls.add(self.Call(addr, sctype))
            arch[faddr].f = f

    def __findCalls(self, seg, sbytes, slength, sctype, arch):
        addr = ida_search.find_binary(
            seg.startEA, seg.endEA, sbytes, 16, ida_search.SEARCH_DOWN)

        while addr != ida_idaapi.BADADDR:
            if (ida_bytes.get_item_head(addr) == addr and
                    ida_bytes.get_item_size(addr) == slength):
                self.__addCall(addr, sctype, arch)

            addr = ida_search.find_binary(
                addr + 1, seg.endEA, sbytes, 16, ida_search.SEARCH_DOWN)

        addr = ida_search.find_binary(
            seg.startEA, seg.endEA, sbytes, 16, ida_search.SEARCH_DOWN)

    def searchSystemCalls(self):
        """ Looks for 'int 80', 'sysenter', 'syscall' and 'gs:[10h]' system calls.
        """
        seg = ida_segment.get_first_seg()

        while seg:

            # Check if segment is executable
            if seg.perm & 1:
                # int 80h. Just on 32bit.
                self.__findCalls(seg, 'CD 80', 2, 0, self.x86)

                # sysenter. Just on 32bit.
                self.__findCalls(seg, '0F 34', 2, 1, self.x86)

                # syscall. 32bit just on AMD. 64bit on AMD and Intel.
                if ida_idp.ph.flag & ida_idp.PR_USE64:
                    self.__findCalls(seg, '0F 05', 2, 2, self.x86_64)
                else:
                    self.__findCalls(seg, '0F 05', 2, 2, self.x86)

                # gs:[10h]. Just on 32bit.
                self.__findCalls(seg, '65 FF 15 10 00 00 00', 7, 3, self.x86)

            seg = ida_segment.get_next_seg(seg.startEA)

    def showView(self):
        self.systemCallView.show()


class SystemCallPlugin_t(idaapi.plugin_t):
    flags = 0
    comment = ''
    help = ''
    wanted_name = 'System Calls'
    wanted_hotkey = ''

    def init(self):
        global systemCalls

        if idaapi.ph_get_id() == idaapi.PLFM_386:
            # Check if already initialized
            if 'systemCalls' not in globals():
                systemCalls = SystemCall()

            return idaapi.PLUGIN_KEEP
        else:
            return idaapi.PLUGIN_SKIP

    def run(self, arg):
        global systemCalls

        systemCalls.showView()

    def term(self):
        if 'systemCalls' in globals():
            del globals()['systemCalls']


def PLUGIN_ENTRY():
    return SystemCallPlugin_t()