apkcollector.py

import src.virustotalv3
import json
import sys
import getopt
import argparse
import psycopg2
#import os
#import logging
from db.config import config
import multi_label


# Global vars
debug_mode = False

def debug(msg):
    if debug_mode:
        print(msg, flush = True)
        #pwd = os.path.dirname(os.path.realpath(__file__))
        #logging.basicConfig(filename = '{0}/logs/apkcollector.log'.format(pwd), encoding = 'utf-8', level = logging.DEBUG)
        #logging.debug('Debug mode is activated')

def usage():
    print("Max arguments: 2")
    print("API Key (-k) and debug mode (-d)")
    print("Usage: " + sys.argv[0] + " [OPTIONS]")
    print("Example: " + sys.argv[0] + " " + \
    "-k [api key] -d [True/False]")
    return

def insert_scan_result(resp):
    sql = """INSERT INTO vt(hash_id, vhash, sha256, sha1, md5, meaningful_name, first_submission_date, last_submission_date, times_submitted, last_modification_date, 
                            last_analysis_date, last_analysis_stats, type_extension, tags, total_votes, links, scan_result)
             VALUES(%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) RETURNING hash_id;"""
    conn = None
    datavt = None
    try:
        params = config()
        conn = psycopg2.connect(**params)
        cur = conn.cursor()
        cur.execute(sql,    (resp['data']['id'],
                            resp['data']['attributes']['vhash'],
                            resp['data']['attributes']['sha256'],
                            resp['data']['attributes']['sha1'],
                            resp['data']['attributes']['md5'],
                            resp['data']['attributes']['meaningful_name'],
                            resp['data']['attributes']['first_submission_date'],
                            resp['data']['attributes']['last_submission_date'],
                            resp['data']['attributes']['times_submitted'],
                            resp['data']['attributes']['last_modification_date'],
                            resp['data']['attributes']['last_analysis_date'],
                            json.dumps(resp['data']['attributes']['last_analysis_stats']),
                            resp['data']['attributes']['type_extension'],
                            resp['data']['attributes']['tags'],
                            json.dumps(resp['data']['attributes']['total_votes']),
                            resp['data']['links']['self'],
                            json.dumps(resp['data']['attributes']['last_analysis_results'])
                            ))
        datavt = cur.fetchone()[0]
        conn.commit()
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
    finally:
        if conn is not None:
            conn.close()

    return datavt

def insert_apk_result(resp):
    sql = """INSERT INTO apk_file(hash_id, version_code, version_name, package)
             VALUES(%s,%s,%s,%s) RETURNING package;"""

    conn = None
    dataapk = None
    try:
        params = config()
        conn = psycopg2.connect(**params)
        cur = conn.cursor()
        # execute the INSERT statement
        cur.execute(sql,    (resp['data']['id'],
                            resp['data']['attributes']['androguard']['AndroidVersionCode'],
                            resp['data']['attributes']['androguard']['AndroidVersionName'],
                            resp['data']['attributes']['androguard']['Package']))
        dataapk = cur.fetchone()[0]
        conn.commit()
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
    finally:
        if conn is not None:
            conn.close()

    return dataapk

def delete_datafile(hash):
    #sql = """DELETE FROM file where hash_id='%s';"""
    #sql_repeat = """DELETE FROM file WHERE NOT EXISTS(SELECT FROM vt WHERE upper(vt.hash_id) = file.hash_id);"""

    conn = None
    try:
        params = config()
        conn = psycopg2.connect(**params)
        cur = conn.cursor()
        cur.execute(sql, (hash))
        conn.commit()
        cur.close()
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
    finally:
        if conn is not None:
            conn.close()

def apkcollector(api):

    f = open("apks_hashes/list_of_selected_sha256", "r")

    filehash = f.readline()
    while filehash != '':
        x = src.virustotalv3
        resp = x.vtotal(api, filehash.rstrip()) # Remove all kinds of trailing whitespace

        # STATUS_CODE != 200
        if (resp != None):
            insert_scan_result(resp)
            multi_label.insert_results(filehash, resp['data']['attributes']['last_analysis_results'])
            # ANDROGUARD is only available if the file is an apk
            if (resp['data']['attributes']['type_extension'] == "apk"):
                insert_apk_result(resp)

            debug(f"Testing debug. STATUS_CODE = '{resp}'.")
        #else:
        #    delete_datafile(filehash)

        filehash = f.readline()

    f.close()

def main():
    global debug_mode
    debug_mode = False

    try:
        opts, args = getopt.getopt(sys.argv[1:], "k:d:h", )
        if len(opts) == 1:
            if opts[0][0] == "-h":
                usage()
            elif opts[0][0] == "-d":
                print("No API key supplied")
                sys.exit(0)
            elif opts[0][0] == "-k":
                key = opts[0][1]
                print("Starting...")
                apkcollector(key)
        elif (len(opts) > 1) and (len(opts) < 3):
            for x, y in opts:
                if x == "-k":
                    key = y
                if x == "-d":
                    if y.lower() in ['true', 'y', 'yes', 'debug']:
                        debug_mode = True
            print("Starting...")
            apkcollector(key)
        else:
            print("Incorrect number of arguments")
            usage()
    except getopt.GetoptError as err:
        print(err)
        sys.exit(-1)


if __name__ == "__main__" :
    main()
    sys.exit(0)