#!/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']

    # 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()