api.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

'''
Copyright (c) 2012, Claudio "nex" Guarnieri
Copyright (c) 2013, Michael Boman <michael@michaelboman.org>

All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
this list of conditions and the following disclaimer in the documentation
and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

import sys
import argparse
import StringIO
import hashlib
import logging
import re
import datetime

from utils import Config, get_type, del_file, get_file, get_chunks, \
    jsonize

try:
    from interruptingcow import timeout
except ImportError:
    sys.exit('ERROR: interruptingcow library is missing')

try:
    import requests
except ImportError:
    sys.exit('ERROR: requests library is missing')

try:
    from pymongo import MongoClient
    import gridfs
except ImportError:
    sys.exit('ERROR: pymongo library is missing')

try:
    # from bottle import route, request
    from bottle import Bottle, route, run, request, server_names, \
        ServerAdapter, hook, response, HTTPError
except ImportError:
    sys.exit('ERROR: Bottle library is missing')

# INIT

logging.basicConfig(format='%(levelname) -10s %(asctime)s %(message)s',
                    level=logging.DEBUG)

client = MongoClient(host=Config().api.dbhost, port=Config().api.dbport)
db = client.vxcage
fs = gridfs.GridFS(db)

@route('/malware/addtags', method='POST')
def add_tags(key=None, tags=None):
    '''
    @key  : dict to identify what sample to update
    @tags : list of tags to add to the sample

    returns : JSON status message 
    '''

    if not key:
        key  = request.forms.get('key')
    if not tags:
        tags = request.forms.get('tags').split(',')

    if key and tags:
        logging.info('Adding tags: %s' % ','.join(tags))
        db.fs.files.update(key, {'$addToSet': {'tags': {'$each': tags}}})

        return jsonize({'message': 'added'})
    else:
        logging.info('No key (%s) or no tags (%s) to add' % (key, tags))


@route('/malware/add', method='POST')
def add_malware():
    '''
    Adds a sample to the repository. Performs hashing and filemagic
    analysis of the uploaded sample.

    @tags : comma seperated tags list
    @file : binary sample stream

    returns : JSON status message 
    '''

    try:
        with timeout(Config().api.timeout * 60, exception=RuntimeError):
            tags = request.forms.get('tags').split(',')
            data = request.files.file
            data.file.seek(0)

            filename = data.filename
            sampleData = data.file.read()

            logging.debug('[%s] Generating hashes' % sampleEntry)
            md5 = hashlib.md5(sampleData).hexdigest()
            sha1 = hashlib.sha1(sampleData).hexdigest()
            sha256 = hashlib.sha256(sampleData).hexdigest()
            sha512 = hashlib.sha512(sampleData).hexdigest()
            filetype = get_type(sampleData)

            key = {'md5': md5}

            logging.debug('Quering database for already existing file (hash=%s)'
                           % md5)
            existing = db.fs.files.find_one({'md5': md5})

            upload_sample = True
            if existing:
                logging.info('Sample already exists')
                logging.info('Verifying contents')
                if not md5 == existing['md5']:
                    logging.warning('Checksum not matching')
                    upload_sample = True
                else:
                    logging.info('Checksum matching')
                    upload_sample = False
            else:
                upload_sample = True

            if upload_sample:
                logging.debug('Uploading sample')
                new = fs.new_file(filename=filename, sha1=sha1,
                                  sha256=sha256, sha512=sha512,
                                  filetype=filetype)
                for chunk in get_chunks(sampleData):
                    logging.debug('writing chunk')
                    new.write(chunk)
                new.close()
                logging.info('Uploaded sample')

            add_tags(key=key, tags=tags)

            logging.debug('Reclaiming memory')
            del sampleData

            response.content_type = 'application/json'
            return jsonize({'message': 'added'})
    except RuntimeError:
        response.content_type = 'application/json'
        return (jsonize({'error': 'timeout'}), 504)


@route('/malware/get', method='POST')
@route('/malware/get/<filehash>', method='GET')
def get_malware(filehash):
    '''
    Retreive sample object by hash
    @md5 : md5 hash
    @sha1 : sha1 hash
    @sha256 : sha256 hash
    @sha512 : sha512 hash
    @filehash : any of the above hash methods, the method will try to
                identify the hash algorithm.
    returns : binary sample or a JSON status message 
    '''

    if not filehash:
        md5 = request.forms.get('md5')
        sha1 = request.forms.get('sha1')
        sha256 = request.forms.get('sha256')
        sha512 = request.forms.get('sha512')

        if md5 != None:
            filehash = md5
        if sha1 != None:
            filehash = sha1
        if sha256 != None:
            filehash = sha256
        if sha512 != None:
            filehash = sha512

    if re.findall(r"([a-fA-F\d]{32})", filehash):

        # MD5

        sampleData = get_file(db, md5=filehash)
    elif re.findall(r"([a-fA-F\d]{40})", filehash):

        # SHA1

        sampleData = get_file(db, sha1=filehash)
    elif re.findall(r"([a-fA-F\d]{64})", filehash):

        # SHA256

        sampleData = get_file(db, sha256=filehash)
    elif re.findall(r"([a-fA-F\d]{128})", filehash):

        # SHA512

        sampleData = get_file(db, sha512=filehash)
    else:

        # Hash not recognized

        response.content_type = 'application/json'
        return (jsonize({'error': 'unknown_hash'}), 400)

    if sampleData:
        response.content_type = 'application/octet-stream'
        return sampleData
    else:
        response.content_type = 'application/json'
        return (jsonize({'error': 'sample_not_found'}), 404)


@route('/malware/del', method='POST')
@route('/malware/del/<filehash>', method='GET')
def del_malware(filehash):
    '''
    Delete sample object by hash
    @md5 : md5 hash
    @sha1 : sha1 hash
    @sha256 : sha256 hash
    @sha512 : sha512 hash
    @filehash : any of the above hash methods, the method will try to
                identify the hash algorithm.
    returns : JSON status message 
    '''

    if not filehash:
        md5 = request.forms.get('md5')
        sha1 = request.forms.get('sha1')
        sha256 = request.forms.get('sha256')
        sha512 = request.forms.get('sha512')

        if md5 != None:
            filehash = md5
        if sha1 != None:
            filehash = sha1
        if sha256 != None:
            filehash = sha256
        if sha512 != None:
            filehash = sha512

    if re.findall(r"([a-fA-F\d]{32})", filehash):

        # MD5

        sampleData = del_file(db, md5=filehash)
    elif re.findall(r"([a-fA-F\d]{40})", filehash):

        # SHA1

        sampleData = del_file(db, sha1=filehash)
    elif re.findall(r"([a-fA-F\d]{64})", filehash):

        # SHA256

        sampleData = del_file(db, sha256=filehash)
    elif re.findall(r"([a-fA-F\d]{128})", filehash):

        # SHA512

        sampleData = del_file(db, sha512=filehash)
    else:

        # Hash not recognized

        response.content_type = 'application/json'
        return (jsonize({'error': 'unknown_hash'}), 400)

    if sampleData:
        response.content_type = 'application/json'
        return (jsonize({'ok': 'sample_deleted'}), 200)
    else:
        response.content_type = 'application/json'
        return (jsonize({'error': 'sample_not_found'}), 404)


@route('/malware/dbsearch', method='POST')
def search_malware():
    '''
    search the database using user-supplied key
    returns : JSON data or JSON status message
    '''

    dblist = list()
    for (key, val) in request.forms.iteritems():
        logging.debug('Looking for %s : %s' % (key, val))
        dblist.extend(db.fs.files.find({key: {'$regex': val}}))

    logging.debug('Found %s items in total' % len(dblist))

    for index in xrange(len(dblist)):
        del dblist[index]['_id']
        dblist[index]['uploadDate'] = str(dblist[index]['uploadDate'])

    return jsonize(dblist)


@route('/malware/find', method='POST')
@route('/malware/find/<filehash>', method='GET')
def find_malware(
    md5=None,
    sha1=None,
    sha256=None,
    sha512=None,
    ssdeep=None,
    tag=None,
    created=None,
    filehash='',
    ):
    '''
    search the database using hash, tag or date uploaded
    returns : JSON metadata or JSON status message
    '''

    if not filehash:
        md5 = request.forms.get('md5')
        sha1 = request.forms.get('sha1')
        sha256 = request.forms.get('sha256')
        sha512 = request.forms.get('sha512')
        if md5 != None:
            filehash = md5
        if sha1 != None:
            filehash = sha1
        if sha256 != None:
            filehash = sha256
        if sha512 != None:
            filehash = sha512
        if filehash == None:
            filehash = ''

    ssdeep = request.forms.get('ssdeep')
    tag = request.forms.get('tag')
    created = request.forms.get('created')
    filehash = request.forms.get('filehash')

    if re.search(r"([a-fA-F\d]{128})", filehash):

        # SHA512

        logging.debug('Looking for SHA512 hash %s' % filehash)
        metadata = db.fs.files.find_one({'sha512': filehash})
    elif re.search(r"([a-fA-F\d]{64})", filehash):

        # SHA256

        logging.debug('Looking for SHA256 hash %s' % filehash)
        metadata = db.fs.files.find_one({'sha256': filehash})
    elif re.search(r"([a-fA-F\d]{40})", filehash):

        # SHA1

        logging.debug('Looking for SHA1 hash %s' % filehash)
        metadata = db.fs.files.find_one({'sha1': filehash})
    elif re.search(r"([a-fA-F\d]{32})", filehash):

        # MD5

        logging.debug('Looking for MD5 hash %s' % filehash)
        metadata = db.fs.files.find_one({'md5': filehash})
    else:
        if ssdeep:
            metadata = db.fs.files.find({'ssdeep': {'$regex': '.*'
                    + ssdeep + '.*'}})
        elif tag:
            metadata = db.fs.files.find({'tags': tag})
        elif created:
            start = datetime.datetime.strptime(created, '%Y-%m-%d')
            end = datetime.datetime.strptime(created, '%Y-%m-%d') \
                + datetime.timedelta(days=1)
            metadata = db.fs.files.find({'uploadDate': {'$gte': start,
                    '$lt': end}})
        else:
            response.content_type = 'application/json'
            return jsonize({'error': 'invalid_search_term'})

    if not metadata:
        response.content_type = 'application/json'
        return (jsonize({'error': 'file_not_found'}), 404)

    response.content_type = 'application/json'
    del metadata['_id']
    metadata['uploadDate'] = str(metadata['uploadDate'])
    return jsonize(metadata)


@route('/tags/list', method='GET')
def list_tags():
    '''
    List all tags
    returns : list of tags in JSON format
    '''

    response.content_type = 'application/json'
    return jsonize(db.fs.files.distinct('tags'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-H',
        '--host',
        help='Host to bind the API server on',
        default='localhost',
        action='store',
        required=False,
        )
    parser.add_argument(
        '-p',
        '--port',
        help='Port to bind the API server on',
        default=8090,
        action='store',
        required=False,
        )
    args = parser.parse_args()

    run(host=args.host, port=args.port)