# Copyright 2004-present Facebook. All Rights Reserved.

from __future__ import print_function

import argparse
import base64
import cStringIO
import hashlib
import json
import os
import sys
import time
import zipfile

from urllib import urlencode
from urllib2 import urlopen
from urlparse import urlunparse


def get_app_creds():
    """Reads your App ID and App Secret from env variables"""
    return {
        'AppID': os.environ['TX_APP_ID'],
        'AppSecret': os.environ['TX_APP_SECRET'],
    }


def get_parser():
    parser = argparse.ArgumentParser(
        description='Query ThreatExchange for malware samples'
    )
    # sample type?
    parser.add_argument('-x', '--hex', type=str)
    parser.add_argument('-i', '--input')
    parser.add_argument('-o', '--output')
    parser.add_argument('-s', '--since')
    parser.add_argument('-u', '--until')
    return parser


def get_query(creds, hexHash):
    # TODO verify hexHash is hex and not something else
    if len(hexHash) == 32:
        hash_type = 'MD5'
    elif len(hexHash) == 40:
        hash_type = 'SHA1'
    elif len(hexHash) == 64:
        hash_type = 'SHA256'
    else:
        raise Exception('Unrecognized hash type for "%s"' % hexHash)
    access_token = '%s|%s' % (creds['AppID'], creds['AppSecret'])
    query = urlencode({
        'access_token': access_token,
        'query_type': 'HASH',
        'hash_type': hash_type,
        'hash_value': hexHash,
        'fields': 'password,sample',
    })
    return urlunparse((
        'https',
        'graph.facebook.com',
        '/' + creds['AppID'] + '/malware_analyses',
        '',
        query,
        ''
    ))


def run_query(url):
    try:
        # the response is a 'data' object, which is a list of objects
        resp = json.loads(urlopen(url).read())['data'][0]
    except Exception as e:
        print('Failed to load md5: %s' % str(e))
        sys.exit(-1)
    return (resp['sample'], resp['password'])


def process_results(outdir, pass_zipped_content, password):
    zipfilehandle = cStringIO.StringIO()
    zipfilehandle.write(base64.b64decode(pass_zipped_content))
    with zipfile.ZipFile(zipfilehandle, 'r') as zf:
        entry = zf.infolist()[0]  # There should only ever be one entry
        try:
          hashobj = hashlib.new('md5', zf.read(entry.filename, password))
        except RuntimeError:
          print('Bad password for %s' % entry.filename)
        filehash = hashobj.hexdigest()
        outfilename = os.path.join(outdir, "%s_%s" % (filehash, entry.filename))
        with open(outfilename, 'w') as f:
            print('Writing to %s' % outfilename)
            f.write(zf.read(entry.filename, password))

if __name__ == '__main__':
    creds = get_app_creds()
    args = get_parser().parse_args()
    # inputs
    if args.input is None and args.hex is None:
        print('Hex hash or input file is needed')
        sys.exit(-1)
    elif args.input is not None and args.hex is not None:
        print('Can only take hex hash or input file')
        sys.exit(-1)
    elif args.input is not None:
        query_hashes = []
        for line in open(args.input, 'r'):
            query_hashes.append(line.strip())
    elif args.hex is not None:
        query_hashes = [args.hex]
    # output directories
    if args.output is not None:
        outdir = args.output
    else:
        outdir = '.'
    if not os.path.exists(outdir):
        print('%s does not exist' % outdir)
        sys.exit(-1)
    if not os.path.isdir(outdir):
        print('%s is not a directory' % outdir)
        sys.exit(-1)
    # time based query
    if args.since is not None:
        print('--since not yet implemented')
        sys.exit(-1)
    if args.until is not None:
        print('--until not yet implemented')
        sys.exit(-1)
    # main functionality
    start = int(time.time())
    for hexHash in query_hashes:
        # TODO how to do this for possibly several hashes
        url = get_query(creds, hexHash)
        results = run_query(url)
        process_results(outdir, results[0], results[1])
    end = int(time.time())
    if len(query_hashes) > 1:
        print('Total time elapsed: %d seconds' % (end - start))

@@ -119,16 +119,34 @@
  base64%0A
+import contextlib%0A
 import c
@@ -2138,16 +2138,35 @@
    with 
+contextlib.closing(
 zipfile.
@@ -2192,16 +2192,17 @@
 le, 'r')
+)
  as zf:%0A

Allow get_samples.py to run on Python 2.6