#!/usr/bin/python
# Greg Martin
# Threat Stream Inc. Copyright (c) 2013.  All rights reserved. 
# www.threatstream.com
#
import sys, os
import boto
from boto.s3.key import Key
import hashlib, magic, gzip

unique_string = 'ADD_YOUR_UNIQUE_STRING_HERE'

if len(sys.argv) > 1:
    arg = sys.argv[1]
else:
    print 'Usage: ./malware_s3_upload <filename>\n'
    sys.exit()

# get magic value
def get_magic(file):
    fm = magic.from_file(file)
    return fm

# Get hash digests
def get_hashes(file):
    md5 = hashlib.md5()
    sha256 = hashlib.sha256()
    with open(file,'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            md5.update(chunk)
            sha256.update(chunk)
    hashes = {'md5' : md5.hexdigest(), 'sha256' : sha256.hexdigest(),}
    return hashes


def upload_s3(filename,md5,sha256,filetype):
    try:
        # Connect and upload to S3 bucket
        conn = boto.connect_s3()
        bucket = conn.create_bucket(unique_string.lower()+'_malware')
        k = Key(bucket)
        # Set unique_id
        k.key = sha256
        k.set_metadata('md5_hash', md5)
        k.set_metadata('file_type', filetype)
        k.set_contents_from_filename(filename)
        return 'success'
    except:
        return 'upload to s3 failed, make sure your aws keys are set in ~/.boto'

def gzip_file(filename):
    f_in = open(filename, 'rb')
    f_out = gzip.open(filename+'.gz', 'wb')
    f_out.writelines(f_in)
    f_out.close()
    f_in.close()
    

def main():
    # main code
    if len(sys.argv) > 1:
        arg = sys.argv[1]
    else:
        print 'Usage: ./malware_s3_upload <filename>\n'
        sys.exit()

    # get file type and hashes before compression
    file_type = get_magic(arg)
    md5_hash = get_hashes(arg)['md5']
    sha256_hash = get_hashes(arg)['sha256']

    print arg
    print file_type
    print 'md5: '+md5_hash
    print 'sha256: '+sha256_hash

    # gzip file
    gzip_file(arg)

    # upload to s3
    print 'Uploading gzipped '+arg+' to s3:\n'
    print upload_s3(arg+'.gz',md5_hash,sha256_hash,file_type)

    # delete local gzipped file
    os.unlink(arg+'.gz')



if __name__ == "__main__":
    main()
