diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 260 |
1 files changed, 22 insertions, 238 deletions
| @@ -1,23 +1,10 @@ | |||
| 1 | import os | 1 | import os |
| 2 | import hmac | ||
| 3 | import mimetypes as mtype | ||
| 4 | from uuid import uuid4 | ||
| 5 | import jinja2 | 2 | import jinja2 |
| 6 | import base64 | ||
| 7 | import io | ||
| 8 | import binascii | ||
| 9 | import zipfile | ||
| 10 | 3 | ||
| 11 | from cerberus import Validator | 4 | from matweb import utils, rest_api, frontend |
| 12 | import utils | 5 | from flask import Flask |
| 13 | import file_removal_scheduler | 6 | from flask_restful import Api |
| 14 | from libmat2 import parser_factory | ||
| 15 | from flask import Flask, flash, request, redirect, url_for, render_template, send_from_directory, after_this_request | ||
| 16 | from flask_restful import Resource, Api, reqparse, abort | ||
| 17 | from werkzeug.utils import secure_filename | ||
| 18 | from werkzeug.datastructures import FileStorage | ||
| 19 | from flask_cors import CORS | 7 | from flask_cors import CORS |
| 20 | from urllib.parse import urljoin | ||
| 21 | 8 | ||
| 22 | 9 | ||
| 23 | def create_app(test_config=None): | 10 | def create_app(test_config=None): |
| @@ -32,235 +19,32 @@ def create_app(test_config=None): | |||
| 32 | if test_config is not None: | 19 | if test_config is not None: |
| 33 | app.config.update(test_config) | 20 | app.config.update(test_config) |
| 34 | 21 | ||
| 22 | # Non JS Frontend | ||
| 35 | app.jinja_loader = jinja2.ChoiceLoader([ # type: ignore | 23 | app.jinja_loader = jinja2.ChoiceLoader([ # type: ignore |
| 36 | jinja2.FileSystemLoader(app.config['CUSTOM_TEMPLATES_DIR']), | 24 | jinja2.FileSystemLoader(app.config['CUSTOM_TEMPLATES_DIR']), |
| 37 | app.jinja_loader, | 25 | app.jinja_loader, |
| 38 | ]) | 26 | ]) |
| 27 | app.register_blueprint(frontend.routes) | ||
| 39 | 28 | ||
| 29 | # Restful API hookup | ||
| 40 | api = Api(app) | 30 | api = Api(app) |
| 41 | CORS(app, resources={r"/api/*": {"origins": utils.get_allow_origin_header_value()}}) | 31 | CORS(app, resources={r"/api/*": {"origins": utils.get_allow_origin_header_value()}}) |
| 42 | 32 | api.add_resource( | |
| 43 | @app.route('/info') | 33 | rest_api.APIUpload, |
| 44 | def info(): | 34 | '/api/upload', |
| 45 | get_supported_extensions() | 35 | resource_class_kwargs={'upload_folder': app.config['UPLOAD_FOLDER']} |
| 46 | return render_template( | 36 | ) |
| 47 | 'info.html', extensions=get_supported_extensions() | 37 | api.add_resource( |
| 48 | ) | 38 | rest_api.APIDownload, |
| 49 | 39 | '/api/download/<string:key>/<string:filename>', | |
| 50 | @app.route('/download/<string:key>/<string:filename>') | 40 | resource_class_kwargs={'upload_folder': app.config['UPLOAD_FOLDER']} |
| 51 | def download_file(key: str, filename: str): | 41 | ) |
| 52 | if filename != secure_filename(filename): | 42 | api.add_resource( |
| 53 | return redirect(url_for('upload_file')) | 43 | rest_api.APIBulkDownloadCreator, |
| 54 | 44 | '/api/download/bulk', | |
| 55 | complete_path, filepath = get_file_paths(filename) | 45 | resource_class_kwargs={'upload_folder': app.config['UPLOAD_FOLDER']} |
| 56 | file_removal_scheduler.run_file_removal_job(app.config['UPLOAD_FOLDER']) | 46 | ) |
| 57 | 47 | api.add_resource(rest_api.APISupportedExtensions, '/api/extension') | |
| 58 | if not os.path.exists(complete_path): | ||
| 59 | return redirect(url_for('upload_file')) | ||
| 60 | if hmac.compare_digest(utils.hash_file(complete_path), key) is False: | ||
| 61 | return redirect(url_for('upload_file')) | ||
| 62 | @after_this_request | ||
| 63 | def remove_file(response): | ||
| 64 | if os.path.exists(complete_path): | ||
| 65 | os.remove(complete_path) | ||
| 66 | return response | ||
| 67 | return send_from_directory(app.config['UPLOAD_FOLDER'], filepath, as_attachment=True) | ||
| 68 | |||
| 69 | @app.route('/', methods=['GET', 'POST']) | ||
| 70 | def upload_file(): | ||
| 71 | utils.check_upload_folder(app.config['UPLOAD_FOLDER']) | ||
| 72 | mimetypes = get_supported_extensions() | ||
| 73 | |||
| 74 | if request.method == 'POST': | ||
| 75 | if 'file' not in request.files: # check if the post request has the file part | ||
| 76 | flash('No file part') | ||
| 77 | return redirect(request.url) | ||
| 78 | |||
| 79 | uploaded_file = request.files['file'] | ||
| 80 | if not uploaded_file.filename: | ||
| 81 | flash('No selected file') | ||
| 82 | return redirect(request.url) | ||
| 83 | |||
| 84 | filename, filepath = save_file(uploaded_file) | ||
| 85 | parser, mime = get_file_parser(filepath) | ||
| 86 | |||
| 87 | if parser is None: | ||
| 88 | flash('The type %s is not supported' % mime) | ||
| 89 | return redirect(url_for('upload_file')) | ||
| 90 | |||
| 91 | meta = parser.get_meta() | ||
| 92 | |||
| 93 | if parser.remove_all() is not True: | ||
| 94 | flash('Unable to clean %s' % mime) | ||
| 95 | return redirect(url_for('upload_file')) | ||
| 96 | |||
| 97 | key, meta_after, output_filename = cleanup(parser, filepath) | ||
| 98 | |||
| 99 | return render_template( | ||
| 100 | 'download.html', mimetypes=mimetypes, meta=meta, filename=output_filename, meta_after=meta_after, key=key | ||
| 101 | ) | ||
| 102 | |||
| 103 | max_file_size = int(app.config['MAX_CONTENT_LENGTH'] / 1024 / 1024) | ||
| 104 | return render_template('index.html', max_file_size=max_file_size, mimetypes=mimetypes) | ||
| 105 | |||
| 106 | def get_supported_extensions(): | ||
| 107 | extensions = set() | ||
| 108 | for parser in parser_factory._get_parsers(): | ||
| 109 | for m in parser.mimetypes: | ||
| 110 | extensions |= set(mtype.guess_all_extensions(m, strict=False)) | ||
| 111 | # since `guess_extension` might return `None`, we need to filter it out | ||
| 112 | return sorted(filter(None, extensions)) | ||
| 113 | |||
| 114 | def save_file(file): | ||
| 115 | filename = secure_filename(file.filename) | ||
| 116 | filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) | ||
| 117 | file.save(os.path.join(filepath)) | ||
| 118 | return filename, filepath | ||
| 119 | |||
| 120 | def get_file_parser(filepath: str): | ||
| 121 | parser, mime = parser_factory.get_parser(filepath) | ||
| 122 | return parser, mime | ||
| 123 | |||
| 124 | def cleanup(parser, filepath): | ||
| 125 | output_filename = os.path.basename(parser.output_filename) | ||
| 126 | parser, _ = parser_factory.get_parser(parser.output_filename) | ||
| 127 | meta_after = parser.get_meta() | ||
| 128 | os.remove(filepath) | ||
| 129 | |||
| 130 | key = utils.hash_file(os.path.join(app.config['UPLOAD_FOLDER'], output_filename)) | ||
| 131 | return key, meta_after, output_filename | ||
| 132 | |||
| 133 | def get_file_paths(filename): | ||
| 134 | filepath = secure_filename(filename) | ||
| 135 | |||
| 136 | complete_path = os.path.join(app.config['UPLOAD_FOLDER'], filepath) | ||
| 137 | return complete_path, filepath | ||
| 138 | |||
| 139 | def is_valid_api_download_file(filename, key): | ||
| 140 | if filename != secure_filename(filename): | ||
| 141 | abort(400, message='Insecure filename') | ||
| 142 | |||
| 143 | complete_path, filepath = get_file_paths(filename) | ||
| 144 | |||
| 145 | if not os.path.exists(complete_path): | ||
| 146 | abort(404, message='File not found') | ||
| 147 | |||
| 148 | if hmac.compare_digest(utils.hash_file(complete_path), key) is False: | ||
| 149 | abort(400, message='The file hash does not match') | ||
| 150 | return complete_path, filepath | ||
| 151 | |||
| 152 | class APIUpload(Resource): | ||
| 153 | |||
| 154 | def post(self): | ||
| 155 | utils.check_upload_folder(app.config['UPLOAD_FOLDER']) | ||
| 156 | req_parser = reqparse.RequestParser() | ||
| 157 | req_parser.add_argument('file_name', type=str, required=True, help='Post parameter is not specified: file_name') | ||
| 158 | req_parser.add_argument('file', type=str, required=True, help='Post parameter is not specified: file') | ||
| 159 | |||
| 160 | args = req_parser.parse_args() | ||
| 161 | try: | ||
| 162 | file_data = base64.b64decode(args['file']) | ||
| 163 | except binascii.Error as err: | ||
| 164 | abort(400, message='Failed decoding file: ' + str(err)) | ||
| 165 | |||
| 166 | file = FileStorage(stream=io.BytesIO(file_data), filename=args['file_name']) | ||
| 167 | filename, filepath = save_file(file) | ||
| 168 | parser, mime = get_file_parser(filepath) | ||
| 169 | |||
| 170 | if parser is None: | ||
| 171 | abort(415, message='The type %s is not supported' % mime) | ||
| 172 | |||
| 173 | meta = parser.get_meta() | ||
| 174 | if not parser.remove_all(): | ||
| 175 | abort(500, message='Unable to clean %s' % mime) | ||
| 176 | |||
| 177 | key, meta_after, output_filename = cleanup(parser, filepath) | ||
| 178 | return utils.return_file_created_response( | ||
| 179 | output_filename, | ||
| 180 | mime, | ||
| 181 | key, | ||
| 182 | meta, | ||
| 183 | meta_after, | ||
| 184 | urljoin(request.host_url, '%s/%s/%s/%s' % ('api', 'download', key, output_filename)) | ||
| 185 | ) | ||
| 186 | |||
| 187 | class APIDownload(Resource): | ||
| 188 | def get(self, key: str, filename: str): | ||
| 189 | complete_path, filepath = is_valid_api_download_file(filename, key) | ||
| 190 | # Make sure the file is NOT deleted on HEAD requests | ||
| 191 | if request.method == 'GET': | ||
| 192 | file_removal_scheduler.run_file_removal_job(app.config['UPLOAD_FOLDER']) | ||
| 193 | @after_this_request | ||
| 194 | def remove_file(response): | ||
| 195 | if os.path.exists(complete_path): | ||
| 196 | os.remove(complete_path) | ||
| 197 | return response | ||
| 198 | |||
| 199 | return send_from_directory(app.config['UPLOAD_FOLDER'], filepath, as_attachment=True) | ||
| 200 | |||
| 201 | class APIBulkDownloadCreator(Resource): | ||
| 202 | schema = { | ||
| 203 | 'download_list': { | ||
| 204 | 'type': 'list', | ||
| 205 | 'minlength': 2, | ||
| 206 | 'maxlength': int(os.environ.get('MAT2_MAX_FILES_BULK_DOWNLOAD', 10)), | ||
| 207 | 'schema': { | ||
| 208 | 'type': 'dict', | ||
| 209 | 'schema': { | ||
| 210 | 'key': {'type': 'string', 'required': True}, | ||
| 211 | 'file_name': {'type': 'string', 'required': True} | ||
| 212 | } | ||
| 213 | } | ||
| 214 | } | ||
| 215 | } | ||
| 216 | v = Validator(schema) | ||
| 217 | |||
| 218 | def post(self): | ||
| 219 | utils.check_upload_folder(app.config['UPLOAD_FOLDER']) | ||
| 220 | data = request.json | ||
| 221 | if not self.v.validate(data): | ||
| 222 | abort(400, message=self.v.errors) | ||
| 223 | # prevent the zip file from being overwritten | ||
| 224 | zip_filename = 'files.' + str(uuid4()) + '.zip' | ||
| 225 | zip_path = os.path.join(app.config['UPLOAD_FOLDER'], zip_filename) | ||
| 226 | cleaned_files_zip = zipfile.ZipFile(zip_path, 'w') | ||
| 227 | with cleaned_files_zip: | ||
| 228 | for file_candidate in data['download_list']: | ||
| 229 | complete_path, file_path = is_valid_api_download_file( | ||
| 230 | file_candidate['file_name'], | ||
| 231 | file_candidate['key'] | ||
| 232 | ) | ||
| 233 | try: | ||
| 234 | cleaned_files_zip.write(complete_path) | ||
| 235 | os.remove(complete_path) | ||
| 236 | except ValueError: | ||
| 237 | abort(400, message='Creating the archive failed') | ||
| 238 | |||
| 239 | try: | ||
| 240 | cleaned_files_zip.testzip() | ||
| 241 | except ValueError as e: | ||
| 242 | abort(400, message=str(e)) | ||
| 243 | |||
| 244 | parser, mime = get_file_parser(zip_path) | ||
| 245 | if not parser.remove_all(): | ||
| 246 | abort(500, message='Unable to clean %s' % mime) | ||
| 247 | key, meta_after, output_filename = cleanup(parser, zip_path) | ||
| 248 | return { | ||
| 249 | 'output_filename': output_filename, | ||
| 250 | 'mime': mime, | ||
| 251 | 'key': key, | ||
| 252 | 'meta_after': meta_after, | ||
| 253 | 'download_link': urljoin(request.host_url, '%s/%s/%s/%s' % ('api', 'download', key, output_filename)) | ||
| 254 | }, 201 | ||
| 255 | |||
| 256 | class APISupportedExtensions(Resource): | ||
| 257 | def get(self): | ||
| 258 | return get_supported_extensions() | ||
| 259 | |||
| 260 | api.add_resource(APIUpload, '/api/upload') | ||
| 261 | api.add_resource(APIDownload, '/api/download/<string:key>/<string:filename>') | ||
| 262 | api.add_resource(APIBulkDownloadCreator, '/api/download/bulk') | ||
| 263 | api.add_resource(APISupportedExtensions, '/api/extension') | ||
| 264 | 48 | ||
| 265 | return app | 49 | return app |
| 266 | 50 | ||
