Merge pull request #10 from Shura1oplot/improvements2

Bugfixes

.travis.yml

@@ -7,7 +7,7 @@ python:
   - "2.7"
 
 install:
-  - "sudo apt-get install --force-yes -y python-dev zlib1g-dev p7zip-full libicu-dev fontconfig rubygems-integration poppler-utils ghostscript"
+  - "sudo apt-get install --force-yes -y python-dev zlib1g-dev unzip p7zip-full p7zip-rar libicu-dev fontconfig rubygems-integration poppler-utils ghostscript"
   - "sudo gem install coveralls-lcov"
   - "sudo apt-get remove --purge -y bzip2 libbz2-dev"
   - "pip install -r utils/dfvfs-requirements.txt"
@@ -16,6 +16,10 @@ install:
   - "python setup.py install"
   - "cp gransk/web/tests/package.json /opt && npm install --prefix=/opt /opt"
 
+before_script:
+  - "chmod +x utils/travis/*"
+  - "export PATH=utils/travis:$PATH"
+
 script:
   - "python run_tests.py"
   - "py.test --cov gransk gransk"

Dockerfile

@@ -12,7 +12,7 @@ WORKDIR /app
 
 RUN apt-get update
 RUN apt-get install --force-yes -y \
-      python-dev python-setuptools zlib1g-dev p7zip-full \
+      python-dev python-setuptools zlib1g-dev unzip p7zip-full p7zip-rar \
       python-pip libicu-dev poppler-utils ghostscript && \
   pip install -r utils/dfvfs-requirements.txt && \
   pip install -r requirements.txt && \

gransk/api.py

@@ -5,13 +5,10 @@
 
 import requests
 import time
-import json
-import hashlib
 import logging
-import traceback
+import threading
 import yaml
 import os
-import sys
 import shutil
 
 import six.moves.http_client
@@ -70,16 +67,17 @@ def consume(self, doc, file_object=None):
           return
 
       if not file_object:
-        file_object = open(doc.path)
+        file_object = open(doc.path, "rb")
 
       self.produce(helper.EXTRACT_META, doc, file_object)
       self.produce(helper.PROCESS_FILE, doc, file_object)
 
       file_object.close()
 
     except Exception as err:
+      LOGGER.exception('could not process %s: %s', doc.path, err)
+      doc.status = 'error'
       doc.meta['gransk_error'] = six.text_type(err)
-      traceback.print_exc(file=sys.stdout)
 
 
 class API(object):
@@ -113,12 +111,18 @@ def __init__(self, injector=None, config_path=None):
     self.pipeline = pipeline.build_pipeline(self.config)
     self.entrypoint = Subscriber(self.pipeline)
     self.entrypoint.setup(self.config)
+    self.write_lock = threading.Lock()
 
   def add_file(self, doc, file_object):
-    return self.entrypoint.consume(doc, file_object=file_object)
+    with self.write_lock:
+      return self.entrypoint.consume(doc, file_object=file_object)
 
   def clear_all(self):
     """Clear all processed data."""
+    with self.write_lock:
+      self._clear_all()
+
+  def _clear_all(self):
     try:
         if os.path.exists(self.config[helper.DATA_ROOT]):
           shutil.rmtree(self.config[helper.DATA_ROOT])
@@ -128,7 +132,7 @@ def clear_all(self):
         os.makedirs(os.path.join(self.config[helper.DATA_ROOT], 'archives'))
         os.makedirs(os.path.join(self.config[helper.DATA_ROOT], 'archives', '.tmp'))
     except Exception as err:
-       print (">>", err)
+       LOGGER.error("could not clear data: %s", err)
 
     connection = self.config['injector'].get_http_connection('%s:%s' % (self.config['es_host'][0], 9200))
     connection.request('DELETE', '/gransk', '', {})

gransk/boot/run.py

@@ -8,28 +8,17 @@
 import os
 import argparse
 import logging
-import six.moves.http_client
-import shutil
 import sys
 
-import yaml
-
-import gransk.core.helper as helper
-import gransk.api
-
 from multiprocessing import Queue, Process
 from six.moves.queue import Empty
 import time
-import traceback
-import logging
-import os
-import sys
 import glob
-import six
 from six.moves import range
 
 from tqdm import tqdm
 
+import gransk.api
 import gransk.core.compat as _
 import gransk.core.file_collector as collector
 import gransk.core.pipeline as pipeline

gransk/boot/ui.py

@@ -9,13 +9,10 @@
 import os
 import logging
 import json
-import shutil
-import time
+from functools import wraps
 import requests
 
-from flask import Flask, Response, render_template, request, abort
-from werkzeug import secure_filename
-import yaml
+from flask import Flask, Response, request, abort
 
 import gransk.api
 import gransk.core.helper as helper
@@ -32,9 +29,6 @@
     static_folder=os.path.join(_root, 'gransk', 'web', 'app'))
 
 
-from functools import wraps
-from flask import request, Response
-
 _globals = {}
 
 def check_auth(username, password):
@@ -83,8 +77,9 @@ def upload():
   _file = request.files.get('file')
 
   doc = document.get_document(
-      secure_filename(_file.filename),
-      parent=document.get_document('root'))
+      _file.filename,
+      parent=document.get_document('root'),
+      need_secure_path=True)
 
   doc.tag = 'upload'
 
@@ -104,8 +99,8 @@ def delete_data():
 @app.route('/file', methods=['GET'])
 def get_file():
   """Get original file."""
-  filename = secure_filename(request.args['filename'])
-  ext = secure_filename(request.args['ext'])
+  filename = document.secure_path(request.args['filename'])
+  ext = document.secure_path(request.args['ext'])
   mediatype = request.args['mediatype']
 
   root = os.path.join(_globals['gransk'].config[helper.DATA_ROOT], 'files')
@@ -117,6 +112,7 @@ def get_file():
   with open(file_path, 'rb') as inp:
     return Response(inp.read(), mimetype=mediatype, status=200)
 
+
 @app.route('/search')
 def search():
   query = json.loads(request.args['q'])
@@ -133,7 +129,7 @@ def search():
 @app.route('/picture', methods=['GET'])
 def picture():
   """Get document content as picture."""
-  name = secure_filename(request.args['name'])
+  name = document.secure_path(request.args['name'])
   mediatype = request.args['mediatype']
 
   root = os.path.join(_globals['gransk'].config[helper.DATA_ROOT], 'pictures')
@@ -231,5 +227,5 @@ def setup(args, pipeline, runmod, injector):
   #context = ('/etc/letsencrypt/live/gransk.com/cert.pem', '/etc/letsencrypt/live/gransk.com/privkey.pem')
   context = None
 
-  app.run(host=args.host, port=args.port, debug=args.debug, threaded=True , ssl_context=context)
+  app.run(host=args.host, port=args.port, debug=args.debug, threaded=True, ssl_context=context)
   _globals['gransk'].pipeline.stop()

gransk/core/detect_type.py

@@ -4,7 +4,6 @@
 from __future__ import absolute_import, unicode_literals
 
 import logging
-import traceback
 import sys
 import os
 import six
@@ -92,8 +91,7 @@ def consume(self, doc, payload):
       self.produce(helper.RUN_PIPELINE, doc, None)
 
     except Exception as err:
-      traceback.print_exc(file=sys.stdout)
-      LOGGER.warning('could not process %s: %s', doc.path, err)
+      LOGGER.exception('could not process %s: %s', doc.path, err)
       doc.status = 'error'
       doc.meta['gransk_error'] = six.text_type(err)
       self.produce(helper.ERRORED_FILE, doc, payload)

gransk/core/document.py

@@ -10,6 +10,9 @@
 
 from six import text_type as unicode
 
+from werkzeug import secure_filename
+from slugify import slugify
+
 
 class Entities(object):
   """Class for a set of entities found within documents."""
@@ -142,32 +145,40 @@ def as_obj(self):
     }
 
 
-def get_document(path, parent=None):
+def get_document(path, parent=None, need_secure_path=None):
   """
   Create a new document object from the given path.
 
   :param path: Path to document (does not have to exist on file system).
   :param parent: Parent document (e.g. diskimage or archive).
   :returns: ``gransk.core.Document``
   """
-  if isinstance(path, unicode):
-    bpath, upath = path.encode('utf-8'), path
-  else:
-    bpath, upath = path, path.decode('utf-8')
+  if isinstance(path, bytes):
+    path = path.decode('utf-8')
+
+  original_path = path
+
+  if need_secure_path:
+    path = secure_path(path)
+
+    if not path:
+      path = 'unnamed'
 
   doc = Document()
-  doc.path = upath
+  doc.path = path
+
+  doc.meta['original_path'] = original_path
 
-  if os.path.dirname(doc.path):
-    doc.meta['directory'] = os.path.dirname(doc.path)
+  if os.path.dirname(path):
+    doc.meta['directory'] = os.path.dirname(path)
 
   digest = hashlib.md5()
 
-  digest.update(bpath)
+  digest.update(path.encode('utf-8'))
 
   doc.docid = digest.hexdigest()
 
-  _, ext = os.path.splitext(doc.path)
+  _, ext = os.path.splitext(path)
   doc.ext = ext.lstrip('.').lower() or 'none'
 
   doc.parent = parent
@@ -179,3 +190,26 @@ def get_document(path, parent=None):
 
   doc.added = int(time.time())
   return doc
+
+
+def secure_path(path):
+  dirname = os.path.dirname(path)
+  filename = os.path.basename(path)
+  file_base, file_ext = os.path.splitext(path)
+
+  dirname = secure_filename(slugify(dirname, only_ascii=True))
+  file_base = secure_filename(slugify(file_base, only_ascii=True)) or 'unnamed'
+  file_ext = secure_filename(slugify(file_ext, only_ascii=True))
+
+  if file_ext:
+    filename = '.'.join([file_base, file_ext])
+  else:
+    filename = file_base
+
+  if len(filename) > 200:
+    filename = '%s__%s' % (filename[:99], filename[-99:])
+
+  if dirname:
+    return os.path.join(dirname, filename)
+
+  return filename

gransk/core/helper.py

@@ -55,3 +55,4 @@
 MIN_SCORE = 'min_score'
 MIN_SHARED = 'min_shared'
 MAX_RESULTS = 'max_results'
+ES_INDEX = 'es_index'

gransk/core/pipeline.py

@@ -4,8 +4,6 @@
 from __future__ import absolute_import, unicode_literals
 
 import logging
-import traceback
-import sys
 import inspect
 import os
 from collections import defaultdict
@@ -133,8 +131,7 @@ def init_subscriber(config, subscriber_mod, pipeline):
 
     pipeline.subscribers.append(subscriber)
   except Exception as err:
-    traceback.print_exc(file=sys.stdout)
-    LOGGER.warning('! %s could not be loaded: %s', subscriber_mod, err)
+    LOGGER.exception('! %s could not be loaded: %s', subscriber_mod, err)
 
 
 def build_pipeline(config):

gransk/core/tests/test_helper.py

@@ -54,12 +54,28 @@ class MockHttpConnection(object):
 
   def __init__(self, response_text):
     self.response_text = response_text
+    self.request_headers = None
 
   def request(self, method, uri, payload, headers):
-    pass
+    self.request_headers = headers
 
   def getresponse(self):
-    return BytesIO(self.response_text)
+    return MockHttpResponse(self.response_text)
+
+
+class MockHttpResponse(object):
+
+    def __init__(self, response_data):
+        super(MockHttpResponse, self).__init__()
+        self.stream = BytesIO(response_data)
+        self.status = 200
+        self.reason = 'OK'
+
+    def read(self):
+      return self.stream.read()
+
+    def close(self):
+      self.stream.close()
 
 
 class MockElasticsearchIndex(object):
@@ -129,7 +145,6 @@ def __init__(self, text):
       self.entities.append(MockEntity('I-PER', parts))
 
 
-
 class MockWorker(object):
 
   def __init__(self):
@@ -141,10 +156,11 @@ def boot(self, inject, config, path):
     self.called = True
     return []
 
+
 class MockInjector(object):
 
-  def __init__(self, response_text=None, ner_entities=[]):
-    self.response_text = response_text
+  def __init__(self, response_text=b"", ner_entities=[]):
+    self.http_connection = MockHttpConnection(response_text)
     self.elastic = MockElasticsearch()
     self.elastic_helper = MockElasticsearchHelper()
     self.polyglot = MockPolyglot
@@ -158,7 +174,7 @@ def get_worker(self):
     return self.worker
 
   def get_http_connection(self, url=None):
-    return MockHttpConnection(self.response_text)
+    return self.http_connection
 
   def get_elasticsearch(self):
     return self.elastic

gransk/plugins/analysis/abstract_related.py

@@ -77,7 +77,7 @@ def load_all(self, config):
             else:
               self.buckets[key] = value
         except:
-          pass
+          logging.warning('could not load related_%s data', self.NAME)
 
   def stop(self):
     """Write data to file."""