From: Ricter Z <ricterzheng@gmail.com>
Date: Thu, 4 Feb 2016 02:25:35 +0000 (+0800)
Subject: feature: proxy support
X-Git-Tag: 0.1.4~5
X-Git-Url: https://git.lizzy.rs/?a=commitdiff_plain;h=4a0f09b052776fb04563c89d6b57f2b6b6cf19c8;p=nhentai.git

feature: proxy support
---

diff --git a/.gitignore b/.gitignore
index a56afbc..d8339d7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,5 @@
 .idea/
 build
 dist/
+*.egg-info
 
diff --git a/nhentai/cmdline.py b/nhentai/cmdline.py
index 891d1ab..0ab8830 100644
--- a/nhentai/cmdline.py
+++ b/nhentai/cmdline.py
@@ -8,6 +8,9 @@ except ImportError:
     pass
 
 
+import constant
+
+
 def banner():
     print('''       _   _            _        _
  _ __ | | | | ___ _ __ | |_ __ _(_)
@@ -31,6 +34,8 @@ def cmd_parser():
                       help='thread count of download doujinshi')
     parser.add_option('--timeout', type='int', dest='timeout', action='store', default=30,
                       help='timeout of download doujinshi')
+    parser.add_option('--proxy', type='string', dest='proxy', action='store', default='',
+                      help='use proxy, example: socks5://127.0.0.1:1080')
     args, _ = parser.parse_args()
 
     if args.ids:
@@ -55,4 +60,12 @@ def cmd_parser():
         logger.critical('Maximum number of used threads is 10')
         raise SystemExit
 
+    if args.proxy:
+        import urlparse
+        proxy_url = urlparse.urlparse(args.proxy)
+        if proxy_url.scheme not in ('http', 'https'):
+            logger.error('Invalid protocol \'{}\' of proxy, ignored'.format(proxy_url.scheme))
+        else:
+            constant.PROXY = {proxy_url.scheme: args.proxy}
+
     return args
diff --git a/nhentai/command.py b/nhentai/command.py
index 647952c..e22deff 100644
--- a/nhentai/command.py
+++ b/nhentai/command.py
@@ -48,6 +48,8 @@ def signal_handler(signal, frame):
     logger.error('Ctrl-C signal received. Quit.')
     raise SystemExit
 
+
 signal.signal(signal.SIGINT, signal_handler)
+
 if __name__ == '__main__':
     main()
diff --git a/nhentai/constant.py b/nhentai/constant.py
index 5196bd7..5439384 100644
--- a/nhentai/constant.py
+++ b/nhentai/constant.py
@@ -3,3 +3,4 @@ URL = '%snhentai.net' % SCHEMA
 DETAIL_URL = '%s/g' % URL
 IMAGE_URL = '%si.nhentai.net/galleries' % SCHEMA
 SEARCH_URL = '%s/search/' % URL
+PROXY = {}
diff --git a/nhentai/downloader.py b/nhentai/downloader.py
index 09c1a5c..402a7a8 100644
--- a/nhentai/downloader.py
+++ b/nhentai/downloader.py
@@ -4,6 +4,7 @@ import requests
 import threadpool
 from urlparse import urlparse
 from logger import logger
+from parser import request
 
 
 class Downloader(object):
@@ -27,7 +28,7 @@ class Downloader(object):
         filename = filename if filename else os.path.basename(urlparse(url).path)
         try:
             with open(os.path.join(folder, filename), "wb") as f:
-                response = requests.get(url, stream=True, timeout=self.timeout)
+                response = request('get', url, stream=True, timeout=self.timeout)
                 length = response.headers.get('content-length')
                 if length is None:
                     f.write(response.content)
diff --git a/nhentai/parser.py b/nhentai/parser.py
index 7585a29..01786e9 100644
--- a/nhentai/parser.py
+++ b/nhentai/parser.py
@@ -4,11 +4,18 @@ import sys
 import re
 import requests
 from bs4 import BeautifulSoup
-from constant import DETAIL_URL, SEARCH_URL
+import constant
 from logger import logger
 from tabulate import tabulate
 
 
+def request(method, url, **kwargs):
+    if not hasattr(requests, method):
+        raise AttributeError('\'requests\' object has no attribute \'{}\''.format(method))
+
+    return requests.__dict__[method](url, proxies=constant.PROXY, **kwargs)
+
+
 def doujinshi_parser(id):
     if not isinstance(id, (int, )) and (isinstance(id, (str, )) and not id.isdigit()):
         raise Exception('Doujinshi id(%s) is not valid' % str(id))
@@ -16,10 +23,10 @@ def doujinshi_parser(id):
     logger.debug('Fetching doujinshi information of id %d' % id)
     doujinshi = dict()
     doujinshi['id'] = id
-    url = '%s/%d/' % (DETAIL_URL, id)
+    url = '%s/%d/' % (constant.DETAIL_URL, id)
 
     try:
-        response = requests.get(url).content
+        response = request('get', url).content
     except Exception as e:
         logger.critical('%s%s' % tuple(e.message))
         sys.exit()
@@ -54,7 +61,13 @@ def doujinshi_parser(id):
 def search_parser(keyword, page):
     logger.debug('Searching doujinshis of keyword %s' % keyword)
     result = []
-    response = requests.get(SEARCH_URL, params={'q': keyword, 'page': page}).content
+    try:
+        response = request('get', url=constant.SEARCH_URL, params={'q': keyword, 'page': page}).content
+    except requests.ConnectionError as e:
+        logger.critical(e)
+        logger.warn('If you are in China, please configure the proxy to fu*k GFW.')
+        raise SystemExit
+
     html = BeautifulSoup(response)
     doujinshi_search_result = html.find_all('div', attrs={'class': 'gallery'})
     for doujinshi in doujinshi_search_result: