python-封装得到virustotal扫描结果

最新推荐文章于 2024-06-11 09:53:35 发布

小驹

最新推荐文章于 2024-06-11 09:53:35 发布

阅读量2k

点赞数

分类专栏： python 文章标签： microsoft sqlite database file null import

本文链接：https://blog.csdn.net/xiaocaiju/article/details/7617312

版权

python 专栏收录该内容

86 篇文章 0 订阅

订阅专栏

import simplejson
import urllib
import urllib2
import os, sys
import logging

try:
    import sqlite3
except ImportError:
    sys.stderr.write("ERROR: Unable to locate Python SQLite3 module. " \
                     "Please verify your installation. Exiting...\n")
    sys.exit(-1)
    
MD5 = "5248f774d2ee0a10936d0b1dc89107f1"
MD5 = "12fa5fb74201d9b6a14f63fbf9a81ff6"  #do not have report on virustotal.com
          

APIKEY = "xxxxxxxxxxxxxxxxxx"用自己的
########################################################################
class VirusTotalDatabase:
    """
    Database abstraction layer.
    """
    def __init__(self, db_file):
        log = logging.getLogger("Database.Init")
        self.__dbfile = db_file
        self._conn = None
        self._cursor = None

        # Check if SQLite database already exists. If it doesn't exist I invoke
        # the generation procedure.
        if not os.path.exists(self.__dbfile):
            if self._generate():
                print("Generated database \"%s\" which didn't" \
                         " exist before." % self.__dbfile)
            else:
                print("Unable to generate database")

        # Once the database is generated of it already has been, I can
        # initialize the connection.
        try:
            self._conn = sqlite3.connect(self.__dbfile)
            self._cursor = self._conn.cursor()
        except Exception, why:
            print("Unable to connect to database \"%s\": %s."
                      % (self.__dbfile, why))

        log.debug("Connected to SQLite database \"%s\"." % self.__dbfile)

    def _generate(self):
        """
        Creates database structure in a SQLite file.
        """
        if os.path.exists(self.__dbfile):
            return False

        db_dir = os.path.dirname(self.__dbfile)
        if not os.path.exists(db_dir):
            try:
                os.makedirs(db_dir)
            except (IOError, os.error), why:
                print("Something went wrong while creating database " \
                          "directory \"%s\": %s" % (db_dir, why))
                return False

        conn = sqlite3.connect(self.__dbfile)
        cursor = conn.cursor()

        cursor.execute("CREATE TABLE virustotal (\n"                           \
                       "  id INTEGER PRIMARY KEY,\n"                        \
                       "  md5 TEXT NOT NULL,\n"                     \
                       "  Kaspersky  TEXT DEFAULT NULL,\n"                             \
                       "  McAfee TEXT DEFAULT NULL,\n"                       \
                       "  Symantec TEXT DEFAULT NULL,\n"                         \
                       "  Norman TEXT DEFAULT NULL,\n"                         \
                       "  Avast TEXT DEFAULT NULL,\n"                        \
                       "  NOD32 TEXT DEFAULT NULL,\n"                  \
                       "  BitDefender TEXT DEFAULT NULL,\n"                        \
                       "  Microsoft TEXT DEFAULT NULL,\n"                        \
                       "  Rising TEXT DEFAULT NULL,\n"                     \
                       "  Panda TEXT DEFAULT NULL\n"                      \
                       ");")
        print "create db:%s sucess" % self.__dbfile

        return True

    def _get_task_dict(self, row):
        try:
            task = {}
            task["id"] = row[0]
            task["md5"] = row[1]
            task["Kaspersky"] = row[2]
            task["McAfee"] = row[3]
            task["Symantec"] = row[4]
            task["Norman"] = row[5]
            task["Avast"] = row[6]
            task["NOD32"] = row[7]
            task["BitDefender"] = row[8]
            task["Microsoft"] = row[9]
            task["Rising"] = row[10]
            task["Panda"] = row[11]
            return task
        except Exception, why:
            return None

    def add_sample(self, md5, virus_dict):
        """
        
        """
        task_id = None

        if not self._cursor:
            return None
        if not md5 or md5 == "":
            return None

        Kaspersky = virus_dict.get("Kaspersky", None)
        McAfee = virus_dict.get("McAfee", None)
        Symantec = virus_dict.get("Symantec", None)
        Norman = virus_dict.get("Norman", None)
        Avast = virus_dict.get("Avast", None)
        NOD32 = virus_dict.get("NOD32", None)
        BitDefender = virus_dict.get("BitDefender", None)
        Microsoft = virus_dict.get("Microsoft", None)
        Rising = virus_dict.get("Rising", None)
        Panda = virus_dict.get("Panda", None)
        
        self._conn.text_factory = str
        try:
            self._cursor.execute("SELECT id FROM virustotal WHERE md5 = ?;",
                                 (md5,))
            sample_row = self._cursor.fetchone()
        except sqlite3.OperationalError, why:
            print "sqlite3 error:%s\n" % str(why)
            return False
        
        if sample_row:
            try:
                sample_row = sample_row[0]
                self._cursor.execute("UPDATE virustotal SET Kaspersky=?, McAfee=?, Symantec=?, Norman=?, Avast=?, \
                                      NOD32=?, BitDefender=?, Microsoft=?, Rising=?, Panda=?   WHERE id = ?;",
                                     (Kaspersky, McAfee, Symantec, Norman, Avast, NOD32, BitDefender, Microsoft,\
                                      Rising, Panda, sample_row))
                self._conn.commit()
                task_id = sample_row
            except sqlite3.OperationalError, why:
                print("Unable to update database: %s." % why)
                return False
        else: #the sample not in the database
            try:
                self._cursor.execute("INSERT INTO virustotal " \
                                     "(md5, Kaspersky, McAfee, Symantec, Norman, Avast, NOD32, BitDefender,\
                                       Microsoft, Rising, Panda) " \
                                     "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);",
                                     (md5, Kaspersky, McAfee, Symantec, Norman, Avast, NOD32, BitDefender,\
                                       Microsoft, Rising, Panda))
                self._conn.commit()
                task_id = self._cursor.lastrowid
            except sqlite3.OperationalError, why:
                print "why",str(why)
                return None
            print "add_to_db:%s, task_id:%s" % (str(self.__dbfile), str(task_id))
        return task_id

    def get_sample(self):
        """
        Gets a task from pending queue.
        """
        log = logging.getLogger("Database.GetTask")

        if not self._cursor:
            log.error("Unable to acquire cursor.")
            return None

        # Select one item from the queue table with higher priority and older
        # addition date which has not already been processed.
        try:        
            self._cursor.execute("SELECT * FROM virustotal " \
                                 #"WHERE lock = 0 " \
                                 #"AND status = 0 " \
                                 "ORDER BY id, added_on LIMIT 1;")
        except sqlite3.OperationalError, why:
            log.error("Unable to query database: %s." % why)
            return None

        sample_row = self._cursor.fetchone()

        if sample_row:
            return self._get_task_dict(sample_row)
        else:
            return None

    def search_md5(self, md5):
        """
      
        """
        if not self._cursor:
            return None

        if not md5 or len(md5) != 32:
            return None

        try:
            self._cursor.execute("SELECT * FROM virustotal " \
                                 "WHERE md5 = ? " \
                                 #"AND status = 1 " \
                                 "ORDER BY id DESC;",
                                 (md5,))
        except sqlite3.OperationalError, why:
            return None

        task_dict = {}
        for row in self._cursor.fetchall():
            task_dict = self._get_task_dict(row)
            #if task_dict:
                #tasks.append(task_dict)

        return task_dict

    

class VirusTotal:
    """"""

    #----------------------------------------------------------------------
    def __init__(self, md5):
        """Constructor"""
        self._virus_dict = {}
        self._md5 = md5
        self._db_file = r"./db/virustotal.db"
        self.get_report_dict()
        
    def repr(self):
        return str(self._virus_dict)
    
    def submit_md5(self, file_path):
        import postfile                                                                          
        #submit the file
        FILE_NAME =  os.path.basename(file_path) 
                           
                                                                                                 
        host = "www.virustotal.com"                                                              
        selector = "https://www.virustotal.com/vtapi/v2/file/scan"                               
        fields = [("apikey", APIKEY)]
        file_to_send = open(file_path, "rb").read()                                              
        files = [("file", FILE_NAME, file_to_send)]                                              
        json = postfile.post_multipart(host, selector, fields, files)                            
        print json
        pass
    
    def get_report_dict(self):
        result_dict = {}
        
        url = "https://www.virustotal.com/vtapi/v2/file/report"
        parameters = {"resource": self._md5,
                       "apikey": APIKEY}
        data = urllib.urlencode(parameters)
        req = urllib2.Request(url, data)
        response = urllib2.urlopen(req)
        json = response.read()
        
        response_dict = simplejson.loads(json)
        if response_dict["response_code"]: #has result 
            scans_dict = response_dict.get("scans", {})
            for anti_virus_comany, virus_name in scans_dict.iteritems():
                if virus_name["detected"]:
                    result_dict.setdefault(anti_virus_comany, virus_name["result"])
        return result_dict
    
    #----------------------------------------------------------------------
    def write_to_db(self):
        """"""
        db = VirusTotalDatabase(self._db_file)
        virus_dict = self.get_report_dict()
        db.add_sample(self._md5, virus_dict)

使用时：

config = {'input':"inputMd5s"}
fp = open(config['input'], "r")
content = fp.readlines()
MD5S = []
for md5 in ifilter(lambda x:len(x)>0, imap(string.strip, content)):
    MD5S.append(md5)   
print "MD5S",MD5S
fp.close()


from getVirusTotalInfo import VirusTotal
#得到扫描结果并写入数库
for md5 in MD5S:
    virus_total = VirusTotal(md5)
    virus_total.write_to_db()