python实现封装得到virustotal扫描结果

本文实例讲述了python实现封装得到virustotal扫描结果的方法。分享给大家供大家参考。具体方法如下：

    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() 

希望本文所述对大家的Python程序设计有所帮助。