# coding:utf-8

import os
import zipfile
import hashlib
import subprocess
import re
import threadpool
from tempfile import NamedTemporaryFile
import threading


# Global configuration
GOODWARE_PATH = "../../mnt/storage/yanjie/drebin"
PICKED_PATH = "../../mnt/storage/yanjie/drebin_picked"
tmp_file = "tmpfile.tmp"
debug = 1


def getApkList(rootDir):
    """
    :param rootDir:  root directory of dataset
    :return: A filepath list of sample
    """
    filePath = []
    for parent, dirnames, filenames in os.walk(rootDir):
        # 三个参数：分别返回 1.父目录 2.所有文件夹名字（不含路径） 3.所有文件名字
        for filename in filenames:  # 输出文件信息
            if ".apk" in filename:
                file = os.path.join(parent, filename)
                filePath.append(file)

    return filePath

class Analysis:
    def __init__(self, path):
        self.dir = path
        self.max_jobs = 15
        self.lock = threading.Lock()
        self.total_pick = 0
        self.developer = {}
        self.dex_md5 = {}


    def process_one(self, args):
        apk = args
        try:
            # APK MD5
            apk_item = open(apk, 'rb')
            apk_md5 = hashlib.md5(apk_item.read()).hexdigest()
            apk_item.close()

            # DEX MD5
            dex_md5 = "None"
            z = zipfile.ZipFile(apk)
            if "classes.dex" in z.namelist():
                dex_item = z.open("classes.dex", 'r')
                dex_md5 = hashlib.md5(dex_item.read()).hexdigest()

            # SIG
            sign_md5 = "None"
            rsa_files = filter(lambda x: x.endswith("RSA"), z.namelist())
            for rsa_file in rsa_files:
                tmpfile = NamedTemporaryFile(delete=False)
                tmpfile.write(z.open(rsa_file, 'r').read())
                tmpfile.close()
                cmd = "keytool -printcert -file " + tmpfile.name
                output = subprocess.check_output(cmd, shell=True)

                result = re.findall(r'MD5:\s+([0-9A-Z:]+)', str(output))
                if result:
                    sign_md5 = result[0].strip()
            z.close()

            if (not dex_md5 == "None") and (not sign_md5 == "None"):
                self.lock.acquire()
                if (not dex_md5 in self.dex_md5) and (not sign_md5 in self.developer):
                    # cp
                    os.system("cp " + apk + " " + os.path.join(PICKED_PATH, apk_md5))
                    self.dex_md5[dex_md5] = 1
                    self.developer[sign_md5] = 1
                    self.total_pick += 1
                self.lock.release()

        except Exception as e:
            print(e, apk)
            return None


    def process(self, dir):
        apks = getApkList(dir)
        print("total files ", len(apks))
        args = [(apk) for apk in apks]
        pool = threadpool.ThreadPool(self.max_jobs)
        requests = threadpool.makeRequests(self.process_one, args)
        [pool.putRequest(req) for req in requests]
        pool.wait()

    def start(self):
        self.process(self.dir)



if __name__ == '__main__':
    analysis = Analysis(GOODWARE_PATH)
    analysis.start()
    print(analysis.total_pick)