From a70838348b577a957f685d81685da0a63af62e38 Mon Sep 17 00:00:00 2001
From: BasioMeusPuga <disgruntled.mob@gmail.com>
Date: Fri, 9 Mar 2018 17:25:44 +0530
Subject: [PATCH] Implement support for Amazon ebooks

---
 TODO                          |  4 +++-
 __main__.py                   | 12 ++++++++----
 parsers/{epub.py => ebook.py} | 34 +++++++++++++++++++++++++++++++++-
 sorter.py                     |  8 +++++---
 threaded.py                   |  3 ++-
 5 files changed, 51 insertions(+), 10 deletions(-)
 rename parsers/{epub.py => ebook.py} (83%)

diff --git a/TODO b/TODO
index eb44e45..a7c2b02 100644
--- a/TODO
+++ b/TODO
@@ -54,8 +54,10 @@ TODO
     Filetypes:
         ✓ cbz, cbr support
             ✓ Keep font settings enabled but only for background color
+        ✓ mobi, azw support
+            Limit the extra files produced by KindleUnpack
+                Have them save to memory
         epub support
-        mobi, azw support
     Other:
         ✓ Define every widget in code
     Bugs:
diff --git a/__main__.py b/__main__.py
index 63af4ab..98a028f 100755
--- a/__main__.py
+++ b/__main__.py
@@ -288,7 +288,8 @@ class MainUI(QtWidgets.QMainWindow, mainwindow.Ui_MainWindow):
                 file_list,
                 'addition',
                 self.database_path,
-                self.settings['auto_tags'])
+                self.settings['auto_tags'],
+                self.temp_dir.path())
 
             parsed_books = books.initiate_threads()
             database.DatabaseFunctions(self.database_path).add_to_database(parsed_books)
@@ -704,9 +705,12 @@ class MainUI(QtWidgets.QMainWindow, mainwindow.Ui_MainWindow):
 
             file_md5 = filename[1]
             if not file_md5:
-                with open(filename[0], 'rb') as current_book:
-                    first_bytes = current_book.read(1024 * 32)  # First 32KB of the file
-                    file_md5 = hashlib.md5(first_bytes).hexdigest()
+                try:
+                    with open(filename[0], 'rb') as current_book:
+                        first_bytes = current_book.read(1024 * 32)  # First 32KB of the file
+                        file_md5 = hashlib.md5(first_bytes).hexdigest()
+                except FileNotFoundError:
+                    return
 
             # Remove any already open files
             # Set focus to last file in case only one is open
diff --git a/parsers/epub.py b/parsers/ebook.py
similarity index 83%
rename from parsers/epub.py
rename to parsers/ebook.py
index a7dcc51..940abd6 100644
--- a/parsers/epub.py
+++ b/parsers/ebook.py
@@ -18,23 +18,46 @@
 
 import os
 import re
+import sys
 import zipfile
 import collections
 from urllib.parse import unquote
 
 import ebooklib.epub
+import KindleUnpack.kindleunpack as KindleUnpack
 
 
-class ParseEPUB:
+class ParseEBook:
     def __init__(self, filename, temp_dir, file_md5):
         # TODO
         # Maybe also include book description
         self.filename = filename
         self.book = None
         self.temp_dir = temp_dir
+        self.temp_dir_copy = temp_dir
         self.file_md5 = file_md5
 
+        # This is a crazy lazy thing
+        # But it works for now
+
+        self.use_KindleUnpack = False
+        kindle_extensions = ['.mobi', '.azw', '.azw3', '.azw4', '.prc']
+        if os.path.splitext(self.filename)[1].lower() in kindle_extensions:
+            self.use_KindleUnpack = True
+            self.temp_dir = os.path.join(
+                self.temp_dir, os.path.basename(self.filename))
+
     def read_book(self):
+        if self.use_KindleUnpack:
+            with HidePrinting():
+                KindleUnpack.unpackBook(self.filename, self.temp_dir)
+
+            new_filename_with_ext = os.path.splitext(
+                os.path.basename(self.filename))[0] + '.epub'
+            self.filename = os.path.join(
+                self.temp_dir, 'mobi8', new_filename_with_ext)
+            self.temp_dir = self.temp_dir_copy
+
         try:
             self.book = ebooklib.epub.read_epub(self.filename)
         except (KeyError, AttributeError, FileNotFoundError):
@@ -162,3 +185,12 @@ class ParseEPUB:
             'images_only': False}
 
         return contents, file_settings
+
+
+class HidePrinting:
+    def __enter__(self):
+        self._original_stdout = sys.stdout
+        sys.stdout = None
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout = self._original_stdout
diff --git a/sorter.py b/sorter.py
index 0de5aff..5ed3f24 100644
--- a/sorter.py
+++ b/sorter.py
@@ -44,14 +44,16 @@ from PyQt5 import QtCore, QtGui
 
 import database
 
-from parsers.epub import ParseEPUB
+from parsers.ebook import ParseEBook
 from parsers.cbz import ParseCBZ
 from parsers.cbr import ParseCBR
 
 sorter = {
-    'epub': ParseEPUB,
+    'epub': ParseEBook,
+    'mobi': ParseEBook,
+    'azw': ParseEBook,
     'cbz': ParseCBZ,
-    'cbr': ParseCBR}
+    'cbr': ParseCBR,}
 
 available_parsers = [i for i in sorter]
 progressbar = None  # This is populated by __main__
diff --git a/threaded.py b/threaded.py
index eabea34..b375d56 100644
--- a/threaded.py
+++ b/threaded.py
@@ -56,7 +56,8 @@ class BackGroundBookAddition(QtCore.QThread):
             self.file_list,
             'addition',
             self.database_path,
-            self.parent.settings['auto_tags'])
+            self.parent.settings['auto_tags'],
+            self.parent.temp_dir.path())
         parsed_books = books.initiate_threads()
         self.parent.lib_ref.generate_model('addition', parsed_books, False)
         if self.prune_required: