[or-cvs] [torbrowser/master 2/4] Update script to submit samples to VirusTotal

[erinn@xxxxxxxxxxxxxx]

Author: Steven Murdoch <Steven.Murdoch@xxxxxxxxxxxx>
Date: Sun, 24 Oct 2010 22:48:51 +0100
Subject: Update script to submit samples to VirusTotal
Commit: e3fea2285fc867bb507410b546c4987dacb5b234

- Add copyright header
- Take API key from a file rather than being hard-coded
- Use in-built json module rather than simplejson
- Print out a summary rather than raw JSON object
---
 build-scripts/virus-scan.py |   62 +++++++++++++++++++++++++++++++++++-------
 1 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/build-scripts/virus-scan.py b/build-scripts/virus-scan.py
index 3b14fcb..a585321 100644
--- a/build-scripts/virus-scan.py
+++ b/build-scripts/virus-scan.py
@@ -1,11 +1,27 @@
-import hashlib, httplib, mimetypes, os, pprint, simplejson, sys, urlparse
+###
+### Submit a sample to VirusTotal and check the result.
+### Based on example code by Bryce Boe, downloaded from:
+### http://www.bryceboe.com/2010/09/01/submitting-binaries-to-virustotal/
+###
+### Needs Python 2.6 or 2.7
+###
+### Usage:
+###  - Put your VirusTotal API key in the file virus-scan.key in the current
+###    working directory (obtain from http://www.virustotal.com/advanced.html#publicapi)
+###  - Run PATH/TO/PYTHON virus-scan.py FILE_TO_SCAN
+###
+### Copyright 2010 Steven J. Murdoch <http://www.cl.cam.ac.uk/users/sjm217/>
+### See LICENSE for licensing information
+###
+
+import hashlib, httplib, mimetypes, os, pprint, json, sys, urlparse
 
 DEFAULT_TYPE = 'application/octet-stream'
 
 REPORT_URL = 'https://www.virustotal.com/api/get_file_report.json'
 SCAN_URL = 'https://www.virustotal.com/api/scan_file.json'
 
-API_KEY = 'YOUR KEY HERE'
+API_KEY_FILE = 'virus-scan.key'
 
 # The following function is modified from the snippet at:
 # http://code.activestate.com/recipes/146306/
@@ -58,36 +74,60 @@ def post_multipart(url, fields, files=()):
     h.request('POST', path, data, {'content-type':content_type})
     return h.getresponse().read()
 
-def scan_file(filename):
+def scan_file(filename, api_key):
     files = [('file', filename, open(filename, 'rb').read())]
-    json = post_multipart(SCAN_URL, {'key':API_KEY}, files)
-    return simplejson.loads(json)
+    json_result = post_multipart(SCAN_URL, {'key':api_key}, files)
+    return json.loads(json_result)
 
-def get_report(filename):
+def get_report(filename, api_key):
     md5sum = hashlib.md5(open(filename, 'rb').read()).hexdigest()
-    json = post_multipart(REPORT_URL, {'resource':md5sum, 'key':API_KEY})
-    data = simplejson.loads(json)
+    json_result = post_multipart(REPORT_URL, {'resource':md5sum, 'key':api_key})
+    data = json.loads(json_result)
     if data['result'] != 1:
         print 'Result not found, submitting file.'
-        data = scan_file(filename)
+        data = scan_file(filename, api_key)
         if data['result'] == 1:
             print 'Submit successful.'
             print 'Please wait a few minutes and try again to receive report.'
+            return 1
         else:
             print 'Submit failed.'
             pprint.pprint(data)
+            return 1
     else:
-        pprint.pprint(data['report'])
+        #pprint.pprint(data['report'])
+        scan_date, result_dict = data['report']
+        print "Scanned on:", scan_date
 
+        failures = 0
+        for av_name, result in result_dict.items():
+            if result != '':
+                failures += 1
+                print " %20s: %s"%(av_name, result)
+        if not failures:
+            print 'SUCCESS: no AV detection triggered'
+            return 0
+        else:
+            print 'FAIL: %s AV detection(s)'%failures
+            return 255
 
 if __name__ == '__main__':
     if len(sys.argv) != 2:
         print 'Usage: %s filename' % sys.argv[0]
         sys.exit(1)
 
+    try:
+        key_fh = open(API_KEY_FILE, "rt")
+        api_key = key_fh.readline().strip()
+        key_fh.close()
+    except IOError, e:
+        print 'Failed to open API key file %s: %s' % (API_KEY_FILE, e)
+        sys.exit(1)
+
     filename = sys.argv[1]
     if not os.path.isfile(filename):
         print '%s is not a valid file' % filename
         sys.exit(1)
 
-    get_report(filename)
+    exit_status = get_report(filename, api_key)
+    sys.exit(exit_status)
-- 
1.7.1