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程序设计有所帮助。