#!/usr/bin/python

from Maat.conf.config import *
from Maat.utils.graphics import *
from Maat.utils.misc import *
from Maat.utils.apkid import *
from Maat.shared.constants import * 
import sklearn, numpy, json
import os
from androguard.misc import *

def extractCountFeatures(allSequences, includeTriggers=True):
    """
    Extracts simple numerical features depicting the occurrence counts for every unique (trigger, payload) tuple in a sequence
    :param allSequences: The sequences of all apps
    :type allSequences: list
    :param ignoreTriggers: Whether to include trigger categories as part of the features (default: True)
    :type ignoreTriggers: bool
    :return: A list of lists depicting the feature vectors extracted from every sequence
    :return: An ordered list of (trigger, payload) tuples extracted from sequences i.e. features
    """
    try:
        allTokens = []
        featureVectors = []
        # 1. Retrieve all possible (trigger, payload) tuples to estimate vector size
        ignore = "False" if not includeTriggers else "True"
        prettyPrint("Retrieving (trigger, payload) tuples from sequences. Ignoring triggers set to \"%s\"" % ignore)
        for sequence in allSequences:
            for t in sequence:
                token = t if includeTriggers else t[1]
                if not token in allTokens:
                    prettyPrint("Adding the tuple %s to list" % str(token))
                    allTokens.append(token)
        prettyPrint("Successfully retrieved %s tokens from %s sequences" % (len(allTokens), len(allSequences)))
        # 2. Extract counts of tokens from sequences
        for sequence in allSequences:
            vector = []
            for feature in allTokens:
                vector.append(float(sequence.count(feature)))
            featureVectors.append(vector)
    except Exception as e:
        prettyPrintError(e)
        return [], []

    return featureVectors, allTokens

def extractDroidmonFeatures(logPath, mode="classes", includeArguments=False):
    """
    Extracts numerical features from Droidmon-generated logs
    :param logPath: The path to the JSON-log generated by Droidmon
    :type logPath: str
    :param mode: The type of features to extract (i.e. classes, methods, both)
    :type mode: str
    :param includeArguments: Whether to include the method argument in the trace
    :type includeArguments: bool
    :return: Two lists depicting the trace found in the log and counts of items it contains
    """
    try:
        features = []
        # Parse the droidmon log
        if not os.path.exists(logPath):
            prettyPrint("Unable to locate \"%s\"" % logPath, "warning")
            return [], []
        lines = open(logPath).read().split('\n')
        if VERBOSE == "ON":
            prettyPrint("Successfully retrieved %s lines from log" % len(lines), "debug")
        droidmonLines = [l for l in lines if l.lower().find("droidmon-apimonitor-") != -1]
        # Generate trace from lines
        trace = []
        for line in droidmonLines:
            tmp = line[line.find("{"):].replace('\n','').replace('\r','')
            # Extract class and method
            c, m = "", ""
            pattern = "class\":\""
            index = tmp.find(pattern)
            c = tmp[index+len(pattern):tmp.find("\"", index+len(pattern))]
            pattern = "method\":\""
            index = tmp.find(pattern)
            m = tmp[index+len(pattern):tmp.find("\"", index+len(pattern))]
            if includeArguments:
                pattern = "args\":["
                index = tmp.find(pattern)
                a = tmp[index+len(pattern):tmp.find(']', index+len(pattern))]
            # Prepare to add args
            args = "" if not includeArguments else a
            # Append to trace
            if mode == "classes":
                trace.append("%s(%s)" % (c, args))
            elif mode == "methods":
                trace.append("%s(%s)" % (m, args))
            elif mode == "both":
                trace.append("%s.%s(%s)" % (c, m, args))
        # Go over droidmon classes and count occurrences
        source = []
        if mode == "classes":
            source = droidmonDefaultClasses
        elif mode == "methods":
            source = droidmonDefaultMethods
        elif mode == "both":
            source = droidmonDefaultAPIs

        # The loop
        noArgsTrace = [t[:t.find('(')] for t in trace]
        for i in source:
            features.append(noArgsTrace.count(i))

    except Exception as e:
        prettyPrintError(e)
        return [], []

    return trace, features

def extractStaticFeatures(apkPath, preAPK=None, preDEX=None, preVM=None):
    """Extracts static numerical features from APK using Androguard"""
    try:
        features = [[], [], [], []] # Tuples are immutable
        if os.path.exists(apkPath.replace(".apk",".static")):
            prettyPrint("Found a pre-computed static features file")
            bFeatures, pFeatures, aFeatures, allFeatures = [], [], [], []
            try:
                possibleExtensions = [".basic", ".perm", ".api", ".static"]
                for ext in possibleExtensions:
                    if os.path.exists(apkPath.replace(".apk", ext)):
                        content = open(apkPath.replace(".apk", ext)).read()
                        if len(content) > 0:
                            features[possibleExtensions.index(ext)] = [float(f) for f in content[1:-1].split(',') if len(f) > 0]

                return tuple(features)

            except Exception as e:
                prettyPrintError(e)
                prettyPrint("Could not extract features from \".static\" file. Continuing as usual", "warning")
        prettyPrint("Starting analysis on \"%s\"" % apkPath, "debug")
        # 1. Analyze APK and retrieve its components
        if preAPK and preDEX and preVM:
            apk, dex, vm = preAPK, preDEX, preVM   
        else:
            apk, dex, vm = AnalyzeAPK(apkPath)
            dex = dex[0]
        # 2. Add features to the features vector
        basicFeatures, permissionFeatures, apiCallFeatures, allFeatures = [], [], [], []
        # 2.a. The APK-related features
        prettyPrint("Extracting basic features", "debug")
        minSDKVersion = 0.0 if not apk.get_min_sdk_version() else float(apk.get_min_sdk_version())
        maxSDKVersion = 0.0 if not apk.get_max_sdk_version() else float(apk.get_max_sdk_version())
        basicFeatures.append(minSDKVersion)
        basicFeatures.append(maxSDKVersion)
        basicFeatures.append(float(len(apk.get_activities()))) # No. of activities
        basicFeatures.append(float(len(apk.get_services()))) # No. of services
        basicFeatures.append(float(len(apk.get_receivers()))) # No. of broadcast receivers
        basicFeatures.append(float(len(apk.get_providers()))) # No. of providers
        # 2.b. Harvest permission-related features
        prettyPrint("Extracting permissions-related features", "debug")
        aospPermissions = float(len(apk.get_permissions())) # Android permissions requested by the app
        declaredPermissions = float(len(apk.get_declared_permissions())) # Custom permissions declared by the app
        dangerousPermissions = float(len([p for p in apk.get_requested_aosp_permissions_details().values() if p["protectionLevel"] == "dangerous"]))
        totalPermissions = float(len(apk.get_permissions()))
        permissionFeatures.append(totalPermissions) # No. of permissions
        if totalPermissions > 0:
            permissionFeatures.append(aospPermissions/totalPermissions) # AOSP permissions : Total permissions
            permissionFeatures.append(declaredPermissions/totalPermissions) # Third-party permissions : Total permissions
            permissionFeatures.append(dangerousPermissions/totalPermissions) # Dangerous permissions : Total permissions
        else:
            permissionFeatures.append(0.0)
            permissionFeatures.append(0.0)
            permissionFeatures.append(0.0)
        # 2.c. The DEX-related features (API calls)
        prettyPrint("Extracting API calls from dex code", "debug")
        apiCallFeatures.append(float(len(dex.get_classes()))) # Total number of classes
        apiCallFeatures.append(float(len(dex.get_strings()))) # Total number of strings
        apiCategories = sensitiveAPICalls.keys()
        apiCategoryCount = [0.0] * len(apiCategories)
        for c in dex.classes.get_names():
            currentClass = dex.get_class(c)
            if not currentClass:
                continue
            code = currentClass.get_source()
            if len(code) < 1:
                continue
            for category in apiCategories:
                if code.find(category) != -1:
                    for call in sensitiveAPICalls[category]:
                        apiCategoryCount[apiCategories.index(category)] += float(len(re.findall(call, code)))

        apiCallFeatures += apiCategoryCount
        # All features = combination of all features :)
        allFeatures = basicFeatures + permissionFeatures + apiCallFeatures
        # Add the compiler id as well
        output = scan(apkPath, 60, "yes")
        try:
            compiler = output["files"][0]["results"]["compiler"][0]
            allFeatures.append(allCompilers.index(compiler))
        except Exception as e:
            allFeatures.append(allCompilers.index("n/a"))


    except Exception as e:
        prettyPrintError(e)
        return [], [], [], []

    return basicFeatures, permissionFeatures, apiCallFeatures, allFeatures
