malware_manager.py

# !/usr/bin/env python
# encoding: utf-8
"""
malwarehouse.py

Created by Scott Roberts.
Copyright (c) 2012 TogaFoamParty Studios. All rights reserved.
"""
import utils
import db_controller
import sys
from malware_sample import MalwareSample
from argparse import ArgumentParser
import extensions.plugins
from extensions.plugin_base import get_plugin, PluginBase as plg

# Malwarehouse Options
config = utils.get_configuration("malwarehouse.cfg")
# Initialize our DB controller
db_controller = db_controller.MalwareDbController(config)
db_initialized = db_controller.initialize_db()

def check_prelim():
    """Initial setup code. Eventually this will set options."""
    directory_structure = utils.initialize_environment(config)
    return directory_structure

# Processes the malware sample
def malware_loader(malware_path, source, notes, tags):
    sample = MalwareSample(malware_path, config, source, notes, tags)
    return sample


def load_sample(args, source, notes, tags):
    sample = malware_loader(args.load, source, notes, tags)
    if sample:        
        # Process plugins
        for plugin in plg.plugins:
            plugin_name = plugin.get_name()
            if config.get(plugin_name, 'plugin') == "On":    
                plg.analysis_report_directory = sample.analysis_report_directory
                _plugin = get_plugin(plugin.get_name(), config)
                _plugin.analyze(sample)
                _plugin.create_report()
        if sample.malware_definition and db_initialized:
            db_controller.load_db(sample.malware_definition)
    else:
        print "Errors were encountered during analysis"
        
        
def delete_sample():
    if not db_initialized:
        print "Failed to initialize database\nPlease verify your database settings"
        return
    pass

def recent_samples(args):
    if not db_initialized:
        print "Failed to initialize database\nPlease verify your database settings"
        return
    args.recent = 5 if not args.recent else args.recent
    data = db_controller.recent(args.recent)
    for parsed_data in map(utils.parse_sqlite_result, data):
        print MalwareSample.summary(parsed_data)
        
def find_sample(args):
    if not db_initialized:
        print "Failed to initialize database\nPlease verify your database settings"
        return
    print "> Find called with %s." % (args.find)
    data = db_controller.find_sample(args.find)
    for parsed_data in map(utils.parse_sqlite_result, data):
        print MalwareSample.summary(parsed_data)

def main():
    if not check_prelim():
        sys.exit(1)
    
    parser = ArgumentParser()
  
    parser.add_argument("-r", "--recent",
                        action="store",
                        nargs='?',
                        default='5',
                        help="Display the newest samples (default: 5)") 
    
    parser.add_argument("-s", "--source",
                        action="store",
                        default=None,
                        help="Source of file")
    parser.add_argument("-t", "--tags",
                        action="store",
                        default=None,
                        help="Any characteristics of the malware")
    parser.add_argument("-n", "--notes",
                        action="store",
                        default="",
                        help="Notes about file")
    parser.add_argument("-f", "--find",
                        action="store",
                        default="",
                        help="Find a sample by name, md5, or sha256")
    parser.add_argument("-l", "--load",
                        action="store",
                        default="",
                        help="Load a malware sample for analysis")

    parser.add_argument("-d", "--delete",
                        action="store",
                        default="",
                        help="Delete a sample by name, md5, or sha256")

    args = parser.parse_args()
    cli_arguments = sys.argv
       
       
    # Set optional arguments
    tags = args.tags   if args.tags   else ""
    source = args.source if args.source else ""
    notes = args.notes if args.notes else ""
    
    # Process user commands
    if args.find:
        find_sample(args)
    elif "-r" in cli_arguments:
        recent_samples(args)
    elif args.delete:
        print "> [not implemented] Delete called with %s" % (args.delete)
    elif args.load:
        load_sample(args, source, notes, tags)
    return True
    
if __name__ == "__main__":
    main()