From c426a344303087c8c13391752473c6a22fda8e14 Mon Sep 17 00:00:00 2001
From: rc0r <hlt99@blinkenshell.org>
Date: Thu, 2 Oct 2014 23:07:08 +0200
Subject: [PATCH] HTTP/SOCKS proxy support added (resolves
 DTVD/rainbowstream#60)

---
 README.rst               | 12 ++++++++++++
 docs/index.rst           | 12 ++++++++++++
 rainbowstream/rainbow.py | 40 ++++++++++++++++++++++++++++++++++++++++
 setup.py                 |  1 +
 4 files changed, 65 insertions(+)

diff --git a/README.rst b/README.rst
index 7727f72..b1f6fdb 100644
--- 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
 number to the terminal, the rainbow will start.
 
+You might want to use rainbowstream via an **HTTP/SOCKS proxy**. Proxy settings are
+provided as follows:
+
+.. code:: bash
+
+    rainbowstream --proxy-host localhost --proxy-port 1337 --proxy-type HTTP
+    # or using the short form:
+    rainbowstream -ph localhost -pp 1337 -pt HTTP
+
+Both ``--proxy-port`` and ``--proxy-type`` can be omitted. In this case default
+proxy port ``8080`` and default proxy type ``SOCKS5`` are used.
+
 The interactive mode
 ^^^^^^^^^^^^^^^^^^^^
 
diff --git a/docs/index.rst b/docs/index.rst
index b9dcdb4..cb7da71 100644
--- 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
 number to the terminal, the rainbow will start.
 
+You might want to use rainbowstream via an **HTTP/SOCKS proxy**. Proxy settings are
+provided as follows:
+
+.. code:: bash
+
+    rainbowstream --proxy-host localhost --proxy-port 1337 --proxy-type HTTP
+    # or using the short form:
+    rainbowstream -ph localhost -pp 1337 -pt HTTP
+
+Both ``--proxy-port`` and ``--proxy-type`` can be omitted. In this case default
+proxy port ``8080`` and default proxy type ``SOCKS5`` are used.
+
 The interactive mode
 ^^^^^^^^^^^^^^^^^^^^
 
diff --git a/rainbowstream/rainbow.py b/rainbowstream/rainbow.py
index 427830a..c05183e 100644
--- 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 fae3c8c..472666a 100644
--- a/setup.py
+++ b/setup.py
@@ -13,6 +13,7 @@ install_requires = [
     "pyfiglet",
     "twitter",
     "Pillow",
+    "PySocks"
 ]
 
 # Copy default config if not exists
-- 
2.25.1