Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
103 lines (87 sloc) 3.82 KB
#!/usr/bin/env python
# Copyright (C) 2015 Derek Morton
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
import sys
import os
import json
import hashlib
import argparse
from time import sleep
from virus_total_apis import PublicApi as VirusTotalPublicApi
from pymongo import MongoClient
# Define your VirusTotal API Key
API_KEY = ''
#Define your MongoDB connection deatils
MONGO_IP = ''
MONGO_PORT = 27017
MONGO_DB = ''
MONGO_COLLECTION = ''
#### No more changes needed past this point ####
#Initalize our MongoDB Connection
client = MongoClient(MONGO_IP, MONGO_PORT)
collection = client[MONGO_DB][MONGO_COLLECTION]
#Initalize the VirusTotal Function
vt = VirusTotalPublicApi(API_KEY)
#Init the list of files we're going to scan
file_list = []
#Function to walk the provided directory and and files to the list
def dirwalk(dirname):
for root, dirs, files in os.walk(dirname, topdown=False):
for name in files:
filepath=(os.path.join(root, name))
file_list.append(filepath)
# Parse command line args given to see if we're scanning a file, directory or both
parser = argparse.ArgumentParser(description='This is a script to submit samples to VirusTotal')
parser.add_argument('-f', action="store", dest="filename")
parser.add_argument('-d', action="store", dest="dirname")
parser_results = parser.parse_args()
if parser_results.dirname != None:
dirwalk(parser_results.dirname)
if parser_results.filename != None:
file_list.append(parser_results.filename)
# Loop through the list to scan our files
while file_list:
sample_file = file_list.pop()
print sample_file
f = open(sample_file, "rb")
sample_file_content = f.read()
f.close()
sample_file_md5 = hashlib.md5(sample_file_content).hexdigest()
if collection.find({'results.md5':sample_file_md5}).count() == 0:
vt_response = vt.get_file_report(sample_file_md5)
if vt_response['response_code'] == 200 and vt_response['results']['response_code'] == 1:
collection.insert(vt_response)
print "Sample " + sample_file_md5 + " added to database"
print
elif vt_response['response_code'] == 204:
print "Hit API Limit; cooling off for 15 seconds"
print "The standard VirusTotal API Key allows 4 requests per minute"
file_list.append(sample_file)
sleep(15)
print
else:
print 'Sample not found; submitting to VirusTotal'
vt_response = vt.scan_file(sample_file)
if vt_response['response_code'] == 200 and vt_response['results']['response_code'] == 1:
print "Sample submiited successfully; checking back later"
file_list.insert(0,sample_file)
print
else:
print "Sample submit failed; try again later"
print
else:
print "Sample " + sample_file_md5 + " already exists in database"
print