Refactoring

4 files changed, 333 insertions, 0 deletions

diff --git a/matweb/file_removal_scheduler.py b/matweb/file_removal_scheduler.py
new file mode 100644
index 0000000..2ce7912
--- /dev/null
+++ b/matweb/file_removal_scheduler.py
@@ -0,0 +1,26 @@
+import glob
+import time
+import sys
+import os
+import random
+
+
+def run_file_removal_job(upload_folder_path):
+    if random.randint(0, 10) == 0:
+        for file in glob.glob(upload_folder_path + '/*'):
+            delete_file_when_too_old(file)
+
+
+def delete_file_when_too_old(filepath):
+    file_mod_time = os.stat(filepath).st_mtime
+
+    # time in second since last modification of file
+    last_time = time.time() - file_mod_time
+
+    # if file is older than our configured max timeframe, delete it
+    if last_time > int(os.environ.get('MAT2_MAX_FILE_AGE_FOR_REMOVAL', 15 * 60)):
+        try:
+            os.remove(filepath)
+        except OSError:
+            print('Automatic File Removal failed on file: ' + str(filepath))
+            sys.exit(1)
diff --git a/matweb/frontend.py b/matweb/frontend.py
new file mode 100644
index 0000000..93432b4
--- /dev/null
+++ b/matweb/frontend.py
@@ -0,0 +1,77 @@
+import hmac
+import os
+
+from flask import Blueprint, render_template, url_for, current_app, after_this_request, send_from_directory, request, \
+    flash
+from werkzeug.utils import secure_filename, redirect
+
+from matweb import file_removal_scheduler, utils
+
+routes = Blueprint('routes', __name__)
+
+
+@routes.route('/info')
+def info():
+    utils.get_supported_extensions()
+    return render_template(
+        'info.html', extensions=utils.get_supported_extensions()
+    )
+
+
+@routes.route('/download/<string:key>/<string:filename>')
+def download_file(key: str, filename: str):
+    if filename != secure_filename(filename):
+        return redirect(url_for('routes.upload_file'))
+
+    complete_path, filepath = utils.get_file_paths(filename, current_app.config['UPLOAD_FOLDER'])
+    file_removal_scheduler.run_file_removal_job(current_app.config['UPLOAD_FOLDER'])
+
+    if not os.path.exists(complete_path):
+        return redirect(url_for('routes.upload_file'))
+    if hmac.compare_digest(utils.hash_file(complete_path), key) is False:
+        return redirect(url_for('routes.upload_file'))
+
+    @after_this_request
+    def remove_file(response):
+        if os.path.exists(complete_path):
+            os.remove(complete_path)
+        return response
+    return send_from_directory(current_app.config['UPLOAD_FOLDER'], filepath, as_attachment=True)
+
+
+@routes.route('/', methods=['GET', 'POST'])
+def upload_file():
+    utils.check_upload_folder(current_app.config['UPLOAD_FOLDER'])
+    mime_types = utils.get_supported_extensions()
+
+    if request.method == 'POST':
+        if 'file' not in request.files:  # check if the post request has the file part
+            flash('No file part')
+            return redirect(request.url)
+
+        uploaded_file = request.files['file']
+        if not uploaded_file.filename:
+            flash('No selected file')
+            return redirect(request.url)
+
+        filename, filepath = utils.save_file(uploaded_file, current_app.config['UPLOAD_FOLDER'])
+        parser, mime = utils.get_file_parser(filepath)
+
+        if parser is None:
+            flash('The type %s is not supported' % mime)
+            return redirect(url_for('routes.upload_file'))
+
+        meta = parser.get_meta()
+
+        if parser.remove_all() is not True:
+            flash('Unable to clean %s' % mime)
+            return redirect(url_for('routes.upload_file'))
+
+        key, meta_after, output_filename = utils.cleanup(parser, filepath, current_app.config['UPLOAD_FOLDER'])
+
+        return render_template(
+            'download.html', mimetypes=mime_types, meta=meta, filename=output_filename, meta_after=meta_after, key=key
+        )
+
+    max_file_size = int(current_app.config['MAX_CONTENT_LENGTH'] / 1024 / 1024)
+    return render_template('index.html', max_file_size=max_file_size, mimetypes=mime_types)
\ No newline at end of file
diff --git a/matweb/rest_api.py b/matweb/rest_api.py
new file mode 100644
index 0000000..60d834f
--- /dev/null
+++ b/matweb/rest_api.py
@@ -0,0 +1,139 @@
+import os
+import base64
+import io
+import binascii
+import zipfile
+from uuid import uuid4
+
+from flask import after_this_request, send_from_directory
+from flask_restful import Resource, reqparse, abort, request
+from cerberus import Validator
+from werkzeug.datastructures import FileStorage
+from urllib.parse import urljoin
+
+from matweb import file_removal_scheduler, utils
+
+
+class APIUpload(Resource):
+    
+    def __init__(self, **kwargs):
+        self.upload_folder = kwargs['upload_folder']
+        
+    def post(self):
+        utils.check_upload_folder(self.upload_folder)
+        req_parser = reqparse.RequestParser()
+        req_parser.add_argument('file_name', type=str, required=True, help='Post parameter is not specified: file_name')
+        req_parser.add_argument('file', type=str, required=True, help='Post parameter is not specified: file')
+
+        args = req_parser.parse_args()
+        try:
+            file_data = base64.b64decode(args['file'])
+        except binascii.Error as err:
+            abort(400, message='Failed decoding file: ' + str(err))
+
+        file = FileStorage(stream=io.BytesIO(file_data), filename=args['file_name'])
+        filename, filepath = utils.save_file(file, self.upload_folder)
+        parser, mime = utils.get_file_parser(filepath)
+
+        if parser is None:
+            abort(415, message='The type %s is not supported' % mime)
+
+        meta = parser.get_meta()
+        if not parser.remove_all():
+            abort(500, message='Unable to clean %s' % mime)
+
+        key, meta_after, output_filename = utils.cleanup(parser, filepath, self.upload_folder)
+        return utils.return_file_created_response(
+            output_filename,
+            mime,
+            key,
+            meta,
+            meta_after,
+            urljoin(request.host_url, '%s/%s/%s/%s' % ('api', 'download', key, output_filename))
+        )
+
+
+class APIDownload(Resource):
+
+    def __init__(self, **kwargs):
+        self.upload_folder = kwargs['upload_folder']
+
+    def get(self, key: str, filename: str):
+        complete_path, filepath = utils.is_valid_api_download_file(filename, key, self.upload_folder)
+        # Make sure the file is NOT deleted on HEAD requests
+        if request.method == 'GET':
+            file_removal_scheduler.run_file_removal_job(self.upload_folder)
+
+            @after_this_request
+            def remove_file(response):
+                if os.path.exists(complete_path):
+                    os.remove(complete_path)
+                return response
+
+        return send_from_directory(self.upload_folder, filepath, as_attachment=True)
+
+
+class APIBulkDownloadCreator(Resource):
+
+    def __init__(self, **kwargs):
+        self.upload_folder = kwargs['upload_folder']
+
+    schema = {
+        'download_list': {
+            'type': 'list',
+            'minlength': 2,
+            'maxlength': int(os.environ.get('MAT2_MAX_FILES_BULK_DOWNLOAD', 10)),
+            'schema': {
+                'type': 'dict',
+                'schema': {
+                    'key': {'type': 'string', 'required': True},
+                    'file_name': {'type': 'string', 'required': True}
+                }
+            }
+        }
+    }
+    v = Validator(schema)
+
+    def post(self):
+        utils.check_upload_folder(self.upload_folder)
+        data = request.json
+        if not self.v.validate(data):
+            abort(400, message=self.v.errors)
+        # prevent the zip file from being overwritten
+        zip_filename = 'files.' + str(uuid4()) + '.zip'
+        zip_path = os.path.join(self.upload_folder, zip_filename)
+        cleaned_files_zip = zipfile.ZipFile(zip_path, 'w')
+        with cleaned_files_zip:
+            for file_candidate in data['download_list']:
+                complete_path, file_path = utils.is_valid_api_download_file(
+                    file_candidate['file_name'],
+                    file_candidate['key'],
+                    self.upload_folder
+                )
+                try:
+                    cleaned_files_zip.write(complete_path)
+                    os.remove(complete_path)
+                except ValueError:
+                    abort(400, message='Creating the archive failed')
+
+            try:
+                cleaned_files_zip.testzip()
+            except ValueError as e:
+                abort(400, message=str(e))
+
+        parser, mime = utils.get_file_parser(zip_path)
+        if not parser.remove_all():
+            abort(500, message='Unable to clean %s' % mime)
+        key, meta_after, output_filename = utils.cleanup(parser, zip_path, self.upload_folder)
+        return {
+                   'output_filename': output_filename,
+                   'mime': mime,
+                   'key': key,
+                   'meta_after': meta_after,
+                   'download_link': urljoin(request.host_url, '%s/%s/%s/%s' % ('api', 'download', key, output_filename))
+               }, 201
+
+
+class APISupportedExtensions(Resource):
+    def get(self):
+        return utils.get_supported_extensions()
diff --git a/matweb/utils.py b/matweb/utils.py
new file mode 100644
index 0000000..8dfff45
--- /dev/null
+++ b/matweb/utils.py
@@ -0,0 +1,91 @@
+import hmac
+import os
+import hashlib
+import mimetypes as mtype
+
+from flask_restful import abort
+from libmat2 import parser_factory
+from werkzeug.utils import secure_filename
+
+
+def get_allow_origin_header_value():
+    return os.environ.get('MAT2_ALLOW_ORIGIN_WHITELIST', '*').split(" ")
+
+
+def hash_file(filepath: str) -> str:
+    sha256 = hashlib.sha256()
+    with open(filepath, 'rb') as f:
+        while True:
+            data = f.read(65536)  # read the file by chunk of 64k
+            if not data:
+                break
+            sha256.update(data)
+    return sha256.hexdigest()
+
+
+def check_upload_folder(upload_folder):
+    if not os.path.exists(upload_folder):
+        os.mkdir(upload_folder)
+
+
+def return_file_created_response(output_filename, mime, key, meta, meta_after, download_link):
+    return {
+        'output_filename': output_filename,
+        'mime': mime,
+        'key': key,
+        'meta': meta,
+        'meta_after': meta_after,
+        'download_link': download_link
+    }
+
+
+def get_supported_extensions():
+    extensions = set()
+    for parser in parser_factory._get_parsers():
+        for m in parser.mimetypes:
+            extensions |= set(mtype.guess_all_extensions(m, strict=False))
+    # since `guess_extension` might return `None`, we need to filter it out
+    return sorted(filter(None, extensions))
+
+
+def save_file(file, upload_folder):
+    filename = secure_filename(file.filename)
+    filepath = os.path.join(upload_folder, filename)
+    file.save(os.path.join(filepath))
+    return filename, filepath
+
+
+def get_file_parser(filepath: str):
+    parser, mime = parser_factory.get_parser(filepath)
+    return parser, mime
+
+
+def cleanup(parser, filepath, upload_folder):
+    output_filename = os.path.basename(parser.output_filename)
+    parser, _ = parser_factory.get_parser(parser.output_filename)
+    meta_after = parser.get_meta()
+    os.remove(filepath)
+
+    key = hash_file(os.path.join(upload_folder, output_filename))
+    return key, meta_after, output_filename
+
+
+def get_file_paths(filename, upload_folder):
+    filepath = secure_filename(filename)
+
+    complete_path = os.path.join(upload_folder, filepath)
+    return complete_path, filepath
+
+
+def is_valid_api_download_file(filename, key, upload_folder):
+    if filename != secure_filename(filename):
+        abort(400, message='Insecure filename')
+
+    complete_path, filepath = get_file_paths(filename, upload_folder)
+
+    if not os.path.exists(complete_path):
+        abort(404, message='File not found')
+
+    if hmac.compare_digest(hash_file(complete_path), key) is False:
+        abort(400, message='The file hash does not match')
+    return complete_path, filepath