from xml.etree.ElementTree import ElementTree
import glob
import subprocess
import os
import codecs
import threading
import time

taskParam_noThread = 10
taskParam_apilistfilename = 'dangerAPIs.txt'
taskParam_bytecodelistfilename = 'Bytecodes.txt'
taskParam_apklistfilename = './apklist.txt'
taskParam_resultfilename = 'Methods.txt'
taskParam_rawdatadirname = "./RawData(3000)/"

def searchToEXmethods(dirname, methods, dalist, bytelist):
    flist = os.listdir(dirname)
    for f in flist:
        next = os.path.join(dirname, f)
        if os.path.isdir(next):
            searchToEXmethods(next,methods,dalist,bytelist)
        else:
            (mset)= extractMethods(next,dalist,bytelist)
            methods.update(mset)

def extractMethods(filename, dalist, bytelist):
    ext = os.path.splitext(filename)[-1]
    MethodSet = set()
    ongoing = False
    if ext == '.smali':
        if not os.path.exists(os.path.abspath(filename)):
            return (MethodSet)
        f = open(os.path.abspath(filename),'r')
        try:
            lines = f.readlines()
        except UnicodeDecodeError:

            return (MethodSet)

        for line in lines:
            if line.count('.method')==1:
                Mfeats = list()
                for m in range(217+len(dalist)):
                    Mfeats.append(0)
                ongoing=True
            elif line.count('.end method')==1:
                MethodSet.add(tuple(Mfeats))
                ongoing=False
            elif ongoing==True:
                bytecode = ""
                if not len(line.split()) == 0:
                    bytecode = line.split()[0]
                if bytecode in bytelist:
                    if bytecode == 'invoke-virtual':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-super':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-direct':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-static':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-interface':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-virtual/range':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-super/range':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-direct/range' == 1:
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-static/range':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    elif bytecode == 'invoke-interface/range':
                        for a in dalist:
                            try:
                                if line.split('->')[1].count(a) == 1:
                                    Mfeats[217+dalist.index(a)] +=1
                            except IndexError:
                                pass
                    Mfeats[bytelist.index(bytecode)] += 1                
        f.close()
    return (MethodSet)

class CollectWholeFeatures(threading.Thread):

    def __init__(self, meths, apklist, dalist, bytelist, locklist, flog) :
        threading.Thread.__init__(self)
        self.meths = meths
        self.locklist = locklist
        self.apklist = apklist
        self.dalist = dalist
        self.bytelist = bytelist
        self.flog = flog

    def run(self):
        for apk in self.apklist:

            if apk[0] == 'g': ######
                continue
                
            print("stage 1:  "+apk+' start')
            dirname = taskParam_rawdatadirname+apk
            if not os.path.exists(dirname+'/'+apk+'_smali'):
                continue
            smalidirs = os.listdir(dirname+'/'+apk+'_smali')
            for d in smalidirs:
                smalidir = dirname+'/'+apk+'_smali'+'/'+d
                if not os.path.isdir(smalidir) or smalidir=='.' or smalidir == '/':
                    continue
                Methods = set()
                searchToEXmethods(smalidir, Methods, self.dalist, self.bytelist)
                for m in Methods:
                    self.meths.add(str(m))
                self.meths.update(Methods)
            self.flog.write(apk)
            self.flog.write('\n')
            self.flog.flush()

def getBytecodelist(filename, bytecodelist):
    f = open(filename, 'r')
    for line in f.readlines():
        temp = line[:-1]
        if temp[-1] == '\r':
            bytecode = temp[:-1]
        else:
            bytecode = temp
        bytecodelist.append(bytecode)
    f.close()

def getAPIList(filename , apilist):
    f = open(filename, 'r')
    for line in f.readlines():
        temp = line[:-1]
        if temp[-1] == '\r':
            api = temp[:-1]
        else:
            api = temp
        apilist.append(api)
    f.close()

############################## main ##############################
fapklist = open(taskParam_apklistfilename,'r') # get apk list
lines = fapklist.readlines()
fapklist.close()
noofapks = len(lines)
#seperate the total apks into the serveral group (for assigning the apks to the threads)
listOfapklist = list()
for m in range(0, taskParam_noThread+1):
    listOfapklist.append(list())  
order = 0
for line in lines:
    order = order+1
    listOfapklist[(order%taskParam_noThread)].append(line[:-1])

# get api list from the txt file
apilist = list()
getAPIList(taskParam_apilistfilename, apilist)

bytecodelist = list()
getBytecodelist(taskParam_bytecodelistfilename, bytecodelist)

tMethodSet = set() # total method feature set

loglist = list()
for m in range(0, taskParam_noThread+1):
    fhowmany = open('howmany'+str(m)+'.txt','w')
    loglist.append(fhowmany)


# thread creation (# of thread = taskParam_noThread)
locklist = list()
for m in range(0, taskParam_noThread+1):
    locklist.append(threading.Lock())
t = list()
for m in range(0, taskParam_noThread+1):
    t.append(CollectWholeFeatures(tMethodSet, listOfapklist[m], apilist, bytecodelist, locklist, loglist[m]))
#thread start
for m in range(0, taskParam_noThread+1):
    t[m].start()
#thread end
for m in range(0, taskParam_noThread+1):
    t[m].join()
#log all features
fmlog = open(taskParam_resultfilename,'a')
for mitem in  tMethodSet:
    fmlog.write(str(mitem))
fmlog.close()
############################## end ###############################