cleanup logging and add listener exception handling

bin/margaritashotgun

@@ -5,9 +5,7 @@ import os
 sys.path.append(os.getcwd())
 
 import margaritashotgun
-import logging
 
 name = 'margaritashotgun'
-margaritashotgun.set_stream_logger(name=name)
 client = margaritashotgun.client(library=False, name=name)
 client.run()

margaritashotgun/client.py

@@ -36,7 +36,9 @@ def __init__(self, config=None, library=True, name=None, verbose=False):
             self.config = self.cli.configure(config=config)
 
         if self.verbose is True:
-            margaritashotgun.set_stream_logger(level=logging.DEBUG)
+            margaritashotgun.set_stream_logger(name=self.name, level=logging.DEBUG)
+        else:
+            margaritashotgun.set_stream_logger(name=self.name, level=logging.INFO)
 
 
 
@@ -48,19 +50,23 @@ def run(self):
         try:
             conf = self.map_config()
             workers = Workers(conf, self.config['workers'], name=self.name, library=self.library)
-            results = workers.spawn()
+            description = 'memory capture action'
+            results = workers.spawn(description)
 
             self.statistics(results)
-            logger.info(("{0} hosts processed. completed: {1} "
-                         "failed {2}".format(self.total, self.completed,
-                                             self.failed)))
-            logger.info("completed_hosts: {0}".format(self.completed_addresses))
-            logger.info("failed_hosts: {0}".format(self.failed_addresses))
             if self.library is True:
-                return results
+                return dict([('total', self.total),
+                             ('completed', self.completed_addresses)
+                             ('failed', self.failed_addresses)])
             else:
+                logger.info(("{0} hosts processed. completed: {1} "
+                             "failed {2}".format(self.total, self.completed,
+                                                 self.failed)))
+                logger.info("completed_hosts: {0}".format(self.completed_addresses))
+                logger.info("failed_hosts: {0}".format(self.failed_addresses))
                 quit()
         except KeyboardInterrupt:
+            workers.cleanup()
             quit(1)
 
     def map_config(self):

margaritashotgun/logger.py

@@ -5,18 +5,15 @@
 import time
 from datetime import datetime
 
-def listener(queue, name, log_file):
+def listener(queue, name, log_file, desc):
     root = logging.getLogger(name)
-    root.setLevel(logging.INFO)
 
     # write file header
     with open(log_file, 'w') as f:
         f.write('[\n')
         f.close
 
     # setup log file handler
-    # TODO: add desc as a parameter
-    desc = 'action'
     fileHandler = logging.FileHandler(log_file, mode='a')
     formatter = logging.Formatter(
         "\t{'timestamp': %(unixtime)s, 'message': '%(message)s', " +
@@ -28,10 +25,9 @@ def listener(queue, name, log_file):
     while True:
         try:
             raw_record = queue.get()
-            #print(raw_record)
             if raw_record is None:
                 break
-            logger = logging.getLogger(name)
+            logger = logging.getLogger(raw_record.name)
             record = logger.makeRecord(raw_record.name,
                                        raw_record.levelno,
                                        raw_record.filename,
@@ -41,23 +37,27 @@ def listener(queue, name, log_file):
                                        raw_record.exc_info,
                                        extra=get_times())
             fileHandler.handle(record)
+        except KeyboardInterrupt:
+            break
         except Exception as ex:
             print(ex)
-            print('Whoops! Problem:')
-
-    # write file footer
-    with open(log_file, 'a') as f:
-        f.write(']')
-        f.flush()
-        f.close()
+            break
 
+    cleanup(log_file)
 
 def get_times():
     tm = int(time.time())
     dt = datetime.utcfromtimestamp(tm).isoformat()
     times = {'unixtime': tm, 'isotime': dt}
     return times
 
+def cleanup(log_file):
+    with open(log_file, 'a') as f:
+        f.write(']')
+        f.flush()
+        f.close()
+
+
 class Logger(multiprocessing.Process):
 
     def __init__(self, *args, **kwargs):

margaritashotgun/memory.py

@@ -7,6 +7,8 @@
 import logging
 from margaritashotgun.exceptions import *
 
+logger = logging.getLogger(__name__)
+
 
 class OutputDestinations(Enum):
     local = 'local'
@@ -16,7 +18,7 @@ class OutputDestinations(Enum):
 class Memory():
 
     def __init__(self, remote_addr, mem_size, progressbar=False,
-                 recv_size=1048576, sock_timeout=1, log=None):
+                 recv_size=1048576, sock_timeout=1):
         """
         :type remote_addr: str
         :param remote_addr: hostname or ip address of target server
@@ -29,10 +31,6 @@ def __init__(self, remote_addr, mem_size, progressbar=False,
         :type sock_timeout: int
         :param sock_timeout: transfer socket receive timeout
         """
-
-        global logger
-        logger = log
-
         self.mem_size = mem_size
         self.progressbar = progressbar
         self.recv_size = recv_size
@@ -81,8 +79,8 @@ def capture(self, tunnel_addr, tunnel_port, filename=None,
         if filename is None:
             raise MemoryCaptureAttributeMissingError('filename')
         if destination == OutputDestinations.local:
-            logger.info("{0}: dumping to file:///{1}".format(self.remote_addr,
-                                                             filename))
+            logger.info("{0}: dumping to file://{1}".format(self.remote_addr,
+                                                            filename))
             result = self.to_file(filename, tunnel_addr, tunnel_port)
         elif destination == OutputDestinations.s3:
             if bucket is None:
@@ -207,7 +205,6 @@ def to_s3(self, bucket, filename, tunnel_addr, tunnel_port):
                     else:
                         self.cleanup()
                         raise
-
         self.cleanup()
         logger.info('{0}: capture complete: s3://{1}/{2}'.format(self.remote_addr,
                                                                  bucket,

margaritashotgun/remote_host.py

@@ -38,9 +38,6 @@ def process(conf):
     tunnel_port = random.randint(10000, 30000)
     remote_module_path = '/tmp/lime.ko'
 
-    #TODO: parameterize? do we do this elsewhere?
-    desc = "margaritashotgun action"
-
     queue_handler = QueueHandler(log_queue)
     logger = logging.getLogger('margaritashotgun')
     logger.addHandler(queue_handler)
@@ -54,7 +51,7 @@ def process(conf):
         filename = "{0}_mem.lime".format(remote_addr)
 
     try:
-        host = Host(logger=logger)
+        host = Host()
         host.connect(username, password, key, remote_addr, remote_port)
         host.start_tunnel(tunnel_port, tunnel_addr, tunnel_port)
         # TODO: cleanup, if I can implement the rest
@@ -82,13 +79,15 @@ def process(conf):
         if lime_loaded:
             result = host.capture_memory(dest, filename, bucket, progressbar)
         else:
-            result = None
+            result = False
         host.cleanup()
 
         return (remote_addr, result)
     except KeyboardInterrupt:
+        logger.removeHandler(queue_handler)
+        queue_handler.close()
         host.cleanup()
-        return (remote_addr, result)
+        return (remote_addr, False)
     except Exception as ex:
         # TODO: log other exception, return failure condition
         host.cleanup()
@@ -97,20 +96,19 @@ def process(conf):
 
 class Host():
 
-    def __init__(self, logger=None):
+    def __init__(self):
         """
         """
-        self.logger = logger
         self.memory = None
         self.tunnel = None
         self.shell = None
         self.remote_addr = None
         self.remote_port = None
         self.tunnel_addr = "127.0.0.1"
         self.tunnel_port = None
-        self.shell = RemoteShell(logger=self.logger)
+        self.shell = RemoteShell()
         self.commands = Commands
-        self.tunnel = SSHTunnel(log=self.logger)
+        self.tunnel = SSHTunnel()
 
     def connect(self, username, password, key, address, port):
         """
@@ -249,7 +247,7 @@ def capture_memory(self, destination, filename, bucket, progressbar):
         """
         """
         mem_size = self.mem_size()
-        mem = Memory(self.remote_addr, mem_size, progressbar=progressbar, log=self.logger)
+        mem = Memory(self.remote_addr, mem_size, progressbar=progressbar)
         mem.capture(self.tunnel_addr, self.tunnel_port, destination=destination,
                     filename=filename, bucket=bucket)
 

margaritashotgun/remote_shell.py

@@ -7,7 +7,7 @@
 from margaritashotgun.exceptions import *
 import logging
 
-#logger = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 class Commands(Enum):
@@ -21,14 +21,11 @@ class Commands(Enum):
 
 class RemoteShell():
 
-    def __init__(self, max_async_threads=2, logger=None):
+    def __init__(self, max_async_threads=2):
         """
         :type args: int
         :param args: maximun number of async command executors
         """
-        #if log is not None:
-        #    global logger
-        self.logger = logger
 
         self.ssh = paramiko.SSHClient()
         self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
@@ -50,7 +47,7 @@ def connect(self, auth, address, port):
         self.address = address
         self.port = port
         try:
-            self.logger.debug(("{0}: paramiko client connecting to "
+            logger.debug(("{0}: paramiko client connecting to "
                           "{0}:{1} with {2}".format(address,
                                                     port,
                                                     auth.method)))
@@ -61,7 +58,7 @@ def connect(self, auth, address, port):
                                            address, port)
             else:
                 raise AuthenticationMethodMissingError()
-            self.logger.debug(("{0}: paramiko client connected to "
+            logger.debug(("{0}: paramiko client connected to "
                           "{0}:{1}".format(address, port)))
         except (AuthenticationException, SSHException, SocketError) as ex:
             raise SSHConnectionError("{0}:{1}".format(address, port), ex)
@@ -110,7 +107,7 @@ def execute(self, command):
         :type command: str
         :param command: command to be run on remote host
         """
-        self.logger.debug('{0}: executing "{1}"'.format(self.address, command))
+        logger.debug('{0}: executing "{1}"'.format(self.address, command))
         stdin, stdout, stderr = self.ssh.exec_command(command)
         return dict(zip(['stdin', 'stdout', 'stderr'],
                         [stdin, stdout, stderr]))
@@ -124,7 +121,7 @@ def execute_async(self, command, callback=None):
         :type callback: function
         :param callback: function to call when execution completes
         """
-        self.logger.debug(('{0}: execute async "{1}"'
+        logger.debug(('{0}: execute async "{1}"'
                       'with callback {2}'.format(self.address, command,
                                                  callback)))
         future = self.executor.submit(self.execute, command)
@@ -143,7 +140,7 @@ def decode(self, stream, encoding='utf-8'):
         """
         data = stream.read().decode(encoding).strip("\n")
         if data != "":
-            self.logger.debug(('{0}: decoded "{1}" with encoding '
+            logger.debug(('{0}: decoded "{1}" with encoding '
                           '{2}'.format(self.address, data, encoding)))
         return data
 
@@ -156,15 +153,15 @@ def upload_file(self, local_path, remote_path):
         :type remote_path: str
         :param remote_path: destination path of upload on remote host
         """
-        self.logger.debug("{0}: uploading {1} to {0}:{2}".format(self.address,
+        logger.debug("{0}: uploading {1} to {0}:{2}".format(self.address,
                                                             local_path,
                                                             remote_path))
         try:
             sftp = self.ssh.open_sftp()
             sftp.put(local_path, remote_path)
             sftp.close()
         except SSHException as ex:
-            self.logger.WARN(("{0}: LiME module upload failed with exception:"
+            logger.warn(("{0}: LiME module upload failed with exception:"
                          "{1}".format(self.address, ex)))
 
     def cleanup(self):

margaritashotgun/ssh_tunnel.py

@@ -13,14 +13,12 @@
 except ImportError:
     import SocketServer as socketserver
 
+logger = logging.getLogger(__name__)
 
-class SSHTunnel():
 
-    def __init__(self, log=None):
-        if log is not None:
-            global logger
-            logger = log
+class SSHTunnel():
 
+    def __init__(self):
         self.transport = None
         self.forward = None