import requests,time,pyzipper,os,argparse
from datetime import datetime, timedelta


apikey = 'get api key here https://bazaar.abuse.ch/api/'

malware_bazaar_file = 'malware_bazaar.csv'

# inputfile = sys.argv[1]
parser = argparse.ArgumentParser(description="Download malware samples from https://bazaar.abuse.ch/")
parser.add_argument('-d','--download',action="store_true",help="Download a new copy of the CSV to work with")
# parser.add_argument('-f','--file',help="Filename to read from in order to download malware samples")
parser.add_argument('-t','--threat',help="Download samples associated with the threat name you enter")
args = parser.parse_args()

def download_malware_bazaar():
	res = requests.get('https://bazaar.abuse.ch/export/csv/recent/')
	savefile = open(malware_bazaar_file, 'ab')
	for chunk in res.iter_content(100000):
		savefile.write(chunk)
	print('Successfully downloaded {}'.format(malware_bazaar_file))

def unzip_file(filename):
	with pyzipper.AESZipFile(filename) as zf:
		zf.pwd = b'infected'
		zf.extractall("./malware")
		print("{} downloaded and unzipped.".format(filename.strip('.zip')))


def analyze_data(threat):
	threatcount = 0
	matched_threats = []
	with open (malware_bazaar_file,'r') as data:
		for item in data:
			item = item.strip()
			if "#" in item:
				pass
			else:
				# Parse malware bazaar file
				item = item.replace("\"","")
				item = item.split(',')
				first_seen_utc = item[0]
				sha256 = item[1].strip()
				filename = item[5].strip()
				zip_filename = filename +'.zip'
				file_type_guess = item[6]
				mime_type = item[7]
				signature = item[8]
				clamav = item[9]
				vtpercent = item[10]
				_created = (time.strptime(first_seen_utc,'%Y-%m-%d %H:%M:%S'))
				created =(time.strftime("%Y-%m-%dT%H:%M:%S", _created))

				if threat.lower() in signature.lower():
					threatcount += 1
					matched_threats.append("{},{},{}".format(created,sha256,zip_filename))

		print("-"* 10)
		answer = input("{} entries match {}.\nDo you want to download all these samples? - [y/n <enter>]\n".format(threatcount,threat)).lower()

		if answer == 'y': # Download file
			for item in matched_threats:
				created,sha256,zip_filename = item.split(',')
				headers = { 'API-KEY': apikey }
				data = { 'query': 'get_file','sha256_hash':sha256, }
				response = requests.post('https://mb-api.abuse.ch/api/v1/', data=data, timeout=15, headers=headers, allow_redirects=True)
				open(zip_filename,'wb').write(response.content)

				# Unzip file
				try:
					unzip_file(zip_filename)
					os.remove(zip_filename)
				except:
					print("{} failed to unzip. Zip file saved".format(zip_filename))
		else:
			print("Exiting since you didn't press 'y'")

if args.download:
	malware_bazaar_data = download_malware_bazaar()	# use to download a complete list of URLs from urlhaus - do just once a day

if args.threat:
	analyze_data(args.threat)	# use to download files