package ca.uwaterloo.cs858.analysis

import com.ibm.wala.classLoader.CallSiteReference
import com.ibm.wala.classLoader.IClass
import com.ibm.wala.classLoader.IMethod
import com.ibm.wala.classLoader.Module
import com.ibm.wala.dalvik.classLoader.DexFileModule
import com.ibm.wala.dalvik.classLoader.DexIRFactory
import com.ibm.wala.dalvik.ipa.callgraph.impl.DexEntryPoint
import com.ibm.wala.ipa.callgraph.*
import com.ibm.wala.ipa.callgraph.impl.Util
import com.ibm.wala.ipa.cha.ClassHierarchy
import com.ibm.wala.ssa.*
import com.ibm.wala.types.ClassLoaderReference
import com.ibm.wala.types.TypeReference
import com.ibm.wala.util.debug.Assertions
import com.ibm.wala.util.graph.Graph
import com.ibm.wala.util.graph.impl.SlowSparseNumberedGraph
import com.ibm.wala.util.graph.traverse.BFSPathFinder
import com.ibm.wala.viz.DotUtil
import java.io.File
import java.util.concurrent.Callable

const val ERROR_VAL = Int.MIN_VALUE

class Analysis(private val ApkFiles: Array<String>, private val OutputDir: String) : Callable<Analysis> {
    private lateinit var scope: AnalysisScope
    private lateinit var cha: ClassHierarchy

    @Throws(Exception::class)
    override fun call(): Analysis {
        initialize()
        analyze()
        return this
    }

    @Throws(Exception::class)
    private fun initialize() {
        val util = AnalysisScopeUtil()
        scope = util.makeAnalysisScope()
        cha = ClassHierarchy.make(scope)

        // make class hierarchy for the target module only
        val appScope = util.makeAnalysisScope(ApkFiles[0])
        val appCha = ClassHierarchy.make(appScope)
        val appLoader = appCha!!.factory.getLoader(ClassLoaderReference.Application, appCha, appScope)

        // remove classes (in target module) from existing class hierarchy's class loader
        val toRemove = mutableSetOf<IClass>()
        val appIter = appLoader.iterateAllClasses()
        while (appIter.hasNext()) {
            val k = appIter.next()
            toRemove.add(k)
        }

        val oriLoader = cha.factory.getLoader(ClassLoaderReference.Application, cha, scope)
        oriLoader.removeAll(toRemove)

        // add target module to scope
        scope.setLoaderImpl(ClassLoaderReference.Application, "com.ibm.wala.dalvik.classLoader.WDexClassLoaderImpl")
        val list: ArrayList<Module?> = ArrayList()
        for (apkFile in ApkFiles) {
            val dexModule: Module = DexFileModule.make(File(apkFile))
            scope.addToScope(ClassLoaderReference.Application, dexModule)
            list.add(dexModule)
        }

        // init the module s.t. the classes are included in the scope's class loader
        oriLoader.init(list)
        // make new class hierarchy
        cha = ClassHierarchy.make(scope, cha.factory)
    }

    @Throws(Exception::class)
    private fun analyze() {
        val cache = AnalysisCache(DexIRFactory())
        val entryPoint = buildSearch()

        val paths = findWitness(entryPoint, cache)
        if (paths.isNotEmpty()) {
            println()
            for (path in paths) {
                println("\t\t" + path.joinToString("->") { it.name.toString() })
            }
        } else {
            println("No path")
        }
    }

    private fun buildSearch(): SearchParameters {
        val methods = arrayOf(
                cha.lookupClass(TypeReference.find(ClassLoaderReference.Application, "Lca/uwaterloo/cs858shamirhost/MainActivity"))
                        .declaredMethods
                        .find { it.name.toString() == "onCreate" }!!,

                cha.lookupClass(TypeReference.find(ClassLoaderReference.Application, "Lca/uwaterloo/cs858shamirclient/MainActivity"))
                        .declaredMethods
                        .find { it.name.toString() == "onCreate" }!!,

                cha.lookupClass(TypeReference.find(ClassLoaderReference.Application, "Lca/uwaterloo/cs858shamirhost/MainActivity"))
                        .declaredMethods
                        .find { it.name.toString() == "setData" }!!
        )

        return object : SearchParameters(methods) {
            override fun sourceFinder(graph: Graph<IMethod>): List<IMethod>? {
                for (m in graph) {
                    if (m.name.toString() == "getDeviceId") {
                        return listOf(m)
                    }
                }

                return null
            }

            override fun intermediateFinder(graph: Graph<IMethod>): List<IMethod>? {
                return graph.filter { it.declaringClass.name.toString().endsWith("Intent") && it.name.toString().startsWith("get") }
            }

            override fun sinkFinder(graph: Graph<IMethod>): List<IMethod>? {
                for (m in graph) {
                    if (m.name.toString() == "makeRequest") {
                        return listOf(m)
                    }
                }

                return null
            }
        }
    }

    private fun findWitness(entryPoint: SearchParameters, cache: AnalysisCache): List<List<IMethod>> {
        val entryList = mutableListOf<Entrypoint>()
        for (method in entryPoint.methods) {
            val dep = DexEntryPoint(method, cha)
            entryList.add(dep)
        }
        val options = AnalysisOptions(scope, entryList)
        // context insensitive call graph
        val builder: CallGraphBuilder = Util.makeVanillaZeroOneCFABuilder(options, cache, cha, scope)
        val cg = builder.makeCallGraph(options, null)
        val graph = buildCallSiteGraph(cg)

        AnalysisUtil.writeGraph(graph, OutputDir)

        val listOfPutExtras = graph.filter { it.declaringClass.name.toString().endsWith("Intent") && it.name.toString() == "putExtra" }
        val listOfGetExtras = graph.filter { it.declaringClass.name.toString().endsWith("Intent") && it.name.toString().startsWith("get") }

        for (putExtra in listOfPutExtras) {
            for (getExtra in listOfGetExtras) {
                if (putExtra.getParameterType(2) == getExtra.returnType) {
                    graph.addEdge(putExtra, getExtra)
                }
            }
        }

        val sources = entryPoint.sourceFinder(graph)
        val intermediates = entryPoint.intermediateFinder(graph) ?: listOf(cg.entrypointNodes.first().method)
        val sinks = entryPoint.sinkFinder(graph)

        val paths = mutableListOf<List<IMethod>>()
        if (sources == null || sinks == null || sources.isEmpty() || sinks.isEmpty()) {
            return paths
        }

        for (source in sources) {
            for (intermediate in intermediates) {
                sinks@ for (sink in sinks) {
                    val entryNodes = cg.entrypointNodes.toList()
                    val sourcePath =  BFSPathFinder(graph, entryNodes.first().method, source).find()
                    val intermediatePath1 = BFSPathFinder(graph, entryNodes.first().method, intermediate).find()
                    val intermediatePath2 = BFSPathFinder(graph, (entryNodes[1] ?: entryNodes.first()).method, intermediate).find()
                    val sinkPath = BFSPathFinder(graph, (entryNodes[1] ?: entryNodes.first()).method, sink).find()

                    if (sourcePath != null && sinkPath != null) {
                        val repeated1 = sourcePath.intersect(intermediatePath1).toMutableList()
                        val repeated2 = intermediatePath2.intersect(sinkPath).toMutableList()
                        repeated1.removeAt(0);
                        repeated2.removeAt(0)
                        val path1 = listOf(sourcePath, intermediatePath1.reversed()).flatten().toMutableList()
                        val path2 = listOf(intermediatePath2, sinkPath.reversed()).flatten().toMutableList()
                        path1.removeAll(repeated1)
                        path2.removeAll(repeated2)

                        val path = listOf(path1, path2).flatten()
                        var currentMode = 0
                        var prevParam: Int? = null
                        paths@ for (i in 1 until path.size) {
                            var parent: IMethod
                            var child: IMethod
                            if (currentMode < 2) {
                                parent = path[i]
                                child = path[i - 1]
                            } else {
                                child = path[i]
                                parent = path[i - 1]
                            }

                            when (currentMode) {
                                0 -> {
                                    if (parent.name == path[i + 1].name) {
                                        currentMode = 1
                                        continue@paths
                                    }

                                    if (prevParam == null) {
                                        val result = checkForTaintedReturn(cg, parent) { AnalysisUtil.getMethodReturnSeedVN(it, child) }
                                        if (result == ERROR_VAL) {
                                            continue@sinks
                                        }

                                        val parentNode = cg.find { it.method.name == parent.name }
                                        val node = cg.getPredNodes(parentNode).next()
                                        var parentCall: CallSiteReference? = null
                                        for (callSite in node.iterateCallSites()) {
                                            if (callSite.declaredTarget.name == parent.name) {
                                                parentCall = callSite
                                            }
                                        }
                                        prevParam = node.ir.getCalls(parentCall!!).first().getReturnValue(0)
                                    } else {
                                        prevParam = checkForTaintedReturn(cg, parent) { prevParam!! }
                                        if (prevParam == ERROR_VAL) {
                                            continue@sinks
                                        }
                                    }
                                }
                                1 -> {
                                    Assertions.productionAssertion(parent.name == child.name, "At top?")
                                    val prev = path[i - 2]
                                    val next = path[i + 1]
                                    prevParam = findPassToMethod(cg, parent, next) { AnalysisUtil.getMethodReturnSeedVN(it, prev) }
                                    if (prevParam == ERROR_VAL) {
                                        continue@sinks
                                    }
                                    currentMode = 2
                                }
                                2 -> {
                                    if (parent.declaringClass.name.toString().endsWith("Intent") && parent.declaringClass.name == child.declaringClass.name) {
                                        currentMode = 3
                                        continue@paths
                                    }

                                    prevParam = findPassToMethod(cg, parent, child) { AnalysisUtil.getParamSeedVN(it, prevParam!!) }
                                    if (prevParam == ERROR_VAL) {
                                        continue@sinks
                                    }
                                }
                                3 -> {
                                    Assertions.productionAssertion(parent.declaringClass.name == child.declaringClass.name, "At Intent pass?")
                                    val prev = path[i - 3]
                                    val result = checkIntentTarget(cg, prev)

                                    if (result == 0)  {
                                        System.err.println("No evidence, making the assumption they are linked")
                                    }

                                    prevParam = null
                                    currentMode = 0
                                }
                            }
                        }

                        paths.add(path)
                    }
                }
            }
        }

        return paths.toList()
    }

    private fun checkIntentTarget(cg: CallGraph, parent: IMethod): Int {
        val node = cg.find { it.method == parent } ?: return ERROR_VAL
        val ir = node.ir
        for (vn in 1..ir.symbolTable.maxValueNumber) {
            val value = ir.symbolTable.getValue(vn)?.toString() ?: continue
            val packageName = cg.entrypointNodes.toList()[1].method.declaringClass.name.toString()
            if (value.endsWith(MattUtil.convertToPackageName(packageName))) {
                return 1
            }
        }
        return 0
    }

    private fun findPassToMethod(cg: CallGraph, parent: IMethod, child: IMethod, getSeedVN: (ir: IR) -> Int): Int {
        // Keeping track of candidates that use value numbers derived from our sink
        val candidates = mutableListOf<IMethod>()
        val uses = mutableListOf<Int>()
        val node = cg.find { it.method == parent } ?: return ERROR_VAL

        // We convert the node into the underlying intermediate representation; this is a long list of instructions, essentially.
        val ir = node.ir
        // We will iterate over all the instructions making up the method
        val iter: Iterator<SSAInstruction> = ir.iterateAllInstructions()
        // This is a list of all value numbers that either come from or are our sink
        val seed = getSeedVN(ir)
        if (seed == ERROR_VAL) {
            return ERROR_VAL
        }

        val possibleVNs = mutableListOf(seed)
        while (iter.hasNext()) {
            val s = iter.next()
//            println(s.toString(ir.symbolTable))
            // This basically asks "in this instruction a method call?"
            if (s is SSAAbstractInvokeInstruction) {
                // each instruction "uses" some number of value numbers, i.e. as parameters in this case.
                for (i in 1..s.numberOfUses) {
                    val use = s.getUse(i-1)
                    // Is this parameter in our watchlist?
                    if (possibleVNs.contains(use)) {
                        val declaringKlass = cha.lookupClass(s.declaredTarget.declaringClass) ?: continue
                        val method = declaringKlass.getMethod(s.declaredTarget.selector) ?: continue
                        // It is, so add this method to the list of who has used tainted value numbers
                        candidates.add(method)
                        uses.add(i - 1)

                        // if the method is void, it defs only one return value - a pontential exception. If it's two, it returns some value
                        if (s.numberOfDefs == 2) {
                            // Add the returned value's value number to the watchlist
                            possibleVNs.add(s.getReturnValue(0))
                        }
                    }
                }

                if (s.isDispatch) {
                    possibleVNs.add(s.getUse(0))
                }
            } else {
                for (i in 1..s.numberOfUses) {
                    val use = s.getUse(i-1)
                    // Is this parameter in our watchlist?
                    if (possibleVNs.contains(use)) {
                        possibleVNs.add(s.getDef(0))
                    }
                }
            }
        }

//            println(candidates.map { it.name.toString() }.joinToString(", "))
        // Is the next step in our tree in the list of methods that used a tainted value?
        return if (candidates.contains(child)) {
            uses[candidates.indexOf(child)]
        } else {
            ERROR_VAL
        }
    }

    private fun checkForTaintedReturn(cg: CallGraph, parent: IMethod, getSeedVN: (ir: IR) -> Int): Int {
        // Keeping track of candidates that use value numbers derived from our sink
        val node = cg.find { it.method == parent } ?: return ERROR_VAL

        // We convert the node into the underlying intermediate representation; this is a long list of instructions, essentially.
        val ir = node.ir
        // We will iterate over all the instructions making up the method
        val iter: Iterator<SSAInstruction> = ir.iterateAllInstructions()
        // This is a list of all value numbers that either come from or are our sink
        val seed = getSeedVN(ir)
        if (seed == ERROR_VAL) {
            return ERROR_VAL
        }
        val possibleVNs = mutableListOf(seed)
        while (iter.hasNext()) {
            val s = iter.next()
//            println(s.toString(ir.symbolTable))
            // This basically asks "in this instruction a method call?"
            if (s is SSAAbstractInvokeInstruction) {
                // each instruction "uses" some number of value numbers, i.e. as parameters in this case.
                for (i in 1..s.numberOfUses) {
                    val use = s.getUse(i-1)
                    // Is this parameter in our watchlist?
                    if (possibleVNs.contains(use)) {
                        val declaringKlass = cha.lookupClass(s.declaredTarget.declaringClass) ?: continue
                        val method = declaringKlass.getMethod(s.declaredTarget.selector) ?: continue
                        // It is, so add this method to the list of who has used tainted value numbers

                        // if the method is void, it defs only one return value - a pontential exception. If it's two, it returns some value
                        if (s.numberOfDefs == 2) {
                            // Add the returned value's value number to the watchlist
                            possibleVNs.add(s.getReturnValue(0))
                        }
                    }
                }

                if (s.isDispatch) {
                    println(s.declaredTarget.name.toString())
                    possibleVNs.add(s.getUse(0))
                }
            } else if (s is SSAReturnInstruction) {
                if (possibleVNs.contains(s.result)) {
                    return 1
                }
            } else {
                for (i in 1..s.numberOfUses) {
                    val use = s.getUse(i-1)
                    // Is this parameter in our watchlist?
                    if (possibleVNs.contains(use)) {
                        possibleVNs.add(s.getDef(0))
                    }
                }
            }
        }

        return ERROR_VAL
    }

    private fun buildCallSiteGraph(cg: CallGraph): Graph<IMethod> {
        val graph = SlowSparseNumberedGraph.make<IMethod>()

        if (cg.entrypointNodes.isEmpty()) {
            throw IllegalArgumentException("Call graph must have an entrypoint node")
        }

        for (node in cg) {
            if (node.method.declaringClass.name.toString().startsWith("Ljava")) {
                continue
            }

            if (node.method.declaringClass.name.toString().startsWith("Landroidx")) {
                continue
            }

            if (!graph.containsNode(node.method)) {
                graph.addNode(node.method)
            }

            for (ref in node.iterateCallSites()) {
                val declaringKlass = cha.lookupClass(ref.declaredTarget.declaringClass) ?: continue

                val method = declaringKlass.getMethod(ref.declaredTarget.selector) ?: continue

                if (!graph.containsNode(method)) {
                    graph.addNode(method)
                }
                graph.addEdge(node.method, method)
            }
        }

        val setOnClickListener = graph.find { it.name.toString() == "setOnClickListener" }
        val setData = graph.find { it.name.toString() == "setData" }
        graph.addEdge(setOnClickListener!!, setData!!)

        DotUtil.dotify(graph, { method ->
            val builder = StringBuilder()
            builder.append(method.declaringClass.name.toString())
            builder.append("\n")
            builder.append(method.name.toString())
            builder.append("\n")
            builder.append(method.selector.toString())
            builder.toString()
        }, OutputDir + File.separator + "callSiteGraph.dot", null, null)

        return graph
    }

    fun isCryptoFunction(node: CGNode): Boolean{
        val ir = node.ir
        if(ir == null){
            return false
        }
        val iter: Iterator<SSAInstruction> = ir.iterateAllInstructions()
        var statmentCounter = 0
        var mathCounter = 0
        var crypto = false
        while (iter.hasNext()) {
            val s = iter.next()
            statmentCounter++
            //println(s.toString(ir.symbolTable))
            //Does it use BigInteger or security
            if (s is SSAAbstractInvokeInstruction) {
                if (s.callSite.declaredTarget.declaringClass.name.toString() == "Ljava/math/BigInteger") {
                    println("Found BigInteger call")
                    return true
                }
                if (s.callSite.declaredTarget.declaringClass.name.toString() == "Ljava/math/BigDecimal") {
                    println("Found BigDecimal call")
                    return true
                }
                if (s.callSite.declaredTarget.declaringClass.name.toString().startsWith("Ljava/security/")) {
                    println("Found security call")
                    return true
                }
                if (s.callSite.declaredTarget.declaringClass.name.toString().startsWith("Ljavax/crypto/")) {
                    println("Found crypto call")
                    return true
                }
                if (s.callSite.declaredTarget.declaringClass.name.toString().startsWith("Ljavax/xml/crypto/")) {
                    println("Found xml.crypto call")
                    return true
                }
            }else if(s is SSABinaryOpInstruction){
                mathCounter++
            }

        }
        if(statmentCounter == 0){
            return false
        }
        else if(mathCounter/statmentCounter >= .5){
            print("Contains mostly math")
            crypto = true
        }
        return crypto
    }
}

//TODO I think I need help with this...
//The algorithm is as follows:

//Read sinks from TXT File

//Make Sources like this except not an array list whatever list you use

//var cryptoFuntions = ArrayList<IMethod>()
//for(node in cg){
//    if(isCryptoFunction(node)){
//        cryptoFuntions.add(node.method)
//        println(node.method.toString())
//    }
//}

//For all verified Paths return (maybe print to file) the parameters the parameters of the crypto functions