From 9ef5c012350e09685da61dc986dfc29d28989533 Mon Sep 17 00:00:00 2001
From: j <j@mailb.org>
Date: Fri, 1 Feb 2019 19:19:16 +0530
Subject: [PATCH] merge cbr/cbz cover extraction

---
 oml/media/cbr.py | 47 ++++++++++++++++++++++++++++++++++++++++++++---
 oml/media/cbz.py | 35 +----------------------------------
 2 files changed, 45 insertions(+), 37 deletions(-)

diff --git a/oml/media/cbr.py b/oml/media/cbr.py
index a1840a3..244ee67 100644
--- a/oml/media/cbr.py
+++ b/oml/media/cbr.py
@@ -1,19 +1,39 @@
 # -*- 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] in IMAGE_EXTENSION]
+    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
-    #open rar file and extract first page here
     try:
         from unrar import rarfile
         rar = rarfile.RarFile(path)
@@ -23,9 +43,30 @@ def cover(path):
             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 info(path):
     data = {}
     data['title'] = os.path.splitext(os.path.basename(path))[0]
diff --git a/oml/media/cbz.py b/oml/media/cbz.py
index 0e71475..1428292 100644
--- a/oml/media/cbz.py
+++ b/oml/media/cbz.py
@@ -1,36 +1,3 @@
 # -*- coding: utf-8 -*-
-import logging
-import os
-import zipfile
-
-import ox
-
-from .cbr import filter_images
-
-logger = logging.getLogger(__name__)
-
-def cover(path):
-    data = None
-    logger.debug('cover %s', path)
-    data = None
-    try:
-        z = zipfile.ZipFile(path)
-    except zipfile.BadZipFile:
-        logger.debug('invalid zbc 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 info(path):
-    data = {}
-    data['title'] = os.path.splitext(os.path.basename(path))[0]
-    #data['pages'] = fixme read rar to count pages
-    return data
+from .cbr import cover, info