Merge pull request #5 from ftruzzi/drm

Add DRM check for .epub files, tests

libgen_uploader/__main__.py

@@ -33,8 +33,12 @@ def main(args):
         else "python -m {}".format(__spec__.name.partition(".")[0])  # type: ignore
     )
     upload_type = parser.add_mutually_exclusive_group(required=True)
-    upload_type.add_argument("--scitech", action="store_true")
-    upload_type.add_argument("--fiction", action="store_true")
+    upload_type.add_argument(
+        "--scitech", action="store_true", help="Upload to scitech (non-fiction) library"
+    )
+    upload_type.add_argument(
+        "--fiction", action="store_true", help="Upload to fiction library"
+    )
     parser.add_argument(
         "--metadata-source",
         type=str,

libgen_uploader/helpers.py

@@ -9,11 +9,13 @@
 
 
 class LibgenUploadException(Exception):
-    pass
+    def __init__(self, message: str):
+        self.message = message
 
 
 class LibgenMetadataException(Exception):
-    pass
+    def __init__(self, message: str):
+        self.message = message
 
 
 def calculate_md5(file_path: str):
@@ -74,6 +76,20 @@ def are_forms_equal(first: Form, second: Form) -> bool:
     return True
 
 
+def epub_has_drm(book: Union[str, bytes]) -> bool:
+    from io import BytesIO
+    from zipfile import BadZipFile, ZipFile
+
+    book_file = BytesIO(book) if isinstance(book, bytes) else book
+
+    try:
+        z = ZipFile(book_file)  # type: ignore
+        return any("encryption.xml" in f.filename for f in z.filelist)
+    except BadZipFile:
+        # assuming not .epub
+        return False
+
+
 def validate_metadata(metadata) -> Union[bool, dict]:
     from cerberus import Validator
     from .constants import METADATA_FORM_SCHEMA

libgen_uploader/libgen_uploader.py

@@ -35,6 +35,7 @@
     are_forms_equal,
     check_upload_form_response,
     check_metadata_form_response,
+    epub_has_drm,
     match_language_to_form_option,
     validate_metadata,
 )
@@ -111,12 +112,18 @@ def _submit_and_check_form(self, form: Form) -> Result[str, Exception]:
     def _validate_file(file: Union[str, bytes]) -> Union[str, bytes]:
         if isinstance(file, bytes):
             # TODO add file data validation?
+            if epub_has_drm(file):
+                raise LibgenUploadException("Your .epub file seems to have DRM.")
+
             return file
 
         if isinstance(file, str):
             if not os.path.isfile(file):
                 raise FileNotFoundError(f"Upload failed: {file} is not a file.")
 
+            if file.endswith(".epub") and epub_has_drm(file):
+                raise LibgenUploadException("Your .epub file seems to have DRM.")
+
             logging.info(f"Selected file {basename(file)}")
             return file
 

tests/test_upload.py

@@ -1,15 +1,16 @@
-import libgen_uploader
-from libgen_uploader.libgen_uploader import LibgenMetadata, LibgenUploader
+from libgen_uploader import LibgenMetadata, LibgenUploader
+from libgen_uploader.helpers import LibgenUploadException, check_upload_form_response
 
 from functools import partial
-from helpers import get_return_value
+from .helpers import get_return_value
 
 import os
 
 import pytest
 
-from returns.result import Success, Failure
 from returns.contrib.pytest import ReturnsAsserts
+from returns.pipeline import is_successful
+from returns.result import Success, Failure
 
 
 files_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "files")
@@ -27,12 +28,16 @@ def test_validate_file_path(uploader: LibgenUploader):
     assert uploader._validate_file(file_path) == Success(file_path)
 
 
+def test_epub_drm(uploader: LibgenUploader):
+    file_path = os.path.join(files_path, "minimal_drm.epub")
+    result = uploader.upload_fiction(file_path=file_path)
+    assert is_successful(result) is False and "drm" in str(result.failure()).lower()
+
+
 @pytest.mark.vcr(record_mode="none")
 def test_file_upload(uploader: LibgenUploader, returns: ReturnsAsserts):
     file_path = os.path.join(files_path, "minimal.epub")
-    with returns.assert_trace(
-        Success, libgen_uploader.helpers.check_upload_form_response
-    ):
+    with returns.assert_trace(Success, check_upload_form_response):
         uploader.upload_fiction(file_path)
 
 
@@ -42,10 +47,8 @@ def test_bytes_upload(uploader: LibgenUploader, returns: ReturnsAsserts):
     with open(file_path, "rb") as f:
         data = f.read()
 
-    with returns.assert_trace(
-        Success, libgen_uploader.helpers.check_upload_form_response
-    ):
-        uploader.upload_fiction(file_path)
+    with returns.assert_trace(Success, check_upload_form_response):
+        uploader.upload_fiction(data)
 
 
 @pytest.mark.vcr(record_mode="none")
@@ -126,7 +129,3 @@ def test_custom_metadata_overrides_fetched(uploader: LibgenUploader):
         "Dante Alighieri" in form["authors"].value
         and form["title"].value == "custom title"
     )
-
-
-# @pytest.mark.vcr(record_mode="none")
-# def upload_book(uploader):