Merge pull request #201 from mitre/feature/config-refactor

Feature/config refactor

docker_utils/api_config.ini

@@ -16,6 +16,7 @@ password =
 vhost = /
 flush_every = 100
 flush_interval = 10
+tz = US/Eastern
 
 [Database]
 db_type = sqlite

multiscanner/__init__.py

@@ -4,12 +4,12 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from .config import (  # noqa F401
-    PY3, MS_WD, CONFIG, MODULESDIR
+    MSConfigParser, MS_WD, PY3, config_init, update_ms_config, update_ms_config_file
 )
 
 from .ms import (  # noqa F401
-    config_init, multiscan, parse_reports, _ModuleInterface,
-    _GlobalModuleInterface, _Thread, _run_module, _main
+    multiscan, parse_reports, _ModuleInterface,
+    _GlobalModuleInterface, _Thread, _run_module, _main, _get_main_modules
 )
 
 from .version import __version__  # noqa F401

multiscanner/analytics/ssdeep_analytics.py

@@ -22,7 +22,6 @@
 '''
 
 import argparse
-import configparser
 import json
 import logging
 import sys
@@ -37,20 +36,13 @@
     ssdeep = False
 
 
-from multiscanner import CONFIG as MS_CONFIG
-from multiscanner.common import utils
 from multiscanner.storage import storage
 
 
 class SSDeepAnalytic:
 
     def __init__(self):
-        storage_conf = utils.get_config_path(MS_CONFIG, 'storage')
-        config_object = configparser.ConfigParser()
-        config_object.optionxform = str
-        config_object.read(storage_conf)
-        conf = utils.parse_config(config_object)
-        storage_handler = storage.StorageHandler(configfile=storage_conf)
+        storage_handler = storage.StorageHandler()
         es_handler = storage_handler.load_required_module('ElasticSearchStorage')
 
         if not es_handler:
@@ -59,7 +51,7 @@ def __init__(self):
 
         # probably not ideal...
         self.es = es_handler.es
-        self.index = conf['ElasticSearchStorage']['index']
+        self.index = es_handler.index
         self.doc_type = '_doc'
 
     def ssdeep_compare(self):

multiscanner/common/dir_monitor.py

@@ -19,9 +19,8 @@
 from watchdog.events import FileSystemEventHandler
 from watchdog.observers import Observer
 
-from multiscanner import CONFIG as MS_CONFIG
 from multiscanner import multiscan, parse_reports
-from multiscanner.common import utils
+from multiscanner import config as msconf
 from multiscanner.storage import storage
 
 logger = logging.getLogger(__name__)
@@ -81,8 +80,7 @@ def start_observer(directory, work_queue, recursive=False):
 def multiscanner_process(work_queue, config, batch_size, wait_seconds, delete, exit_signal):
     filelist = []
     time_stamp = None
-    storage_conf = utils.get_config_path(config, 'storage')
-    storage_handler = storage.StorageHandler(configfile=storage_conf)
+    storage_handler = storage.StorageHandler()
     while not exit_signal.value:
         time.sleep(1)
         try:
@@ -102,7 +100,7 @@ def multiscanner_process(work_queue, config, batch_size, wait_seconds, delete, e
             else:
                 continue
 
-        resultlist = multiscan(filelist, configfile=config)
+        resultlist = multiscan(filelist, config=config)
         results = parse_reports(resultlist, python=True)
         if delete:
             for file_name in results:
@@ -118,13 +116,16 @@ def multiscanner_process(work_queue, config, batch_size, wait_seconds, delete, e
 
 def _main():
     args = _parse_args()
+    if args.config != msconf.CONFIG_FILEPATH:
+        msconf.update_ms_config_file(args.config)
+
     work_queue = multiprocessing.Queue()
     exit_signal = multiprocessing.Value('b')
     exit_signal.value = False
     observer = start_observer(args.Directory, work_queue, args.recursive)
     ms_process = multiprocessing.Process(
         target=multiscanner_process,
-        args=(work_queue, args.config, args.batch, args.seconds, args.delete, exit_signal))
+        args=(work_queue, msconf.MS_CONFIG, args.batch, args.seconds, args.delete, exit_signal))
     ms_process.start()
     try:
         while True:
@@ -141,7 +142,7 @@ def _main():
 def _parse_args():
     parser = argparse.ArgumentParser(description='Monitor a directory and submit new files to MultiScanner')
     parser.add_argument("-c", "--config", help="The config file to use", required=False,
-                        default=MS_CONFIG)
+                        default=msconf.CONFIG_FILEPATH)
     parser.add_argument("-s", "--seconds", help="The number of seconds to wait for additional files",
                         required=False, default=120, type=int)
     parser.add_argument("-b", "--batch", help="The max number of files per batch", required=False,

multiscanner/common/pdf_generator/__init__.py

@@ -1,27 +1,26 @@
 from __future__ import (division, absolute_import, with_statement,
                         print_function, unicode_literals)
 
-import json
 import os
 
 from reportlab.lib import colors, units
 from reportlab.platypus import TableStyle
 
 from multiscanner.common.pdf_generator import generic_pdf
+from multiscanner import config as msconf
 
 
 def create_pdf_document(DIR, report):
     '''
     Method to create a PDF report based of a multiscanner JSON report.
 
     Args:
-        DIR: Represents the a directory containing the 'pdf_config.json' file.
+        DIR: Represents the a directory containing the 'pdf_config.ini' file.
         report: A JSON object.
 
     '''
-    with open(os.path.join(os.path.split(DIR)[0], 'pdf_config.json')) as data_file:
-        pdf_components = json.load(data_file)
-
+    pdf_config = os.path.join(DIR, 'pdf_config.ini')
+    pdf_components = msconf.read_config(pdf_config).get_section('pdf')
     gen_pdf = generic_pdf.GenericPDF(pdf_components)
 
     notice = []

multiscanner/common/utils.py

@@ -3,14 +3,11 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import (absolute_import, division, unicode_literals, with_statement)
 
-import ast
-import configparser
 import imp
 import logging
 import os
-import sys
 
-from multiscanner.config import PY3
+from six import PY3
 
 logger = logging.getLogger(__name__)
 
@@ -76,42 +73,6 @@ def convert_encoding(data, encoding='UTF-8', errors='replace'):
         return data
 
 
-def parse_config(config_object):
-    """Take a config object and returns it as a dictionary"""
-    return_var = {}
-    for section in config_object.sections():
-        section_dict = dict(config_object.items(section))
-        for key in section_dict:
-            try:
-                section_dict[key] = ast.literal_eval(section_dict[key])
-            except Exception as e:
-                logger.debug(e)
-        return_var[section] = section_dict
-    return return_var
-
-
-def get_config_path(config_file, component):
-    """Gets the location of the config file for the given multiscanner component
-    from the multiscanner config file
-
-    Components:
-        storage
-        api
-        web"""
-    conf = configparser.ConfigParser()
-    conf.read(config_file)
-    conf = parse_config(conf)
-    try:
-        return conf['main']['%s-config' % component]
-    except KeyError:
-        logger.error(
-            "Couldn't find '{}-config' value in 'main' section "
-            "of config file. Have you run 'python multiscanner.py init'?"
-            .format(component)
-        )
-        sys.exit()
-
-
 def dirname(path):
     """OS independent version of os.path.dirname"""
     split = path.split('/')
@@ -134,7 +95,7 @@ def basename(path):
         return split[-1]
 
 
-def parseDir(directory, recursive=False, exclude=['__init__']):
+def parse_dir(directory, recursive=False, exclude=['__init__']):
     """
     Returns a list of files in a directory.
 
@@ -148,7 +109,7 @@ def parseDir(directory, recursive=False, exclude=['__init__']):
         item = os.path.join(directory, item)
         if os.path.isdir(item):
             if recursive:
-                filelist.extend(parseDir(item, recursive))
+                filelist.extend(parse_dir(item, recursive))
             else:
                 continue
         else:
@@ -162,7 +123,7 @@ def parseDir(directory, recursive=False, exclude=['__init__']):
     return filelist
 
 
-def parseFileList(FileList, recursive=False):
+def parse_file_list(FileList, recursive=False):
     """
     Takes a list of files and directories and returns a list of files.
 
@@ -173,7 +134,7 @@ def parseFileList(FileList, recursive=False):
     filelist = []
     for item in FileList:
         if os.path.isdir(item):
-            filelist.extend(parseDir(item, recursive))
+            filelist.extend(parse_dir(item, recursive))
         elif os.path.isfile(item):
             if not PY3:
                 filelist.append(item.decode('utf8'))