openmedialibrary/oml/item/handlers.py

# -*- coding: utf-8 -*-


from datetime import datetime
import mimetypes
import os
from urllib.request import quote
import zipfile
import base64

import ox

from .models import Item, File
from user.models import List
from .scan import add_file
import db
import settings
import tornado.web
import tornado.gen
from concurrent.futures import ThreadPoolExecutor
from tornado.concurrent import run_on_executor

from oxtornado import json_dumps, json_response

from media import get_id
import state

import logging
logger = logging.getLogger(__name__)

MAX_WORKERS = 4


class OptionalBasicAuthMixin(object):
    class SendChallenge(Exception):
        pass

    def prepare(self):
        if settings.preferences.get('authentication'):
            try:
                self.authenticate_user()
            except self.SendChallenge:
                self.send_auth_challenge()

    def send_auth_challenge(self):
        realm = "Open Media Library"
        hdr = 'Basic realm="%s"' % realm
        self.set_status(401)
        self.set_header('www-authenticate', hdr)
        self.finish()
        return False

    def authenticate_user(self):
        auth_header = self.request.headers.get('Authorization')
        if not auth_header or not auth_header.startswith('Basic '):
            raise self.SendChallenge()

        auth_data = auth_header.split(None, 1)[-1]
        auth_data = base64.b64decode(auth_data).decode('ascii')
        username, password = auth_data.split(':', 1)

        auth = settings.preferences.get('authentication')
        if auth.get('username') == username and auth.get('password') == password:
            self._current_user = username
        else:
            raise self.SendChallenge()

class OMLHandler(OptionalBasicAuthMixin, tornado.web.RequestHandler):

    def initialize(self):
        pass

class EpubHandler(OMLHandler):

    def get(self, id, filename):
        with db.session():
            item = Item.get(id)
            path = item.get_path()
        if not item or item.info['extension'] != 'epub' or not path:
            self.set_status(404)
            self.write('')
        else:
            z = zipfile.ZipFile(path)
            if filename == '':
                self.write('<br>\n'.join([f.filename for f in z.filelist]))
            elif filename not in [f.filename for f in z.filelist]:
                self.set_status(404)
                self.write('')
            else:
                content_type = {
                    'xpgt': 'application/vnd.adobe-page-template+xml'
                }.get(filename.split('.')[0], mimetypes.guess_type(filename)[0]) or 'text/plain'
                self.set_header('Content-Type', content_type)
                self.write(z.read(filename))

def serve_static(handler, path, mimetype, include_body=True, disposition=None):
    handler.set_header('Content-Type', mimetype)
    size = os.stat(path).st_size
    handler.set_header('Accept-Ranges', 'bytes')
    if disposition:
        handler.set_header('Content-Disposition', "attachment; filename*=UTF-8''%s" % quote(disposition.encode('utf-8')))
    if include_body:
        if 'Range' in handler.request.headers:
            handler.set_status(206)
            r = handler.request.headers.get('Range').split('=')[-1].split('-')
            start = int(r[0])
            end = int(r[1]) if r[1] else (size - 1)
            length = end - start + 1
            handler.set_header('Content-Length', str(length))
            handler.set_header('Content-Range', 'bytes %s-%s/%s' % (start, end, size))
            with open(path, 'rb') as fd:
                fd.seek(start)
                handler.write(fd.read(length))
        else:
            handler.set_header('Content-Length', str(size))
            with open(path, 'rb') as fd:
                handler.write(fd.read())
    else:
        handler.set_header('Content-Length', str(size))
    return

class FileHandler(OMLHandler):

    def initialize(self, attachment=False):
        self._attachment = attachment

    def head(self, id):
        self.get(id, include_body=False)

    def get(self, id, include_body=True):
        with db.session():
            item = Item.get(id)
            path = item.get_path() if item else None
            if not item or not path:
                self.set_status(404)
                return
            mimetype = {
                'cbr': 'application/x-cbr',
                'cbz': 'application/x-cbz',
                'epub': 'application/epub+zip',
                'pdf': 'application/pdf',
                'txt': 'text/plain',
            }.get(path.split('.')[-1], None)
            if mimetype == 'text/plain':
                try:
                    open(path, 'rb').read().decode('utf-8')
                    mimetype = 'text/plain; charset=utf-8'
                except:
                    mimetype = 'text/plain; charset=latin-1'
            if self._attachment:
                disposition = os.path.basename(path)
            else:
                disposition = None
        return serve_static(self, path, mimetype, include_body, disposition=disposition)

class ReaderHandler(OMLHandler):

    def get(self, id):
        with db.session():
            item = Item.get(id)
            if not item:
                self.set_status(404)
                return
            if item.info['extension'] in ('cbr', 'cbz'):
                html = 'html/cbr.html'
            elif item.info['extension'] == 'epub':
                html = 'html/epub.html'
            elif item.info['extension'] == 'pdf':
                html = 'html/pdf.html'
            elif item.info['extension'] == 'txt':
                html = 'html/txt.html'
            else:
                self.set_status(404)
                return
            item.accessed = datetime.utcnow()
            item.timesaccessed = (item.timesaccessed or 0) + 1
            item.update_sort()
            item.save()
            path = os.path.join(settings.static_path, html)
        return serve_static(self, path, 'text/html')

class UploadHandler(OMLHandler):
    executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)

    def initialize(self, context=None):
        self._context = context

    def get(self):
        self.write('use POST')

    @run_on_executor
    def save_files(self, request):
        listname = request.arguments.get('list', None)
        if listname:
            listname = listname[0]
            if isinstance(listname, bytes):
                listname = listname.decode('utf-8')
        with self._context():
            prefs = settings.preferences
            ids = []
            for upload in request.files.get('files', []):
                filename = upload.filename
                id = get_id(data=upload.body)
                ids.append(id)
                file = File.get(id)
                if not file or not os.path.exists(file.fullpath()):
                    logger.debug('add %s to library', id)
                    prefix_books = os.path.join(os.path.expanduser(prefs['libraryPath']), 'Books' + os.sep)
                    prefix_imported = os.path.join(prefix_books, '.import' + os.sep)
                    ox.makedirs(prefix_imported)
                    import_name = os.path.join(prefix_imported, filename)
                    n = 1
                    while os.path.exists(import_name):
                        n += 1
                        name, extension = filename.rsplit('.', 1)
                        if extension == 'kepub':
                            extension = 'epub'
                        import_name = os.path.join(prefix_imported, '%s [%d].%s' % (name, n, extension))
                    with open(import_name, 'wb') as fd:
                        fd.write(upload.body)
                    file = add_file(id, import_name, prefix_books)
                    file.move()
                else:
                    user = state.user()
                    if not file.item:
                        item = Item.get_or_create(id=file.sha1, info=file.info)
                        file.item_id = item.id
                        state.db.session.add(file)
                        state.db.session.commit()
                    else:
                        item = file.item
                    if user not in item.users:
                        logger.debug('add %s to local user', id)
                        item.add_user(user)
                        add_record('additem', item.id, file.info)
                        add_record('edititem', item.id, item.meta)
                        item.update()
        if listname and ids:
            list_ = List.get(settings.USER_ID, listname)
            if list_:
                list_.add_items(ids)
        response = json_response({'ids': ids})
        return response

    @tornado.gen.coroutine
    def post(self):
        if 'origin' in self.request.headers and self.request.host not in self.request.headers['origin']:
            logger.debug('reject cross site attempt to access api %s', self.request)
            self.set_status(403)
            self.write('')
            return

        response = yield self.save_files(self.request)
        if 'status' not in response:
            response = json_response(response)
        response = json_dumps(response)
        self.set_header('Content-Type', 'application/json')
        self.write(response)