feat: Prevent zipbomb files (#231)

* feat: Prevent zipbomb files

- Added a limit of the zip file content to be 5MB max
- Added a unittest
- Added to the pytest ignore directories to node_modules directory

Author: orronai

lms/extractors/ziparchive.py

@@ -5,6 +5,8 @@
 from zipfile import BadZipFile, ZipFile
 
 from lms.extractors.base import Extractor, File
+from lms.models.errors import FileSizeError
+from lms.utils.consts import MAX_ZIP_CONTENT_SIZE
 from lms.utils.log import log
 
 
@@ -29,6 +31,11 @@ def __init__(self, **kwargs):
     def can_extract(self) -> bool:
         return self.is_zipfile
 
+    def check_files_size(self) -> bool:
+        return sum(
+            f.file_size for f in self.archive.infolist()
+        ) <= MAX_ZIP_CONTENT_SIZE
+
     @staticmethod
     def _extract(archive: ZipFile, filename: str, dirname: str = '') -> File:
         with archive.open(filename) as current_file:
@@ -57,7 +64,7 @@ def get_exercises_by_dirs(
     ) -> Iterator[Tuple[int, List[File]]]:
         for dirname in filenames:
             if len(dirname.strip(os.path.sep).split(os.path.sep)) == 1:
-                # Checking if the dirname is in the first dir in the zipfile
+                # Checking if the dirname is in the first dir of the zipfile
                 parent_name, _ = os.path.split(dirname)
                 exercise_id, _ = self._clean(parent_name)
                 if exercise_id:
@@ -79,6 +86,11 @@ def get_exercise(self, file: ZipFile) -> Iterator[Tuple[int, List[File]]]:
                 yield from self.get_exercises_by_dirs(archive, filenames)
 
     def get_exercises(self) -> Iterator[Tuple[int, List[File]]]:
+        if not self.check_files_size():
+            raise FileSizeError(
+                'File content is too big. '
+                f'{MAX_ZIP_CONTENT_SIZE // 1000000}MB allowed.',
+            )
         for exercise_id, files in self.get_exercise(self.archive):
             if exercise_id and files and any(file.code for file in files):
                 yield (exercise_id, files)

lms/lmsweb/views.py

@@ -26,7 +26,7 @@
     MAX_REQUEST_SIZE, PERMISSIVE_CORS, get_next_url, login_manager,
 )
 from lms.models import comments, notifications, share_link, solutions, upload
-from lms.models.errors import LmsError, UploadError, fail
+from lms.models.errors import FileSizeError, LmsError, UploadError, fail
 from lms.utils.consts import RTL_LANGUAGES
 from lms.utils.files import get_language_name_by_extension
 from lms.utils.log import log
@@ -292,7 +292,7 @@ def upload_page():
         return fail(404, 'User not found.')
     if request.content_length > MAX_REQUEST_SIZE:
         return fail(
-            413, f'File is too big. {MAX_REQUEST_SIZE // 1000000}MB allowed',
+            413, f'File is too big. {MAX_REQUEST_SIZE // 1000000}MB allowed.',
         )
 
     file: Optional[FileStorage] = request.files.get('file')
@@ -304,6 +304,9 @@ def upload_page():
     except UploadError as e:
         log.debug(e)
         return fail(400, str(e))
+    except FileSizeError as e:
+        log.debug(e)
+        return fail(413, str(e))
 
     return jsonify({
         'exercise_matches': matches,

lms/models/errors.py

@@ -17,6 +17,10 @@ class BadUploadFile(LmsError):
     pass
 
 
+class FileSizeError(LmsError):
+    pass
+
+
 def fail(status_code: int, error_msg: str):
     data = {
         'status': 'failed',

lms/tests/samples/zipbomb.zip

@@ -1,6 +1,6 @@
 from io import BufferedReader, BytesIO
 from tempfile import SpooledTemporaryFile
-from typing import Iterator, Tuple
+from typing import Iterator, List, Tuple
 from zipfile import ZipFile
 
 from flask import json
@@ -18,6 +18,7 @@ class TestExtractor:
     IGNORE_FILES_ZIP_NAME = 'Upload_123.zip'
     PY_NAMES = ('code1.py', 'code2.py')
     ZIP_FILES = ('Upload_1.zip', 'zipfiletest.zip')
+    ZIP_BOMB_FILE = 'zipbomb.zip'
 
     def setup(self):
         self.ipynb_file = self.ipynb_file()
@@ -32,11 +33,18 @@ def setup(self):
             self.zipfile_file, self.IGNORE_FILES_ZIP_NAME,
         )
         self.zipfiles_extractor_files = list(self.zip_files(self.ZIP_FILES))
-        self.zipfiles_extractors_bytes_io = list(self.get_bytes_io_zip_files())
+        self.zipfiles_extractors_bytes_io = list(self.get_bytes_io_zip_files(
+            self.zipfiles_extractor_files, self.ZIP_FILES,
+        ))
+        self.zipbomb_file_list = list(self.zip_files((self.ZIP_BOMB_FILE,)))
+        self.zipbomb_bytes_io = next(self.get_bytes_io_zip_files(
+            self.zipbomb_file_list, (self.ZIP_BOMB_FILE,),
+        ))
 
     def teardown(self):
         self.ipynb_file.close()
         self.zipfile_file.close()
+        self.zipbomb_file_list[0].close()
         for py_file in self.pyfiles_files:
             py_file.close()
         for zip_file in self.zipfiles_extractor_files:
@@ -67,6 +75,13 @@ def create_zipfile_storage(
         opened_file.seek(0)
         return zip_file_storage
 
+    @staticmethod
+    def get_bytes_io_zip_files(
+        files: List[BufferedReader], filesnames: Tuple[str, ...],
+    ) -> Iterator[Tuple[BytesIO, str]]:
+        for file, name in zip(files, filesnames):
+            yield BytesIO(file.read()), name
+
     def get_zip_filenames(self):
         the_zip = ZipFile(f'{SAMPLES_DIR}/{self.IGNORE_FILES_ZIP_NAME}')
         return the_zip.namelist()
@@ -106,10 +121,6 @@ def test_zip_ignore_files(self):
             for filename in original_zip_filenames
         )
 
-    def get_bytes_io_zip_files(self) -> Iterator[Tuple[BytesIO, str]]:
-        for file, name in zip(self.zipfiles_extractor_files, self.ZIP_FILES):
-            yield BytesIO(file.read()), name
-
     def test_zip(self, student_user: User):
         conftest.create_exercise()
         conftest.create_exercise()
@@ -134,3 +145,14 @@ def test_zip(self, student_user: User):
             file=self.zipfiles_extractors_bytes_io[0],
         ))
         assert second_upload_response.status_code == 400
+
+    def test_zip_bomb(self, student_user: User):
+        conftest.create_exercise()
+
+        client = conftest.get_logged_user(username=student_user.username)
+
+        # Trying to upload a zipbomb file
+        upload_response = client.post('/upload', data={
+            'file': self.zipbomb_bytes_io,
+        })
+        assert upload_response.status_code == 413

lms/tests/test_extractor.py

@@ -2,3 +2,6 @@
     'he', 'ar', 'arc', 'dv', 'fa', 'ha',
     'khw', 'ks', 'ku', 'ps', 'ur', 'yi',
 }
+
+
+MAX_ZIP_CONTENT_SIZE = 5_000_000  # 5MB (in bytes)

lms/utils/consts.py

@@ -1,3 +1,4 @@
 [pytest]
 env =
     LOCAL_SETUP=true
+norecursedirs = node_modules/*