HTTP/SOCKS proxy support added (resolves DTVD/rainbowstream#60)

diff --git a/README.rst b/README.rst
index 7727f720c57ed94c7aaa6fb1813c4ce60e6d69f7..b1f6fdb55019a4ced4fb3a15ccd8495c505ef755 100644 (file)
--- a/README.rst
+++ b/README.rst
@@ -105,6 +105,18 @@ In the first time you will be asked for authorization of Rainbow Stream
 app at Twitter. Just click the “Authorize access” button and paste PIN\r
 number to the terminal, the rainbow will start.\r
 \r
+You might want to use rainbowstream via an **HTTP/SOCKS proxy**. Proxy settings are\r
+provided as follows:\r
+\r
+.. code:: bash\r
+\r
+    rainbowstream --proxy-host localhost --proxy-port 1337 --proxy-type HTTP\r
+    # or using the short form:\r
+    rainbowstream -ph localhost -pp 1337 -pt HTTP\r
+\r
+Both ``--proxy-port`` and ``--proxy-type`` can be omitted. In this case default\r
+proxy port ``8080`` and default proxy type ``SOCKS5`` are used.\r
+\r
 The interactive mode\r
 ^^^^^^^^^^^^^^^^^^^^\r
 \r

diff --git a/docs/index.rst b/docs/index.rst
index b9dcdb42c5f635034b6121e18285e5266b3fbdb5..cb7da71d7d9239c1606145f2c57f9f687b099503 100644 (file)
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -97,6 +97,18 @@ In the first time you will be asked for authorization of Rainbow Stream
 app at Twitter. Just click the “Authorize access” button and paste PIN\r
 number to the terminal, the rainbow will start.\r
 \r
+You might want to use rainbowstream via an **HTTP/SOCKS proxy**. Proxy settings are\r
+provided as follows:\r
+\r
+.. code:: bash\r
+\r
+    rainbowstream --proxy-host localhost --proxy-port 1337 --proxy-type HTTP\r
+    # or using the short form:\r
+    rainbowstream -ph localhost -pp 1337 -pt HTTP\r
+\r
+Both ``--proxy-port`` and ``--proxy-type`` can be omitted. In this case default\r
+proxy port ``8080`` and default proxy type ``SOCKS5`` are used.\r
+\r
 The interactive mode\r
 ^^^^^^^^^^^^^^^^^^^^\r
 \r

diff --git a/rainbowstream/rainbow.py b/rainbowstream/rainbow.py
index 427830aa020d60f56b3d8980b9505cd243bf9fe8..c05183e74c6af1ee2529514bae48e27f85b9571d 100644 (file)
--- a/rainbowstream/rainbow.py
+++ b/rainbowstream/rainbow.py
@@ -9,6 +9,8 @@ import requests
 import webbrowser
 import traceback
 import pkg_resources
+import socks
+import socket
 
 from twitter.stream import TwitterStream, Timeout, HeartbeatTimeout, Hangup
 from twitter.api import *
@@ -16,6 +18,8 @@ from twitter.oauth import OAuth, read_token_file
 from twitter.oauth_dance import oauth_dance
 from twitter.util import printNicely
 
+from urllib import error
+
 from .draw import *
 from .colors import *
 from .config import *
@@ -58,6 +62,20 @@ def parse_arguments():
         '--image-on-term',
         action='store_true',
         help='Display all image on terminal.')
+    parser.add_argument(
+        '-ph',
+        '--proxy-host',
+        help='Use HTTP/SOCKS proxy for network connections.')
+    parser.add_argument(
+        '-pp',
+        '--proxy-port',
+        default=8080,
+        help='HTTP/SOCKS proxy port (Default: 8080).')
+    parser.add_argument(
+        '-pt',
+        '--proxy-type',
+        default='SOCKS4',
+        help='Proxy type (HTTP, SOCKS4, SOCKS5; Default: SOCKS5).')
     return parser.parse_args()
 
 
@@ -2024,6 +2042,20 @@ def fly():
     # Initial
     args = parse_arguments()
     try:
+        if args.proxy_host:
+            # Setup proxy by monkeypatching the standard lib
+            # You might want to check https://github.com/Anorov/PySocks for further
+            # further info.
+            if args.proxy_type.lower() == "socks5" or not args.proxy_type:
+                socks.set_default_proxy(socks.SOCKS5, args.proxy_host, int(args.proxy_port))
+            elif args.proxy_type.lower() == "http":
+                socks.set_default_proxy(socks.HTTP, args.proxy_host, int(args.proxy_port))
+            elif args.proxy_type.lower() == "socks4":
+                socks.set_default_proxy(socks.SOCKS4, args.proxy_host, int(args.proxy_port))
+            else:
+                printNicely(magenta("Sorry, wrong proxy type specified! Aborting..."))
+                sys.exit()
+            socket.socket = socks.socksocket
         init(args)
     except TwitterHTTPError:
         printNicely('')
@@ -2032,6 +2064,14 @@ def fly():
         printNicely(magenta("Let's try again later."))
         save_history()
         sys.exit()
+    except (ConnectionRefusedError, socks.ProxyConnectionError, error.URLError):
+        printNicely(
+            magenta("There seems to be a connection problem."))
+        printNicely(
+            magenta("You might want to check your proxy settings (host, port and type)!"))
+        save_history()
+        sys.exit()
+
     # Spawn stream thread
     th = threading.Thread(
         target=stream,

diff --git a/setup.py b/setup.py
index fae3c8c88dc16e3cab72168a6e9a9ea995277591..472666ad37552768df358b949cc15d6a8de53b4e 100644 (file)
--- a/setup.py
+++ b/setup.py
@@ -13,6 +13,7 @@ install_requires = [
     "pyfiglet",
     "twitter",
     "Pillow",
+    "PySocks"
 ]
 
 # Copy default config if not exists