Change scan_file.py to peekaboo-util.py

Add subcommand scan_file with same behaviour as scan_file.py
previously. Moved everything into class PeekabooUtil.

Allow more commands to be added to the API.
Server can now easily be extended to understand and handle new API requests

4 files changed, 155 insertions, 117 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ed110c9..1837167 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,8 @@ See documentation for details.
 - Improve REST API access robustness by introducing configurable urllib3 retry
   handling with backoff and defined endless retry or failure report to client.
   (#43)
+- Introduce peekaboo-util.py as a super charged replacement for scan_file.py
+  (#107)
 
 ## 1.7
 
diff --git a/bin/peekaboo-util.py b/bin/peekaboo-util.py
new file mode 100755
index 0000000..e9b3cbf
--- /dev/null
+++ b/bin/peekaboo-util.py
@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+# encoding: utf-8
+
+###############################################################################
+#                                                                             #
+# Peekaboo Extended Email Attachment Behavior Observation Owl                 #
+#                                                                             #
+# peekaboo-util.py                                                            #
+###############################################################################
+#                                                                             #
+# Copyright (C) 2016-2019  science + computing ag                             #
+#                                                                             #
+# This program is free software: you can redistribute it and/or modify        #
+# it under the terms of the GNU General Public License as published by        #
+# the Free Software Foundation, either version 3 of the License, or (at       #
+# your option) any later version.                                             #
+#                                                                             #
+# This program is distributed in the hope that it will be useful, but         #
+# WITHOUT ANY WARRANTY; without even the implied warranty of                  #
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU           #
+# General Public License for more details.                                    #
+#                                                                             #
+# You should have received a copy of the GNU General Public License           #
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
+#                                                                             #
+###############################################################################
+
+
+from __future__ import print_function
+from os import path
+from argparse import ArgumentParser
+import socket
+import re
+import logging
+
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+
+class PeekabooUtil(object):
+    """ Utility fo interface with Peekaboo API over the socket connection """
+    def __init__(self, socket_file):
+        logger.debug('Initialising PeekabooUtil')
+        self.peekaboo = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        logger.debug('Opening socket %s', socket_file)
+        self.peekaboo.connect(socket_file)
+
+    def send_receive(self, request, output=False):
+        """ Send request to peekaboo and return its answer """
+        logger.debug('Sending request: %s', request)
+
+        self.peekaboo.send(request)
+        print ('Waiting for response...')
+
+        buf = ''
+        while True:
+            data = self.peekaboo.recv(1024)
+            if data:
+                buf += data
+                if output:
+                    print(data, end='')
+            else:
+                self.peekaboo.close()
+                break
+        logger.debug('Received from peekaboo: %s', buf)
+        return buf
+
+    def scan_file(self, filename):
+        """ Scan the supplied filenames with peekaboo and output result """
+        result_regex = re.compile(r'.*wurde als',
+                                  re.MULTILINE + re.DOTALL + re.UNICODE)
+        file_snippets = []
+        for filename in filename:
+            file_snippets.append('{ "full_name": "%s" }' % path.abspath(filename))
+        request = '[ %s ]' % ', '.join(file_snippets)
+
+        buf = self.send_receive(request)
+
+        for result in buf.splitlines():
+            output = result_regex.search(result)
+            if output:
+                if 'bad' in result:
+                    print(result)
+                logger.info(result)
+
+
+def main():
+    parser = ArgumentParser()
+    subparsers = parser.add_subparsers(help='commands')
+
+    parser.add_argument('-v', '--verbose', action='store_true', required=False,
+                        help='List results of all files not only bad ones')
+    parser.add_argument('-vv', '--verbose2', action='store_true', required=False,
+                        help='List detailed analysis results of every rule')
+    parser.add_argument('-d', '--debug', action='store_true', required=False,
+                        help='Output additional diagnostics')
+    parser.add_argument('-s', '--socket-file', action='store', required=True,
+                        help='Path to Peekaboo\'s socket file')
+
+    scan_file_parser = subparsers.add_parser('scan-file',
+                                             help='Scan a file and report it')
+    scan_file_parser.add_argument('-f', '--filename', action='append', required=True,
+                                  help='Path to the file to scan. Can be given more '
+                                       'than once to scan multiple files.')
+    scan_file_parser.set_defaults(func=command_scan_file)
+
+    args = parser.parse_args()
+
+    logger.setLevel(logging.ERROR)
+    if args.verbose:
+        logger.setLevel(logging.INFO)
+    if args.verbose2 or args.debug:
+        logger.setLevel(logging.DEBUG)
+
+    args.func(args)
+
+def command_scan_file(args):
+    """ Handler for command scan_file """
+    util = PeekabooUtil(args.socket_file)
+    util.scan_file(args.filename)
+
+if __name__ == "__main__":
+    main()
diff --git a/bin/scan_file.py b/bin/scan_file.py
deleted file mode 100755
index cfbb5fb..0000000
--- a/bin/scan_file.py
+++ /dev/null
@@ -1,91 +0,0 @@
-#!/usr/bin/env python
-# encoding: utf-8
-
-###############################################################################
-#                                                                             #
-# Peekaboo Extended Email Attachment Behavior Observation Owl                 #
-#                                                                             #
-# scan_file.py                                                                #
-###############################################################################
-#                                                                             #
-# Copyright (C) 2016-2019  science + computing ag                             #
-#                                                                             #
-# This program is free software: you can redistribute it and/or modify        #
-# it under the terms of the GNU General Public License as published by        #
-# the Free Software Foundation, either version 3 of the License, or (at       #
-# your option) any later version.                                             #
-#                                                                             #
-# This program is distributed in the hope that it will be useful, but         #
-# WITHOUT ANY WARRANTY; without even the implied warranty of                  #
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU           #
-# General Public License for more details.                                    #
-#                                                                             #
-# You should have received a copy of the GNU General Public License           #
-# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
-#                                                                             #
-###############################################################################
-
-
-from __future__ import print_function
-from os import path, linesep
-from argparse import ArgumentParser
-import socket
-import re
-
-
-def main():
-    parser = ArgumentParser()
-    parser.add_argument('-v', '--verbose', action='store_true', required=False,
-                        help='List results of all files not only bad ones')
-    parser.add_argument('-vv', '--verbose2', action='store_true', required=False,
-                        help='List detailed analysis results of every rule')
-    parser.add_argument('-d', '--debug', action='store_true', required=False,
-                        help='Output additional diagnostics')
-    parser.add_argument('-s', '--socket_file', action='store', required=True,
-                        help='Path to Peekaboo\'s socket file')
-    parser.add_argument('-f', '--filename', action='append', required=True,
-                        help='Path to the file to scan. Can be given more '
-                        'than once to scan multiple files.')
-    args = parser.parse_args()
-
-    result_regex = re.compile(r'.*wurde als',
-                              re.MULTILINE + re.DOTALL + re.UNICODE)
-    peekaboo = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    peekaboo.connect(args.socket_file)
-
-    file_snippets = []
-    for filename in args.filename:
-        file_snippets.append('{ "full_name": "%s" }' % path.abspath(filename))
-    request = '[ %s ]' % ', '.join(file_snippets)
-
-    if args.debug:
-        print ('Sending request: %s' % request)
-
-    peekaboo.send(request)
-    print ('Waiting for response...')
-
-    if args.verbose2:
-        args.verbose = True
-
-    buf = ''
-    while True:
-        data = peekaboo.recv(1024)
-        if data:
-            buf += data.decode('utf-8')
-        else:
-            peekaboo.close()
-            break
-
-    for result in buf.splitlines():
-        output = result_regex.search(result)
-        if output:
-            if 'bad' in result:
-                print(result)
-            elif args.verbose:
-                print(result)
-        elif args.verbose2:
-            print(result)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/peekaboo/server.py b/peekaboo/server.py
index dc27d4d..417c625 100644
--- a/peekaboo/server.py
+++ b/peekaboo/server.py
@@ -153,6 +153,30 @@ class PeekabooStreamRequestHandler(socketserver.StreamRequestHandler):
         # here we know that all samples have reported back
         self.report(submitted)
 
+    def submit_sample(self, api_data):
+        """ Submit API supplied file as Sample """
+        path = api_data['full_name']
+        logger.info("Got run_analysis request for %s", path)
+        if not os.path.exists(path):
+            self.talk_back(_('ERROR: Path does not exist or no '
+                             'permission to access it.'))
+            logger.error('Path does not exist or no permission '
+                         'to access it.')
+            return None
+
+        if not os.path.isfile(path):
+            self.talk_back(_('ERRROR: Input is not a file'))
+            logger.error('Input is not a file')
+            return None
+
+        sample = self.sample_factory.make_sample(
+            path, status_change=self.status_change, metainfo=api_data)
+        if not self.job_queue.submit(sample, self.__class__):
+            self.talk_back(_('Error submitting sample to job queue'))
+            # submit will have logged an error
+            return None
+        return sample
+
     def parse(self):
         """ Reads and parses an analysis request. This is expected to be a JSON
         structure containing the path of the directory / file to analyse.
@@ -182,35 +206,14 @@ class PeekabooStreamRequestHandler(socketserver.StreamRequestHandler):
 
         submitted = []
         for part in parts:
-            if 'full_name' not in part:
+            if 'full_name' in part:
+                sample = self.submit_sample(part)
+                submitted.append(sample)
+                logger.debug('Created and submitted sample %s', sample)
+            else:
                 self.talk_back(_('ERROR: Incomplete data structure.'))
                 logger.error('Incomplete data structure.')
                 return None
-
-            path = part['full_name']
-            logger.info("Got run_analysis request for %s", path)
-            if not os.path.exists(path):
-                self.talk_back(_('ERROR: Path does not exist or no '
-                                 'permission to access it.'))
-                logger.error('Path does not exist or no permission '
-                             'to access it.')
-                return None
-
-            if not os.path.isfile(path):
-                self.talk_back(_('ERRROR: Input is not a file'))
-                logger.error('Input is not a file')
-                return None
-
-            sample = self.sample_factory.make_sample(
-                path, status_change=self.status_change, metainfo=part)
-            if not self.job_queue.submit(sample, self.__class__):
-                self.talk_back(_('Error submitting sample to job queue'))
-                # submit will have logged an error
-                return None
-
-            submitted.append(sample)
-            logger.debug('Created and submitted sample %s', sample)
-
         return submitted
 
     def wait(self, to_be_analysed):