"""@package unlinker
Documentation for this module.
This is a IDA script file built to output Section, Symbol, and Xref information


If run as script in IDA should create csv files called Xref, Symbols, and Sections
"""
import json
from idautils import *
from idc import *
from idaapi import *
import sqlite3
import sys

#Python Quicksort from interactivepython.org/courselib/static/pythonds/SortSearch/TheQuickSort.html
def quickSort(alist, blist):
    """Documentation for a function.
    @param alist list to be sorted
    @param blist an additional list that is to be sorted not by its order but by lista's order
    @brief Simple function that sorts both lists in the order of lowest to highest based on the first list
    """
    quickSortHelper(alist,0,len(alist)-1, blist)

def quickSortHelper(alist,first,last, blist):
    """Documentation for a function.
    @param alist list to be sorted
    @param first is the index of first value in quicksort
    @param last is the index of last value in quicksort
    @param blist an additional list that is to be sorted not by its order but by lista's order
    @brief Helper Function that can be called recursively
    """
    if first<last:

        splitpoint = partition(alist,first,last, blist)

        quickSortHelper(alist,first,splitpoint-1, blist)
        quickSortHelper(alist,splitpoint+1,last, blist)


def partition(alist,first,last, blist):
    """Documentation for a function.
    @param alist list to be sorted
    @param first is the index of first value in quicksort
    @param last is the index of last value in quicksort
    @param blist an additional list that is to be sorted not by its order but by lista's order
    @brief Partitions and sorts per quicksort rules. Pivotvalue just chosen as first right now
    """
    pivotvalue = alist[first]

    leftmark = first+1
    rightmark = last

    done = False
    while not done:

        while leftmark <= rightmark and alist[leftmark] <= pivotvalue:
            leftmark = leftmark + 1

        while alist[rightmark] >= pivotvalue and rightmark >= leftmark:
            rightmark = rightmark -1

        if rightmark < leftmark:
            done = True
        else:
            temp = alist[leftmark]
            alist[leftmark] = alist[rightmark]
            alist[rightmark] = temp
            temp = blist[leftmark]
            blist[leftmark] = blist[rightmark]
            blist[rightmark] = temp

    temp = alist[first]
    alist[first] = alist[rightmark]
    alist[rightmark] = temp
    temp = blist[first]
    blist[first] = blist[rightmark]
    blist[rightmark] = temp

    return rightmark

def write_all(cur):
    """Documentation for a function.
    @param cur Cursor to the sqlite database to write to
    @brief Discovers xrefs and symbols and gets information to print them. Calls print sections
    """
    #Random Declarations
    funcs = Functions()
    #Record all xref.to's to symbol list and their type to corresponding index in symbolListType
    symbolList = []
    symbolListType = []
    ptrList = []

    #Prints the section
    add_sections_to_db(cur)

    #Make sure all functions are added to symbol list. Must be prior so funcs don't get identified as objects
    for x in funcs:
        func = get_func(x)
        if not func is None:
            symbolList.append(func.startEA)
            symbolListType.append(17)
    funcs = Functions()
    #This main loop goes through all functions discovering all xrefs
    for x in funcs:
        func = get_func(x)
        if not func is None:
            fname = Name(func.startEA)
            items = FuncItems(func.startEA)
            #Every item is a line of assembly
            for i in items:
                #There can be multiple xrefs in a line of assembly if two pointers are referenced
                for xref in XrefsFrom(i, 0):
                    add_xref_info_to_db(cur, i, xref, 0)
                    #Detecting if this is a symbol and has not been added to symbol table yet
                    if xref.to not in symbolList and int(xref.type) != 21:
                        symbolList.append(xref.to)
                        symbolListType.append(int(xref.type))

                #Both the next two if statements are for detecting pointer's intial values
                if(get_operand_type(i, 0) == 2):
                    #we don't need to store its intial value more than once so created a list
                    if int(GetOperandValue(i, 0)) not in ptrList:
                        ptrList.append(int(GetOperandValue(i,0)))
                        for xrefHidden in XrefsFrom(GetOperandValue(i, 0), 0):
                            #add the new xref to xref file
                            add_xref_info_to_db(cur, GetOperandValue(i, 0), xrefHidden, 1)
                            #add the xref.to to the symbol file
                            if xrefHidden.to not in symbolList and int(xrefHidden.type) != 21:
                                symbolList.append(xrefHidden.to)
                                symbolListType.append(int(xrefHidden.type))
                #Same as above but with second operand
                if(get_operand_type(i, 1) == 2):
                    if int(GetOperandValue(i, 1)) not in ptrList:
                        ptrList.append(int(GetOperandValue(i, 1)))
                        for xrefHidden in XrefsFrom(GetOperandValue(i, 1), 0):
                            add_xref_info_to_db(cur, GetOperandValue(i, 1), xrefHidden, 1)
                            if xrefHidden.to not in symbolList and int(xrefHidden.type) != 21:
                                symbolList.append(xrefHidden.to)
                                symbolListType.append(int(xrefHidden.type))


    #Use xrefs to calculate size of symbols if they are data
    quickSort(symbolList, symbolListType)
    i = 0
    #Calculates the size of each symbol by subtraction
    while i < len(symbolList):
        #Either subtracts next Xref or Seg End, whichever comes first
        segEnd = get_segm_end(symbolList[i])
        #Last xref has to be linked against segment because no xref after it
        if i != len(symbolList)-1:
            nextXref = symbolList[i + 1]
        else:
            nextXref = segEnd
        if(nextXref < segEnd):
            diff = nextXref - symbolList[i]
        else:
            diff = segEnd - symbolList[i]
        add_symbol_to_db(cur, symbolList[i], symbolListType[i], diff)
        i += 1

#    fileX.close()
#    fileS.close()
    return;

def create_tables(cur):
    """Documentation for a function.
    @param cur Cursor to the sqlite database to write to
    @brief Clears any existing database information and sets up new tables
    """
    #Purge any old data.
    #SQLite will not create new tables if there are existing tables with the same name.
    cur.execute("DROP TABLE IF EXISTS XREF")
    cur.execute("DROP TABLE IF EXISTS SYMBOLS")
    cur.execute("DROP TABLE IF EXISTS SECTIONS")
    cur.execute("CREATE TABLE XREF(ADDRESS INTEGER, SIZE INTEGER, TARGET_ADDRESS INTEGER, TYPE STRING)")
    cur.execute("CREATE TABLE SYMBOLS(ADDRESS INTEGER, SIZE INTEGER, TYPE STRING, NAME STRING)")
    cur.execute("CREATE TABLE SECTIONS(NAME STRING, PERM INTEGER, VADDR INTEGER, VSIZE INTEGER)")

#Replacement for write_sections_csv
def add_sections_to_db(cur):
    """Documentation for a function.
    @param cur Cursor to the sqlite database to write to
    @brief Ouput section information to section.csv
    """
    for s in Segments():
        seg_name = get_segm_name(getseg(s))
        if(seg_name != 'LOAD'):
            sqlString = "INSERT INTO SECTIONS VALUES('{}',{}, {}, {})"
            sqlString = sqlString.format(seg_name, get_segm_attr(s, SEGATTR_PERM), get_segm_start(s), get_segm_end(s) - get_segm_start(s))
            cur.execute(sqlString)

#Replacement for addSymbol
def add_symbol_to_db(cur, address, xrefType, size):
    """Documentation for a function.
    @param cur Cursor to the sqlite database to write to
    @param address Is the address of the symbol to be added
    @param xrefType Is an integer representing the type of xref the symbol was
    @param size Is the size of the symbol
    @brief should output symbol at address information to a symbols.txt file
    """
    if(xrefType != 21):
        funcInside = get_func(address)
        if(funcInside is None or int(funcInside.startEA) == int(address)):
            if(xrefType == 17 and funcInside is not None):
                size = funcInside.endEA - funcInside.startEA;
            if(xrefType == 1 or xrefType == 2 or xrefType == 3):
                typeChanged = "object"
            else:
                typeChanged = "func"
            sqlString = "INSERT INTO SYMBOLS VALUES({}, {}, '{}', '{}')"
            sqlString = sqlString.format(address, size, typeChanged, str(idc.get_name(address, 0)))
            cur.execute(sqlString)

#Replacement for printXrefInfo
def add_xref_info_to_db(cur, address, xref, ptr):
    """Documentation for a function.
    @param cur Cursor to the sqlite database to write to
    @param address the address of the xref
    @param xref the xref object from idapython
    @param ptr Is 1 if ptr is 0 if not
    @brief It will then output to the xref file (Xref.csv) all the information about that xref.
    """
    #xref type 21 is some sort of dataFlow xref that is not legit
    if(xref.type != 21):
        size = 4
        typeChanged = ""
        if (int(xref.to) == int(GetOperandValue(address, 0))):
            size = len(str(hex(GetOperandValue(address, 0)))) - 1
            if(get_operand_type(address, 0) == 2):
                typeChanged = "disp"
            elif(get_operand_type(address, 0) == 5):
                if(get_operand_type(address, 1) == 2):
                    typeChanged = "imm"
                else:
                    typeChanged = "imm_alone"
        if(int(xref.to) == int(GetOperandValue(address, 1))):
            size = len(str(hex(GetOperandValue(address, 1)))) - 1
            if(get_operand_type(address, 1) == 2):
                typeChanged = "disp"
            elif(get_operand_type(address, 1) == 5):
                if(get_operand_type(address, 0) == 2):
                    typeChanged = "imm"
                else:
                    typeChanged = "imm_alone"
        if(xref.type == 17):
            typeChanged = "call"
        elif(xref.type == 19):
            typeChanged = "jmp"
        #Passed in whether this is a pointer
        if(ptr == 1):
            typeChanged = "ptr"
        sqlString = "INSERT INTO XREF VALUES({}, {}, {}, '{}')"
        sqlString = sqlString.format(address, size, xref.to, typeChanged)
        cur.execute(sqlString)



#Technically this is the main function:
print("\n\n\n\nBegin script call")
#fileName = "database.sqlite"+sys.argv[1]
#fileName = "database.sqlite"
targetNum = get_root_filename()[-1]
fileName = "unlinker"+str(targetNum)+".db"
conn = sqlite3.connect(fileName)
cur = conn.cursor()
create_tables(cur)
write_all(cur)
conn.commit()
conn.close()
qexit(0)
print("Database has been generated.")