Tests updated for blacklist config file

vumi_http_proxy/proxy_blacklist.yaml → proxy_blacklist.yml

@@ -2,8 +2,7 @@
 #   Config file for vumi_http_proxy.
 #   Add all ip addresses to blacklist here
 #----------------------------------------------
-
-PROXY_BLACKLIST:
+proxy-blacklist:
     # IPv4 ADDRESSES
     - '69.16.230.117'
     - '93.184.216.34'

vumi_http_proxy/clickme.py

@@ -8,7 +8,8 @@
 @click.option('--interface', default='0.0.0.0', help='eg 0.0.0.0')
 @click.option('--port', default=8080, help='eg 80')
 @click.option(
-    '--blacklist', default='proxy_blacklist.yml', help='proxy_blacklist.yml')
+    '--blacklist', default='proxy_blacklist.yml',
+    help='proxy_blacklist.yml')
 def cli(interface, port, blacklist):
     """This script runs vumi-http-proxy on <interface>:<port>
     with the specified blacklist"""

vumi_http_proxy/http_proxy.py

@@ -1,5 +1,3 @@
-#!usr/bin/env python
-
 from twisted.python import log
 from twisted.web import http, proxy
 from twisted.internet import reactor
@@ -8,7 +6,6 @@
 from twisted.web.client import Agent, readBody
 from urlparse import urlparse, urlunparse
 from twisted.internet.defer import inlineCallbacks, succeed
-
 import yaml
 # blacklist of disallowed domains (move to proxy_blacklist.py)
 
@@ -106,13 +103,8 @@ def __init__(self, blacklist, resolver, http_client):
 
 
 class Initialize(object):
-
-    def __init__(self, blacklist, ip, port):
-        if not blacklist:
-            blacklist = DEFAULT_BLACKLIST
-        else:
-            with open("proxy_blacklist.yml", 'r') as blstream:
-                print yaml.load(blstream)
+    def __init__(self, blacklistfile, ip, port):
+        blacklist = self.read_file(blacklistfile)
         self.blacklist = blacklist
         self.ip = ip
         self.port = port
@@ -125,3 +117,13 @@ def main(self):
             reactor, "tcp:%d:interface=%s" % (self.port, self.ip))
         endpoint.listen(factory)
         reactor.run()
+
+    def read_file(self, blacklistfile):
+        blacklist = []
+        if not blacklistfile:
+            log.err("No blacklist config file provided.")
+        else:
+            with open(str(blacklistfile), 'r') as blstream:
+                bufferlist = yaml.load(blstream)
+                blacklist = bufferlist.get('proxy-blacklist')
+        return blacklist

vumi_http_proxy/proxy_blacklist.yml

@@ -1,13 +1,11 @@
-"""Config file for vumi_http_proxy.
-    Add all ip addresses to blacklist here"""
-
-PROXY_BLACKLIST = [
-
+#----------------------------------------------
+#   Config file for vumi_http_proxy.
+#   Add all ip addresses to blacklist here
+#----------------------------------------------
+proxy-blacklist:
     # IPv4 ADDRESSES
-    '69.16.230.117',
-    '93.184.216.34',
+    - '69.16.230.117'
+    - '93.184.216.34'
 
     # IPv6 ADDRESSES
-    '2606:2800:220:1:248:1893:25c8:1946'
-
-]
+    - '2606:2800:220:1:248:1893:25c8:1946'

vumi_http_proxy/servicemaker.py

@@ -15,7 +15,9 @@
 class Options(usage.Options):
     optParameters = [["port", None, 8080,
                      "The port number to start the proxy"],
-                     ["interface", None, "0.0.0.0", "IP to start proxy on"]]
+                     ["interface", None, "0.0.0.0", "IP to start proxy on"],
+                     ["blacklist", None, "proxy_blacklist.yml",
+                     "Name of the YAML config file for blacklist"]]
 
 
 class ProxyWorkerServiceMaker(object):
@@ -26,6 +28,6 @@ class ProxyWorkerServiceMaker(object):
 
     def makeService(self, options):
         factory = ProxyFactory(
-            ["zombo.com"], client.createResolver(), Agent(reactor))
+            options["blacklist"], client.createResolver(), Agent(reactor))
         return strports.service("tcp:%d:interface=%s" % (
                 options["port"], str(options["interface"])), factory)

vumi_http_proxy/test/test_clickme.py

@@ -12,12 +12,25 @@ def test_click(self):
         self.initializers = []
         self.patch(http_proxy.Initialize, 'main',
                    lambda x: self.initializers.append(x))
-        result = runner.invoke(clickme.cli)
-        [initializer] = self.initializers
-        self.assertEqual(result.exit_code, 0)
+        fake_blacklist = ["foo", "bar"]
+        filename = self.make_blacklist(fake_blacklist)
+        result = runner.invoke(clickme.cli, ['--blacklist', filename])
+        self.assertEqual(result.exception, None)
         self.assertEqual(result.output.splitlines(), [
            'Starting connection to 0.0.0.0:8080',
         ])
+        self.assertEqual(result.exit_code, 0)
+
+        [initializer] = self.initializers
         self.assertEquals(type(initializer), http_proxy.Initialize)
         self.assertEquals(initializer.port, 8080)
         self.assertEquals(initializer.ip, "0.0.0.0")
+        self.assertEquals(initializer.blacklist, ["foo", "bar"])
+
+    def make_blacklist(self, blacklist):
+        filename = self.mktemp()
+        with open(filename, 'w') as stream:
+            stream.write("proxy-blacklist:\n")
+            for i in blacklist:
+                stream.write(" - "+i+"\n")
+        return filename

vumi_http_proxy/test/test_servicemaker.py

@@ -11,13 +11,16 @@ def test_defaults(self):
         options.parseOptions([])
         self.assertEqual(options["port"], 8080)
         self.assertEqual(str(options["interface"]), "0.0.0.0")
+        self.assertEqual(str(options["blacklist"]), "proxy_blacklist.yml")
 
     def test_override(self):
         options = Options()
         options.parseOptions(["--port", 8000])
         options.parseOptions(["--interface", '127.0.0.1'])
+        options.parseOptions(["--blacklist", "fake_blacklist.yml"])
         self.assertEqual(options["port"], "8000")
         self.assertEqual(str(options["interface"]), "127.0.0.1")
+        self.assertEqual(str(options["blacklist"]), "fake_blacklist.yml")
 
 
 class TestProxyWorkerServiceMaker(unittest.TestCase):