openmedialibrary/oml/media/cbr.py

# -*- coding: utf-8 -*-
import logging
import os
import zipfile

import ox

logger = logging.getLogger(__name__)
IMAGE_EXTENSIONS = ['.jpg', '.png', '.gif']


def filter_images(files):
    return [f for f in files if os.path.splitext(f)[-1].lower() in IMAGE_EXTENSIONS]

def detect_format(path):
    with open(path, 'rb') as fd:
        head = fd.read(10)
    if head[:2] == b'PK':
        return 'cbz'
    if head[:3] == b'Rar':
        return 'cbr'
    logger.debug('unknown cbr/cbz file %s - %s', head, path)
    return 'unknown'

def cover(path):
    format = detect_format(path)
    if format == 'cbz':
        cover = cover_cbz(path)
    elif format == 'cbr':
        cover = cover_cbr(path)
    else:
        cover = None
    return cover

def cover_cbr(path):
    data = None
    try:
        from unrar import rarfile
        rar = rarfile.RarFile(path)
        files = rar.namelist()
        files = filter_images(files)
        if files:
            cover = ox.sorted_strings(files)[0]
            data = rar.read(cover)
    except:
        logger.debug('invalid cbr file %s', path)
        data = None
    return data

def cover_cbz(path):
    data = None
    logger.debug('cover %s', path)
    data = None
    try:
        z = zipfile.ZipFile(path)
    except zipfile.BadZipFile:
        logger.debug('invalid cbz file %s', path)
        return data
    files = [f.filename for f in z.filelist]
    files = filter_images(files)
    if files:
        cover = ox.sorted_strings(files)[0]
        try:
            data = z.read(cover)
        except:
            data = None
    return data

def get_pages(path):
    files = []
    format = detect_format(path)
    if format == 'cbz':
        try:
            z = zipfile.ZipFile(path)
        except zipfile.BadZipFile:
            logger.debug('invalid cbz file %s', path)
            return data
        files = [f.filename for f in z.filelist]
    elif format == 'cbr':
        try:
            from unrar import rarfile
            rar = rarfile.RarFile(path)
            files = rar.namelist()
        except:
            pass
    files = filter_images(files)
    return len(files)


def info(path):
    data = {}
    data['title'] = os.path.splitext(os.path.basename(path))[0]
    data['pages'] = get_pages(path)
    return data