1#!/usr/local/bin/python3.8
2# vim:fileencoding=UTF-8:ts=4:sw=4:sta:et:sts=4:ai
3# License: GPLv3 Copyright: 2015, Kovid Goyal <kovid at kovidgoyal.net>
4
5import os
6import shutil
7import tempfile
8import time
9from qt.core import (
10    QApplication, QCursor, QFileSystemWatcher, QObject, Qt, QTimer, pyqtSignal
11)
12from threading import Event, Thread
13
14from calibre import prints
15from calibre.db.adding import compile_rule, filter_filename
16from calibre.ebooks import BOOK_EXTENSIONS
17from calibre.gui2 import gprefs
18from calibre.gui2.dialogs.duplicates import DuplicatesQuestion
19from calibre.utils.tdir_in_cache import tdir_in_cache
20
21AUTO_ADDED = frozenset(BOOK_EXTENSIONS) - {'pdr', 'mbp', 'tan'}
22
23
24class AllAllowed:
25
26    def __init__(self):
27        self.disallowed = frozenset(gprefs['blocked_auto_formats'])
28
29    def __contains__(self, x):
30        return x not in self.disallowed
31
32
33def allowed_formats():
34    ' Return an object that can be used to test if a format (lowercase) is allowed for auto-adding '
35    if gprefs['auto_add_everything']:
36        allowed = AllAllowed()
37    else:
38        allowed = AUTO_ADDED - frozenset(gprefs['blocked_auto_formats'])
39    return allowed
40
41
42class Worker(Thread):
43
44    def __init__(self, path, callback):
45        Thread.__init__(self)
46        self.daemon = True
47        self.keep_running = True
48        self.wake_up = Event()
49        self.path, self.callback = path, callback
50        self.staging = set()
51        self.allowed = allowed_formats()
52        self.read_rules()
53
54    def read_rules(self):
55        try:
56            self.compiled_rules = tuple(map(compile_rule, gprefs.get('add_filter_rules', ())))
57        except Exception:
58            self.compiled_rules = ()
59            import traceback
60            traceback.print_exc()
61
62    def is_filename_allowed(self, filename):
63        allowed = filter_filename(self.compiled_rules, filename)
64        if allowed is None:
65            ext = os.path.splitext(filename)[1][1:].lower()
66            allowed = ext in self.allowed
67        return allowed
68
69    def run(self):
70        self.tdir = tdir_in_cache('aa')
71        try:
72            while self.keep_running:
73                self.wake_up.wait()
74                self.wake_up.clear()
75                if not self.keep_running:
76                    break
77                try:
78                    self.auto_add()
79                except:
80                    import traceback
81                    traceback.print_exc()
82        finally:
83            shutil.rmtree(self.tdir, ignore_errors=True)
84
85    def auto_add(self):
86        from calibre.ebooks.metadata.meta import metadata_from_filename
87        from calibre.ebooks.metadata.opf2 import metadata_to_opf
88        from calibre.utils.ipc.simple_worker import WorkerError, fork_job
89
90        files = [x for x in os.listdir(self.path) if
91                    # Must not be in the process of being added to the db
92                    x not in self.staging and
93                    # Firefox creates 0 byte placeholder files when downloading
94                    os.stat(os.path.join(self.path, x)).st_size > 0 and
95                    # Must be a file
96                    os.path.isfile(os.path.join(self.path, x)) and
97                    # Must have read and write permissions
98                    os.access(os.path.join(self.path, x), os.R_OK|os.W_OK) and
99                    # Must be a known ebook file type
100                    self.is_filename_allowed(x)
101                ]
102        data = []
103        # Give any in progress copies time to complete
104        time.sleep(2)
105
106        def safe_mtime(x):
107            try:
108                return os.path.getmtime(os.path.join(self.path, x))
109            except OSError:
110                return time.time()
111
112        for fname in sorted(files, key=safe_mtime):
113            f = os.path.join(self.path, fname)
114
115            # Try opening the file for reading, if the OS prevents us, then at
116            # least on windows, it means the file is open in another
117            # application for writing. We will get notified by
118            # QFileSystemWatcher when writing is completed, so ignore for now.
119            try:
120                open(f, 'rb').close()
121            except:
122                continue
123            tdir = tempfile.mkdtemp(dir=self.tdir)
124            try:
125                fork_job('calibre.ebooks.metadata.meta',
126                        'forked_read_metadata', (f, tdir), no_output=True)
127            except WorkerError as e:
128                prints('Failed to read metadata from:', fname)
129                prints(e.orig_tb)
130            except:
131                import traceback
132                traceback.print_exc()
133
134            # Ensure that the pre-metadata file size is present. If it isn't,
135            # write 0 so that the file is rescanned
136            szpath = os.path.join(tdir, 'size.txt')
137            try:
138                with open(szpath, 'rb') as f:
139                    int(f.read())
140            except:
141                with open(szpath, 'wb') as f:
142                    f.write(b'0')
143
144            opfpath = os.path.join(tdir, 'metadata.opf')
145            try:
146                if os.stat(opfpath).st_size < 30:
147                    raise Exception('metadata reading failed')
148            except:
149                mi = metadata_from_filename(fname)
150                with open(opfpath, 'wb') as f:
151                    f.write(metadata_to_opf(mi))
152            self.staging.add(fname)
153            data.append((fname, tdir))
154        if data:
155            self.callback(data)
156
157
158class AutoAdder(QObject):
159
160    metadata_read = pyqtSignal(object)
161    auto_convert = pyqtSignal(object)
162
163    def __init__(self, path, parent):
164        QObject.__init__(self, parent)
165        if path and os.path.isdir(path) and os.access(path, os.R_OK|os.W_OK):
166            self.watcher = QFileSystemWatcher(self)
167            self.worker = Worker(path, self.metadata_read.emit)
168            self.watcher.directoryChanged.connect(self.dir_changed,
169                    type=Qt.ConnectionType.QueuedConnection)
170            self.metadata_read.connect(self.add_to_db,
171                    type=Qt.ConnectionType.QueuedConnection)
172            QTimer.singleShot(2000, self.initialize)
173            self.auto_convert.connect(self.do_auto_convert,
174                    type=Qt.ConnectionType.QueuedConnection)
175        elif path:
176            prints(path,
177                'is not a valid directory to watch for new ebooks, ignoring')
178
179    def read_rules(self):
180        if hasattr(self, 'worker'):
181            self.worker.read_rules()
182
183    def initialize(self):
184        try:
185            if os.listdir(self.worker.path):
186                self.dir_changed()
187        except:
188            pass
189        self.watcher.addPath(self.worker.path)
190
191    def dir_changed(self, *args):
192        if os.path.isdir(self.worker.path) and os.access(self.worker.path,
193                os.R_OK|os.W_OK):
194            if not self.worker.is_alive():
195                self.worker.start()
196            self.worker.wake_up.set()
197
198    def stop(self):
199        if hasattr(self, 'worker'):
200            self.worker.keep_running = False
201            self.worker.wake_up.set()
202
203    def wait(self):
204        if hasattr(self, 'worker'):
205            self.worker.join()
206
207    def __enter__(self):
208        QApplication.setOverrideCursor(QCursor(Qt.CursorShape.WaitCursor))
209
210    def __exit__(self, *args):
211        QApplication.restoreOverrideCursor()
212
213    def add_to_db(self, data):
214        with self:
215            self.do_add(data)
216
217    def do_add(self, data):
218        from calibre.ebooks.metadata.opf2 import OPF
219
220        gui = self.parent()
221        if gui is None:
222            return
223        m = gui.library_view.model()
224        count = 0
225
226        needs_rescan = False
227        duplicates = []
228        added_ids = set()
229
230        for fname, tdir in data:
231            path_to_remove = os.path.join(self.worker.path, fname)
232            paths = [path_to_remove]
233            fpath = os.path.join(tdir, 'file_changed_by_plugins')
234            if os.path.exists(fpath):
235                with open(fpath) as f:
236                    paths[0] = f.read()
237            sz = os.path.join(tdir, 'size.txt')
238            try:
239                with open(sz, 'rb') as f:
240                    sz = int(f.read())
241                if sz != os.stat(paths[0]).st_size:
242                    raise Exception('Looks like the file was written to after'
243                            ' we tried to read metadata')
244            except:
245                needs_rescan = True
246                try:
247                    self.worker.staging.remove(fname)
248                except KeyError:
249                    pass
250
251                continue
252
253            mi = os.path.join(tdir, 'metadata.opf')
254            if not os.access(mi, os.R_OK):
255                continue
256            mi = OPF(open(mi, 'rb'), tdir, populate_spine=False).to_book_metadata()
257            if gprefs.get('tag_map_on_add_rules'):
258                from calibre.ebooks.metadata.tag_mapper import map_tags
259                mi.tags = map_tags(mi.tags, gprefs['tag_map_on_add_rules'])
260            if gprefs.get('author_map_on_add_rules'):
261                from calibre.ebooks.metadata.author_mapper import (
262                    compile_rules, map_authors
263                )
264                new_authors = map_authors(mi.authors, compile_rules(gprefs['author_map_on_add_rules']))
265                if new_authors != mi.authors:
266                    mi.authors = new_authors
267                    mi.author_sort = gui.current_db.new_api.author_sort_from_authors(mi.authors)
268            mi = [mi]
269            dups, ids = m.add_books(paths,
270                    [os.path.splitext(fname)[1][1:].upper()], mi,
271                    add_duplicates=not gprefs['auto_add_check_for_duplicates'],
272                    return_ids=True)
273            added_ids |= set(ids)
274            num = len(ids)
275            if dups:
276                path = dups[0][0]
277                with open(os.path.join(tdir, 'dup_cache.'+dups[1][0].lower()),
278                        'wb') as dest, open(path, 'rb') as src:
279                    shutil.copyfileobj(src, dest)
280                    dups[0][0] = dest.name
281                duplicates.append(dups)
282
283            try:
284                os.remove(path_to_remove)
285                self.worker.staging.remove(fname)
286            except:
287                import traceback
288                traceback.print_exc()
289            count += num
290
291        if duplicates:
292            paths, formats, metadata = [], [], []
293            for p, f, mis in duplicates:
294                paths.extend(p)
295                formats.extend(f)
296                metadata.extend(mis)
297            dups = [(mic, mic.cover, [p]) for mic, p in zip(metadata, paths)]
298            d = DuplicatesQuestion(m.db, dups, parent=gui)
299            dups = tuple(d.duplicates)
300            if dups:
301                paths, formats, metadata = [], [], []
302                for mi, cover, book_paths in dups:
303                    paths.extend(book_paths)
304                    formats.extend([p.rpartition('.')[-1] for p in book_paths])
305                    metadata.extend([mi for i in book_paths])
306                ids = m.add_books(paths, formats, metadata,
307                        add_duplicates=True, return_ids=True)[1]
308                added_ids |= set(ids)
309                num = len(ids)
310                count += num
311
312        for fname, tdir in data:
313            try:
314                shutil.rmtree(tdir)
315            except:
316                pass
317
318        if added_ids and gprefs['auto_add_auto_convert']:
319            self.auto_convert.emit(added_ids)
320
321        if count > 0:
322            m.books_added(count)
323            gui.status_bar.show_message(
324                (_('Added a book automatically from {src}') if count == 1 else _('Added {num} books automatically from {src}')).format(
325                    num=count, src=self.worker.path), 2000)
326            gui.refresh_cover_browser()
327
328        if needs_rescan:
329            QTimer.singleShot(2000, self.dir_changed)
330
331    def do_auto_convert(self, added_ids):
332        gui = self.parent()
333        gui.iactions['Convert Books'].auto_convert_auto_add(added_ids)
334