nhentai/utils.py

   1 # coding: utf-8
   2 from __future__ import unicode_literals, print_function
   3
   4 import sys
   5 import re
   6 import os
   7 import string
   8 import zipfile
   9 import shutil
  10 import requests
  11 import sqlite3
  12 import img2pdf
  13
  14 from nhentai import constant
  15 from nhentai.logger import logger
  16 from nhentai.serializer import serialize_json, serialize_comicxml, set_js_database
  17
  18
  19 def request(method, url, **kwargs):
  20     session = requests.Session()
  21     session.headers.update({
  22         'Referer': constant.LOGIN_URL,
  23         'User-Agent': 'nhentai command line client (https://github.com/RicterZ/nhentai)',
  24         'Cookie': constant.COOKIE
  25     })
  26     return getattr(session, method)(url, proxies=constant.PROXY, verify=False, **kwargs)
  27
  28
  29 def check_cookie():
  30     response = request('get', constant.BASE_URL).text
  31     username = re.findall('"/users/\d+/(.*?)"', response)
  32     if not username:
  33         logger.error('Cannot get your username, please check your cookie or use `nhentai --cookie` to set your cookie')
  34     else:
  35         logger.info('Login successfully! Your username: {}'.format(username[0]))
  36
  37
  38 class _Singleton(type):
  39     """ A metaclass that creates a Singleton base class when called. """
  40     _instances = {}
  41
  42     def __call__(cls, *args, **kwargs):
  43         if cls not in cls._instances:
  44             cls._instances[cls] = super(_Singleton, cls).__call__(*args, **kwargs)
  45         return cls._instances[cls]
  46
  47
  48 class Singleton(_Singleton(str('SingletonMeta'), (object,), {})):
  49     pass
  50
  51
  52 def urlparse(url):
  53     try:
  54         from urlparse import urlparse
  55     except ImportError:
  56         from urllib.parse import urlparse
  57
  58     return urlparse(url)
  59
  60
  61 def readfile(path):
  62     loc = os.path.dirname(__file__)
  63
  64     with open(os.path.join(loc, path), 'r') as file:
  65         return file.read()
  66
  67
  68 def generate_html(output_dir='.', doujinshi_obj=None):
  69     image_html = ''
  70
  71     if doujinshi_obj is not None:
  72         doujinshi_dir = os.path.join(output_dir, doujinshi_obj.filename)
  73     else:
  74         doujinshi_dir = '.'
  75
  76     file_list = os.listdir(doujinshi_dir)
  77     file_list.sort()
  78
  79     for image in file_list:
  80         if not os.path.splitext(image)[1] in ('.jpg', '.png'):
  81             continue
  82
  83         image_html += '<img src="{0}" class="image-item"/>\n'\
  84             .format(image)
  85     html = readfile('viewer/index.html')
  86     css = readfile('viewer/styles.css')
  87     js = readfile('viewer/scripts.js')
  88
  89     if doujinshi_obj is not None:
  90         serialize_json(doujinshi_obj, doujinshi_dir)
  91         name = doujinshi_obj.name
  92         if sys.version_info < (3, 0):
  93             name = doujinshi_obj.name.encode('utf-8')
  94     else:
  95         name = {'title': 'nHentai HTML Viewer'}
  96
  97     data = html.format(TITLE=name, IMAGES=image_html, SCRIPTS=js, STYLES=css)
  98     try:
  99         if sys.version_info < (3, 0):
 100             with open(os.path.join(doujinshi_dir, 'index.html'), 'w') as f:
 101                 f.write(data)
 102         else:
 103             with open(os.path.join(doujinshi_dir, 'index.html'), 'wb') as f:
 104                 f.write(data.encode('utf-8'))
 105
 106         logger.log(15, 'HTML Viewer has been written to \'{0}\''.format(os.path.join(doujinshi_dir, 'index.html')))
 107     except Exception as e:
 108         logger.warning('Writing HTML Viewer failed ({})'.format(str(e)))
 109
 110
 111 def generate_main_html(output_dir='./'):
 112     """
 113     Generate a main html to show all the contain doujinshi.
 114     With a link to their `index.html`.
 115     Default output folder will be the CLI path.
 116     """
 117
 118     image_html = ''
 119
 120     main = readfile('viewer/main.html')
 121     css = readfile('viewer/main.css')
 122     js = readfile('viewer/main.js')
 123
 124     element = '\n\
 125             <div class="gallery-favorite">\n\
 126                 <div class="gallery">\n\
 127                     <a href="./{FOLDER}/index.html" class="cover" style="padding:0 0 141.6% 0"><img\n\
 128                             src="./{FOLDER}/{IMAGE}" />\n\
 129                         <div class="caption">{TITLE}</div>\n\
 130                     </a>\n\
 131                 </div>\n\
 132             </div>\n'
 133
 134     os.chdir(output_dir)
 135     doujinshi_dirs = next(os.walk('.'))[1]
 136
 137     for folder in doujinshi_dirs:
 138         files = os.listdir(folder)
 139         files.sort()
 140
 141         if 'index.html' in files:
 142             logger.info('Add doujinshi \'{}\''.format(folder))
 143         else:
 144             continue
 145
 146         image = files[0]  # 001.jpg or 001.png
 147         if folder is not None:
 148             title = folder.replace('_', ' ')
 149         else:
 150             title = 'nHentai HTML Viewer'
 151
 152         image_html += element.format(FOLDER=folder, IMAGE=image, TITLE=title)
 153     if image_html == '':
 154         logger.warning('No index.html found, --gen-main paused.')
 155         return
 156     try:
 157         data = main.format(STYLES=css, SCRIPTS=js, PICTURE=image_html)
 158         if sys.version_info < (3, 0):
 159             with open('./main.html', 'w') as f:
 160                 f.write(data)
 161         else:
 162             with open('./main.html', 'wb') as f:
 163                 f.write(data.encode('utf-8'))
 164         shutil.copy(os.path.dirname(__file__)+'/viewer/logo.png', './')
 165         set_js_database()
 166         logger.log(
 167             15, 'Main Viewer has been written to \'{0}main.html\''.format(output_dir))
 168     except Exception as e:
 169         logger.warning('Writing Main Viewer failed ({})'.format(str(e)))
 170
 171
 172 def generate_cbz(output_dir='.', doujinshi_obj=None, rm_origin_dir=False, write_comic_info=False):
 173     if doujinshi_obj is not None:
 174         doujinshi_dir = os.path.join(output_dir, doujinshi_obj.filename)
 175         if write_comic_info:
 176             serialize_comicxml(doujinshi_obj, doujinshi_dir)
 177         cbz_filename = os.path.join(os.path.join(doujinshi_dir, '..'), '{}.cbz'.format(doujinshi_obj.filename))
 178     else:
 179         cbz_filename = './doujinshi.cbz'
 180         doujinshi_dir = '.'
 181
 182     file_list = os.listdir(doujinshi_dir)
 183     file_list.sort()
 184
 185     logger.info('Writing CBZ file to path: {}'.format(cbz_filename))
 186     with zipfile.ZipFile(cbz_filename, 'w') as cbz_pf:
 187         for image in file_list:
 188             image_path = os.path.join(doujinshi_dir, image)
 189             cbz_pf.write(image_path, image)
 190
 191     if rm_origin_dir:
 192         shutil.rmtree(doujinshi_dir, ignore_errors=True)
 193
 194     logger.log(15, 'Comic Book CBZ file has been written to \'{0}\''.format(doujinshi_dir))
 195
 196
 197 def generate_pdf(output_dir='.', doujinshi_obj=None, rm_origin_dir=False):
 198     """Write images to a PDF file using img2pdf."""
 199     if doujinshi_obj is not None:
 200         doujinshi_dir = os.path.join(output_dir, doujinshi_obj.filename)
 201         pdf_filename = os.path.join(
 202             os.path.join(doujinshi_dir, '..'),
 203             '{}.pdf'.format(doujinshi_obj.filename)
 204         )
 205     else:
 206         pdf_filename = './doujinshi.pdf'
 207         doujinshi_dir = '.'
 208
 209     file_list = os.listdir(doujinshi_dir)
 210     file_list.sort()
 211
 212     logger.info('Writing PDF file to path: {}'.format(pdf_filename))
 213     with open(pdf_filename, 'wb') as pdf_f:
 214         full_path_list = (
 215             [os.path.join(doujinshi_dir, image) for image in file_list]
 216         )
 217         pdf_f.write(img2pdf.convert(full_path_list))
 218
 219     if rm_origin_dir:
 220         shutil.rmtree(doujinshi_dir, ignore_errors=True)
 221
 222     logger.log(15, 'PDF file has been written to \'{0}\''.format(doujinshi_dir))
 223
 224
 225 def format_filename(s):
 226     """Take a string and return a valid filename constructed from the string.
 227 Uses a whitelist approach: any characters not present in valid_chars are
 228 removed. Also spaces are replaced with underscores.
 229
 230 Note: this method may produce invalid filenames such as ``, `.` or `..`
 231 When I use this method I prepend a date string like '2009_01_15_19_46_32_'
 232 and append a file extension like '.txt', so I avoid the potential of using
 233 an invalid filename.
 234
 235 """
 236     if sys.platform.startswith('win32'):
 237         invalid_chars = '\/:*?"<>|.'
 238         for char in invalid_chars:
 239             s = s.replace(char, '_')
 240     return s
 241
 242     # maybe you can use `--format` to select a suitable filename
 243     valid_chars = "-_.()[] %s%s" % (string.ascii_letters, string.digits)
 244     filename = ''.join(c for c in s if c in valid_chars)
 245     if len(filename) > 100:
 246         filename = filename[:100] + '...]'
 247
 248     # Remove [] from filename
 249     filename = filename.replace('[]', '').strip()
 250     return filename
 251
 252
 253 def signal_handler(signal, frame):
 254     logger.error('Ctrl-C signal received. Stopping...')
 255     exit(1)
 256
 257
 258 class DB(object):
 259     conn = None
 260     cur = None
 261
 262     def __enter__(self):
 263         self.conn = sqlite3.connect(constant.NHENTAI_HISTORY)
 264         self.cur = self.conn.cursor()
 265         self.cur.execute('CREATE TABLE IF NOT EXISTS download_history (id text)')
 266         self.conn.commit()
 267         return self
 268
 269     def __exit__(self, exc_type, exc_val, exc_tb):
 270         self.conn.close()
 271
 272     def clean_all(self):
 273         self.cur.execute('DELETE FROM download_history WHERE 1')
 274         self.conn.commit()
 275
 276     def add_one(self, data):
 277         self.cur.execute('INSERT INTO download_history VALUES (?)', [data])
 278         self.conn.commit()
 279
 280     def get_all(self):
 281         data = self.cur.execute('SELECT id FROM download_history')
 282         return [i[0] for i in data]