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