build-scripts/virus-scan.py

   1 ###
   2 ### Submit a sample to VirusTotal and check the result.
   3 ### Based on example code by Bryce Boe, downloaded from:
   4 ### http://www.bryceboe.com/2010/09/01/submitting-binaries-to-virustotal/
   5 ###
   6 ### Needs Python 2.6 or 2.7
   7 ###
   8 ### Usage:
   9 ###  - Put your VirusTotal API key in the file virus-scan.key in the current
  10 ###    working directory (obtain from http://www.virustotal.com/advanced.html#publicapi)
  11 ###  - Run PATH/TO/PYTHON virus-scan.py FILE_TO_SCAN
  12 ###
  13 ### Copyright 2010 Steven J. Murdoch <http://www.cl.cam.ac.uk/users/sjm217/>
  14 ### See LICENSE for licensing information
  15 ###
  16
  17 import hashlib, httplib, mimetypes, os, pprint, json, sys, urlparse
  18
  19 DEFAULT_TYPE = 'application/octet-stream'
  20
  21 REPORT_URL = 'https://www.virustotal.com/api/get_file_report.json'
  22 SCAN_URL = 'https://www.virustotal.com/api/scan_file.json'
  23
  24 API_KEY_FILE = 'virus-scan.key'
  25
  26 # The following function is modified from the snippet at:
  27 # http://code.activestate.com/recipes/146306/
  28 def encode_multipart_formdata(fields, files=()):
  29     """
  30     fields is a dictionary of name to value for regular form fields.
  31     files is a sequence of (name, filename, value) elements for data to be
  32     uploaded as files.
  33     Return (content_type, body) ready for httplib.HTTP instance
  34     """
  35     BOUNDARY = '----------ThIs_Is_tHe_bouNdaRY_$'
  36     CRLF = '\r\n'
  37     L = []
  38     for key, value in fields.items():
  39         L.append('--' + BOUNDARY)
  40         L.append('Content-Disposition: form-data; name="%s"' % key)
  41         L.append('')
  42         L.append(value)
  43     for (key, filename, value) in files:
  44         L.append('--' + BOUNDARY)
  45         L.append('Content-Disposition: form-data; name="%s"; filename="%s"' %
  46                  (key, filename))
  47         content_type = mimetypes.guess_type(filename)[0] or DEFAULT_TYPE
  48         L.append('Content-Type: %s' % content_type)
  49         L.append('')
  50         L.append(value)
  51     L.append('--' + BOUNDARY + '--')
  52     L.append('')
  53     body = CRLF.join(L)
  54     content_type = 'multipart/form-data; boundary=%s' % BOUNDARY
  55     return content_type, body
  56
  57 def post_multipart(url, fields, files=()):
  58     """
  59     url is the full to send the post request to.
  60     fields is a dictionary of name to value for regular form fields.
  61     files is a sequence of (name, filename, value) elements for data to be
  62     uploaded as files.
  63     Return body of http response.
  64     """
  65     content_type, data = encode_multipart_formdata(fields, files)
  66     url_parts = urlparse.urlparse(url)
  67     if url_parts.scheme == 'http':
  68         h = httplib.HTTPConnection(url_parts.netloc)
  69     elif url_parts.scheme == 'https':
  70         h = httplib.HTTPSConnection(url_parts.netloc)
  71     else:
  72         raise Exception('Unsupported URL scheme')
  73     path = urlparse.urlunparse(('', '') + url_parts[2:])
  74     h.request('POST', path, data, {'content-type':content_type})
  75     return h.getresponse().read()
  76
  77 def scan_file(filename, api_key):
  78     files = [('file', filename, open(filename, 'rb').read())]
  79     json_result = post_multipart(SCAN_URL, {'key':api_key}, files)
  80     return json.loads(json_result)
  81
  82 def get_report(filename, api_key):
  83     md5sum = hashlib.md5(open(filename, 'rb').read()).hexdigest()
  84     json_result = post_multipart(REPORT_URL, {'resource':md5sum, 'key':api_key})
  85     data = json.loads(json_result)
  86     if data['result'] != 1:
  87         print 'Result not found, submitting file.'
  88         data = scan_file(filename, api_key)
  89         if data['result'] == 1:
  90             print 'Submit successful.'
  91             print 'Please wait a few minutes and try again to receive report.'
  92             return 1
  93         else:
  94             print 'Submit failed.'
  95             pprint.pprint(data)
  96             return 1
  97     else:
  98         #pprint.pprint(data['report'])
  99         scan_date, result_dict = data['report']
 100         print "Scanned on:", scan_date
 101
 102         failures = 0
 103         for av_name, result in result_dict.items():
 104             if result != '':
 105                 failures += 1
 106                 print " %20s: %s"%(av_name, result)
 107         if not failures:
 108             print 'SUCCESS: no AV detection triggered'
 109             return 0
 110         else:
 111             print 'FAIL: %s AV detection(s)'%failures
 112             return 255
 113
 114 if __name__ == '__main__':
 115     if len(sys.argv) != 2:
 116         print 'Usage: %s filename' % sys.argv[0]
 117         sys.exit(1)
 118
 119     try:
 120         key_fh = open(API_KEY_FILE, "rt")
 121         api_key = key_fh.readline().strip()
 122         key_fh.close()
 123     except IOError, e:
 124         print 'Failed to open API key file %s: %s' % (API_KEY_FILE, e)
 125         sys.exit(1)
 126
 127     filename = sys.argv[1]
 128     if not os.path.isfile(filename):
 129         print '%s is not a valid file' % filename
 130         sys.exit(1)
 131
 132     exit_status = get_report(filename, api_key)
 133     sys.exit(exit_status)