Advertisement
Guest User

Search VirusTotal. Added -g and -s options

a guest
Jun 4th, 2012
168
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 8.88 KB | None | 0 0
  1. #!/usr/bin/env python
  2.  
  3. __description__ = 'Program to search VirusTotal reports with search terms (MD5, SHA1, SHA256) found in the argument file'
  4. __author__ = 'Didier Stevens'
  5. __version__ = '0.0.4'
  6. __date__ = '2012/06/05'
  7.  
  8. """
  9.  
  10. Source code put in public domain by Didier Stevens, no Copyright
  11. https://DidierStevens.com
  12. Use at your own risk
  13.  
  14. History:
  15.  2012/04/25: start
  16.  2012/04/27: added serialization of reports
  17.  2012/05/23: emergency fix pkl init bug
  18.  2012/05/26: 0.0.3 added force option and key option; added environment variable; added requested field
  19.  2012/06/05: 0.0.4 added genhash and short options. By TecToDo
  20. Todo:
  21. """
  22.  
  23. import optparse
  24. import urllib
  25. import urllib2
  26. import time
  27. import sys
  28. import pickle
  29. import os
  30. import hashlib
  31.  
  32. try:
  33.     import simplejson
  34. except:
  35.     print('Missing simplejson Python module, please check if it is installed.')
  36.     exit()
  37.  
  38. VIRUSTOTAL_API2_KEY = ''
  39. VIRUSTOTAL_REPORT_URL = "https://www.virustotal.com/vtapi/v2/file/report"
  40.  
  41. PICKLE_FILE = 'virustotal-search.pkl'
  42.  
  43. def Serialize(object):
  44.     try:
  45.         fPickle = open(PICKLE_FILE, 'wb')
  46.     except:
  47.         return False
  48.     try:
  49.         pickle.dump(object, fPickle)
  50.     except:
  51.         return False
  52.     finally:
  53.         fPickle.close()
  54.     return True
  55.  
  56. def DeSerialize():
  57.     import os.path
  58.  
  59.     if os.path.isfile(PICKLE_FILE):
  60.         try:
  61.             fPickle = open(PICKLE_FILE, 'rb')
  62.         except:
  63.             return None
  64.         try:
  65.             object = pickle.load(fPickle)
  66.         except:
  67.             return None
  68.         finally:
  69.             fPickle.close()
  70.         return object
  71.     else:
  72.         return None
  73.  
  74. def Timestamp(epoch=None):
  75.     if epoch == None:
  76.         localTime = time.localtime()
  77.     else:
  78.         localTime = time.localtime(epoch)
  79.     return '%04d%02d%02d-%02d%02d%02d' % localTime[0:6]
  80.  
  81. class CSVLogger():
  82.     def __init__(self, prefix, headers, separator=';'):
  83.         self.separator = separator
  84.         self.filename = '%s-%s.csv' % (prefix, Timestamp())
  85.         self.f = open(self.filename, 'w')
  86.         self.f.write(self.separator.join(headers) + '\n')
  87.         self.f.close()
  88.  
  89.     def PrintAndLog(self, formats, parameters, short=False):
  90.         line = self.separator.join(formats) % parameters
  91.         if short == True:
  92.             result = line.split(';')
  93.             print result[0], result[4]+'/'+result[5], result[6]
  94.         else:
  95.             print(line)
  96.         f = open(self.filename, 'a')
  97.         f.write(line + '\n')
  98.         f.close()
  99.  
  100. def VTHTTPReportRequest(searchTerm):
  101.     req = urllib2.Request(VIRUSTOTAL_REPORT_URL, urllib.urlencode({'resource': searchTerm, 'apikey': VIRUSTOTAL_API2_KEY}))
  102.     try:
  103.         if sys.hexversion >= 0x020601F0:
  104.             hRequest = urllib2.urlopen(req, timeout=15)
  105.         else:
  106.             hRequest = urllib2.urlopen(req)
  107.     except:
  108.         return None
  109.     try:
  110.         data = hRequest.read()
  111.     except:
  112.         return None
  113.     finally:
  114.         hRequest.close()
  115.     return data
  116.  
  117. def InsertIntoTuple(tupleIn, position, value):
  118.     listIn = list(tupleIn)
  119.     listIn.insert(position, value)
  120.     return tuple(listIn)
  121.  
  122. def GetReport(searchTerm, reports, withComment, forceRequest, short):
  123.     global oLogger
  124.  
  125.     if withComment:
  126.         index = searchTerm.find(' ')
  127.         if index == -1:
  128.             comment = ''
  129.         else:
  130.             comment = searchTerm[index+1:]
  131.             searchTerm = searchTerm[:index]
  132.     if not forceRequest and searchTerm in reports:
  133.         issuedRequest = False
  134.         oResult = reports[searchTerm]
  135.     else:
  136.         jsonResponse = VTHTTPReportRequest(searchTerm)
  137.         issuedRequest = True
  138.         if jsonResponse == None:
  139.             formats = ('%s', '%s')
  140.             parameters = (searchTerm, 'Error')
  141.             if withComment:
  142.                 formats = InsertIntoTuple(formats, 1, '%s')
  143.                 parameters = InsertIntoTuple(parameters, 1, comment)
  144.             oLogger.PrintAndLog(formats, parameters)
  145.             return issuedRequest
  146.         else:
  147.             oResult = simplejson.loads(jsonResponse)
  148.             if oResult['response_code'] == 1:
  149.                 reports[searchTerm] = oResult
  150.     if oResult['response_code'] == 1:
  151.         scans = []
  152.         for scan in sorted(oResult['scans']):
  153.             if oResult['scans'][scan]['detected']:
  154.                 scans.append('#'.join((scan, oResult['scans'][scan]['result'], oResult['scans'][scan]['update'], oResult['scans'][scan]['version'])))
  155.         formats = ('%s', '%d', '%d', '%s', '%d', '%d', '%s', '%s')
  156.         parameters = (searchTerm, issuedRequest, oResult['response_code'], oResult['scan_date'], oResult['positives'], oResult['total'], oResult['permalink'], ','.join(scans))
  157.         if withComment:
  158.             formats = InsertIntoTuple(formats, 1, '%s')
  159.             parameters = InsertIntoTuple(parameters, 1, comment)
  160.         oLogger.PrintAndLog(formats, parameters, short)
  161.     else:
  162.         formats = ('%s', '%d', '%d', '%s')
  163.         parameters = (searchTerm, issuedRequest, oResult['response_code'], oResult['verbose_msg'])
  164.         if withComment:
  165.             formats = InsertIntoTuple(formats, 1, '%s')
  166.             parameters = InsertIntoTuple(parameters, 1, comment)
  167.         oLogger.PrintAndLog(formats, parameters, short)
  168.     return issuedRequest
  169.  
  170. def File2Strings(filename):
  171.     try:
  172.         f = open(filename, 'r')
  173.     except:
  174.         return None
  175.     try:
  176.         return map(lambda line:line.rstrip('\n'), f.readlines())
  177.     except:
  178.         return None
  179.     finally:
  180.         f.close()
  181.  
  182. def VirusTotalSearch(filename, options):
  183.     global oLogger
  184.  
  185.     if options.genhash: #if command executed with the -g or --genhash flag.
  186.         searchTerms = genhash(filename)
  187.         if searchTerms == None:
  188.             print('Error reading file %s' % filename)
  189.             return
  190.     else:
  191.         searchTerms = File2Strings(filename)
  192.         if searchTerms == None:
  193.             print('Error reading file %s' % filename)
  194.             return
  195.         elif searchTerms == []:
  196.             print('No searchterms in file %s' % filename)
  197.             return
  198.  
  199.     headers = ('Search Term', 'Requested', 'Response', 'Scan Date', 'Detections', 'Total', 'Permalink', 'AVs')
  200.     if options.comment:
  201.         headers = InsertIntoTuple(headers, 1, 'Comment')
  202.     oLogger = CSVLogger('virustotal-search', headers)
  203.  
  204.     data = DeSerialize()
  205.     if data == None:
  206.         reports = {}
  207.     else:
  208.         reports = data['reports']
  209.  
  210.     while searchTerms != []:
  211.         issuedRequest = GetReport(searchTerms[0], reports, options.comment, options.force, options.short)
  212.         searchTerms = searchTerms[1:]
  213.         if issuedRequest and searchTerms != []:
  214.             time.sleep(options.delay)
  215.     Serialize({'reports': reports})
  216.  
  217. def genhash(malfile):
  218.     hash = None
  219.     try:
  220.         hash = []
  221.         f = open(malfile,'rb')
  222.         m = hashlib.md5()
  223.         while True:
  224.             data = f.read(10240)
  225.             if len(data) == 0:
  226.                 break
  227.             m.update(data)
  228.         hash.append(m.hexdigest())
  229.     except:
  230.         None
  231.     return hash
  232.  
  233. def Main():
  234.     global VIRUSTOTAL_API2_KEY
  235.  
  236.     oParser = optparse.OptionParser(usage='usage: %prog [options] file\n' + __description__, version='%prog ' + __version__)
  237.     oParser.add_option('-d', '--delay', type=int, default=16, help='delay in seconds between queries (default 16s, VT rate limit is 4 queries per minute)')
  238.     oParser.add_option('-c', '--comment', action='store_true', default=False, help='the search term is followed by a comment and separated by a space character')
  239.     oParser.add_option('-f', '--force', action='store_true', default=False, help='force all request to be send to VirusTotal, even if found in local database (pkl file)')
  240.     oParser.add_option('-g', '--genhash', action='store_true', default=False, help='Generate (md5) hash from file passed as argument.')
  241.     oParser.add_option('-s', '--short', action='store_true', default=False, help='Shorten output, return: hash, score, date and url.')
  242.     oParser.add_option('-k', '--key', default='', help='VirusTotal API key')
  243.     (options, args) = oParser.parse_args()
  244.  
  245.     if len(args) != 1:
  246.         oParser.print_help()
  247.         print('')
  248.         print('  Source code put in the public domain by Didier Stevens, no Copyright')
  249.         print('  Use at your own risk')
  250.         print('  https://DidierStevens.com')
  251.         return
  252.     if os.getenv('VIRUSTOTAL_API2_KEY') != None:
  253.         VIRUSTOTAL_API2_KEY = os.getenv('VIRUSTOTAL_API2_KEY')
  254.     if options.key != '':
  255.         VIRUSTOTAL_API2_KEY = options.key
  256.     if VIRUSTOTAL_API2_KEY == '':
  257.         print('You need to get a VirusTotal API key and set environment variable VIRUSTOTAL_API2_KEY, use option -k or add it to this program.\nTo get your API key, you need a VirusTotal account.')
  258.     else:
  259.         VirusTotalSearch(args[0], options)
  260.  
  261. if __name__ == '__main__':
  262.     Main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement