#!/usr/bin/env python3
import os
import sys
import json
from optparse import OptionParser

from malwareconfig import fileparser
from malwareconfig.modules import __decoders__, __preprocessors__

__description__ = 'Malware Config Extractor'
__author__ = '@kevthehermit, https://techanarchy.net, https://malwareconfig.com'
__version__ = '1.0'
__date__ = '2019/10'


def banner():
    banner_text = '''
 __  __       _  ____             __ 
|  \/  | __ _| |/ ___|___  _ __  / _|
| |\/| |/ _` | | |   / _ \| '_ \| |_ 
| |  | | (_| | | |__| (_) | | | |  _|
|_|  |_|\__,_|_|\____\___/|_| |_|_| 

Malware Configuration Parser by @kevthehermit
'''
    print(banner_text)


def process_file(file_path, output_file):
    # Open and parse the file
    print("[+] Loading File: {0}".format(file_path))
    file_info = fileparser.FileParser(file_path=file_path)

    print("  [-] Found: {0}".format(file_info.malware_name))

    # First we preprocesss
    # Check for a packer we can unpack
    if file_info.malware_name in __preprocessors__:
        print("  [+] Running PreProcessor {0}".format(file_info.malware_name))
        module = __preprocessors__[file_info.malware_name]['obj']()
        module.set_file(file_info)
        module.pre_process()


    if file_info.malware_name in __decoders__:
        print("  [-] Running Decoder")
        module = __decoders__[file_info.malware_name]['obj']()
        module.set_file(file_info)
        module.get_config()
        conf = module.config

        if output_file:
            print("  [+] Writing config to file: {0}".format(output_file))
            with open(output_file, 'a') as out:
                out.write('"{0}","{1}","{2}"\n'.format(file_path, file_info.malware_name, conf))

        else:
            print("  [-] Config Output\n")
            json_config = json.dumps(conf, indent=4, sort_keys=True)
            print(json_config)
    else:
        print("[!] No Decoder or File is Packed")


if __name__ == "__main__":
    parser = OptionParser(usage='usage: %prog file / dir\n' + __description__, version='%prog ' + __version__)
    parser.add_option("-r", "--recursive", action='store_true', default=False, help="Recursive Mode")
    parser.add_option("-l", "--list", action="store_true", default=False, help="List Available Decoders")
    parser.add_option("-o", "--output", help="Output Config elements to file.")
    (options, args) = parser.parse_args()

    banner()

    if options.list:
        print("[+] Listing Decoders")
        for name in __decoders__.keys():
            print("  [-] Loaded: {0}".format(name))

        print("[+] Listing PreProcessors")
        for name in __preprocessors__.keys():
            print("  [-] Loaded: {0}".format(name))

        sys.exit()

    # We need at least one arg
    if len(args) < 1:
        print("[!] Not enough Arguments, Need at least file path\n")
        parser.print_help()
        sys.exit()

    # Check for file or dir
    is_file = os.path.isfile(args[0])
    is_dir = os.path.isdir(args[0])

    if options.output:
        output_file = options.output
    else:
        output_file = None


    if options.recursive:
        if not is_dir:
            print("[!] Recursive requires a directory not a file.\n")
            sys.exit()

        # Read all the things
        base_dir = args[0]
        for root, dirs, files in os.walk(base_dir):
            for file_name in files:
                file_path = os.path.join(root, file_name)
                try:
                    process_file(file_path, output_file)
                except Exception as e:
                    print("[!] Unable to process file {0}".format(file_path))

    else:
        if not is_file:
            print("[!] You did not provide a valid file.\n")
            sys.exit()

        # Read in the file.
        process_file(args[0], output_file)
