Headers: Init the feature.

cabu/core.py

@@ -64,7 +64,7 @@ def __init__(self, import_name=__package__, db=None, *args, **kwargs):
                 self.config['S3_SECRET_KEY']
             )
 
-        if self.config.get('FTP_HOST'):
+        if self.config.get('FTP_HOST'):  # pragma: no cover
             self.ftp = ftpretty(
                 self.config['FTP_HOST'],
                 self.config['FTP_LOGIN'],
@@ -75,11 +75,11 @@ def __init__(self, import_name=__package__, db=None, *args, **kwargs):
         self.webdriver = load_driver(self.config, self.vdisplay)
 
     def __del__(self):
-        if hasattr(self, 'ftp'):
+        if hasattr(self, 'ftp') and self.ftp:  # pragma: no cover
             self.ftp.close()
-        if hasattr(self, 'webdriver'):
+        if hasattr(self, 'webdriver') and self.webdriver:
             unload_driver(self.webdriver)
-        if hasattr(self, 'vdisplay'):
+        if hasattr(self, 'vdisplay') and self.vdisplay:
             unload_vdisplay(self.vdisplay)
 
     def __exit__(self, *args, **kwargs):  # pragma: no cover
@@ -124,7 +124,10 @@ def load_config(self, settings='cabu.default_settings'):
         # settings.py are overriden by Env vars.
         for key in self.config.keys():
             if key.isupper() and key in os.environ:
-                if os.environ[key] == 'True':
+                if os.environ[key][:1] == '{':
+                    import ast
+                    self.config[key] = ast.literal_eval(os.environ[key])
+                elif os.environ[key] == 'True':
                     self.config[key] = True
                 elif os.environ[key] == 'False':
                     self.config[key] = False

cabu/default_settings.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-DEBUG = True
+DEBUG = False
 CABU_TEST = None
 LOGGER_LEVEL = 'error'
 LOGGER_PATH = ''
@@ -9,17 +9,18 @@
 ADMIN_USERNAME = 'admin'
 ADMIN_PASSWORD = 'admin'
 
-DRIVER_NAME = 'PhantomJS'
+DRIVER_NAME = None
 DRIVER_BINARY_PATH = None
 # DRIVER_BINARY_PATH = '/Users/t.leveque/Applications/Firefox.app/Contents/MacOS/firefox'
 DRIVER_WINDOWS_WIDTH = 1024
 DRIVER_WINDOWS_HEIGHT = 768
 DRIVER_PAGE_TIMEOUT = 30  # In seconds.
 
-FTP_HOST = 'ftp'
+FTP_HOST = None
 FTP_LOGIN = None
 FTP_PASSWORD = None
 
 HEADLESS = True  # Run the Browser in Xvfd
+HEADERS = None
 
 IO_CONCURRENCY = False  # If True, Python version should be superior or equal to py34

cabu/drivers.py

@@ -5,6 +5,10 @@
 from selenium import webdriver
 from xvfbwrapper import Xvfb
 from cabu.exceptions import DriverException
+from cabu.utils.headers import Headers
+from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
+from selenium import webdriver
+
 try:
     from urllib.parse import urlsplit
 except ImportError:  # pragma: no cover
@@ -33,6 +37,15 @@ def load_vdisplay(config):
     return vdisplay
 
 
+def unload_vdisplay(vdisplay):
+    """Shutdown given Xvfb instance.
+
+    Args:
+        vdisplay (XvfbWrapper): The running virtual X server.
+    """
+    vdisplay.stop()
+
+
 def load_driver(config, vdisplay=None):
     """Initialize a weddriver selected in config with given config.
 
@@ -69,15 +82,6 @@ def unload_driver(driver):
     driver.quit()
 
 
-def unload_vdisplay(vdisplay):
-    """Shutdown given Xvfb instance.
-
-    Args:
-        vdisplay (XvfbWrapper): The running virtual X server.
-    """
-    vdisplay.stop()
-
-
 def load_firefox(config):
     """Start Firefox webdriver with the given configuration.
 
@@ -89,13 +93,12 @@ def load_firefox(config):
 
     """
     binary = None
-    profile = None
+    profile = webdriver.FirefoxProfile()
 
     if os.environ.get('HTTPS_PROXY') or os.environ.get('HTTP_PROXY'):
         proxy_address = os.environ.get('HTTPS_PROXY', os.environ.get('HTTP_PROXY'))
         proxy_port = re.search('\:([0-9]+)$', proxy_address).group(1)
 
-        profile = webdriver.FirefoxProfile()
         profile.set_preference('network.proxy.type', 1)
         profile.set_preference(
             'network.proxy.http',
@@ -104,6 +107,9 @@ def load_firefox(config):
         profile.set_preference('network.proxy.http_port', proxy_port)
         profile.update_preferences()
 
+    if 'HEADERS' in config and config['HEADERS']:
+        profile = Headers(config).set_headers(profile)
+
     if config['DRIVER_BINARY_PATH']:
         from selenium.webdriver.firefox.firefox_binary import FirefoxBinary
         binary = FirefoxBinary(config['DRIVER_BINARY_PATH'])
@@ -134,18 +140,24 @@ def load_phantomjs(config):
         webdriver (selenium.webdriver): An instance of phantomJS webdriver.
 
     """
+    dcap = dict(DesiredCapabilities.PHANTOMJS)
     service_args = [
         '--ignore-ssl-errors=true',
         '--ssl-protocol=any',
         '--web-security=false'
     ]
+
     if os.environ.get('HTTPS_PROXY') or os.environ.get('HTTP_PROXY'):
         proxy_address = os.environ.get('HTTPS_PROXY', os.environ.get('HTTP_PROXY'))
         proxy_ip = re.search('http\:\/\/(.*)$', proxy_address).group(1)
         service_args.append('--proxy=%s' % proxy_ip)
         service_args.append('--proxy-type=http')
 
+    if 'HEADERS' in config and config['HEADERS']:
+        dcap = Headers(config).set_headers(dcap)
+
     return webdriver.PhantomJS(
+        desired_capabilities=dcap,
         service_args=service_args,
         service_log_path=os.path.devnull
     )

cabu/exceptions.py

@@ -27,5 +27,9 @@ class CookieStorageException(Exception):
     pass
 
 
+class HeaderException(Exception):
+    pass
+
+
 class DriverException(Exception):
     pass

cabu/tests/fixtures/custom_settings.py

@@ -17,6 +17,11 @@
 FTP_PASSWORD = None
 
 HEADLESS = True  # Run the Browser in Xvfd
+HEADERS = {
+    'User-Agent':
+        'Mozilla/6.0 (Macintosh; Intel Mac OS X 10_11_3) Apple'
+        'WebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.103 Safari/537.36'
+}
 
 IO_CONCURRENCY = False  # If True, Python version should be superior or equal to py34
 

cabu/tests/test_core.py

@@ -18,12 +18,14 @@ def setUp(self):
             del os.environ['CABU_SETTINGS']
 
     def tearDown(self):
+        if 'HEADERS' in os.environ:
+            del os.environ['HEADERS']
         self.patcher_load_phantomjs.stop()
 
     def test_cabu_default(self):
         cabu = Cabu(__name__)
         self.assertIsInstance(cabu, Cabu)
-        self.assertEquals(cabu.config['DRIVER_NAME'], 'PhantomJS')
+        self.assertEquals(cabu.config['DRIVER_NAME'], None)
         self.assertEquals(cabu.config['DRIVER_WINDOWS_WIDTH'], 1024)
         self.assertEquals(cabu.config['DRIVER_WINDOWS_HEIGHT'], 768)
         self.assertEquals(cabu.config['DRIVER_PAGE_TIMEOUT'], 30)
@@ -54,6 +56,20 @@ def test_cabu_custom(self):
         self.assertEquals(cabu.config['HEADLESS'], True)
         del cabu
 
+    def test_cabu_custom_with_dict_in_environ(self):
+        os.environ['CABU_SETTINGS'] = 'cabu.tests.fixtures.custom_settings'
+        os.environ['HEADERS'] = "{'abc': 1}"
+        cabu = Cabu(__name__)
+        self.assertIsInstance(cabu, Cabu)
+        self.assertEquals(cabu.config['DRIVER_NAME'], 'PhantomJS')
+        self.assertEquals(cabu.config['DRIVER_WINDOWS_WIDTH'], 800)
+        self.assertEquals(cabu.config['DRIVER_WINDOWS_HEIGHT'], 600)
+        self.assertEquals(cabu.config['DRIVER_PAGE_TIMEOUT'], 10)
+        self.assertEquals(cabu.config['IO_CONCURRENCY'], False)
+        self.assertEquals(cabu.config['HEADLESS'], True)
+        self.assertEquals(cabu.config['HEADERS'], {'abc': 1})
+        del cabu
+
     def test_cabu_load_with_db(self):
         cabu = Cabu(__name__, db=PyMongo)
         self.assertIsInstance(cabu.db, Database)

cabu/tests/test_drivers.py

@@ -36,9 +36,6 @@ def test_unload_driver_fail(self, mock_driver):
     def test_load_firefox(self, mock_driver):
         self.config['DRIVER_NAME'] = 'Firefox'
         driver = load_driver(self.config)
-        mock_driver.assert_called_once_with(
-            firefox_binary=None, firefox_profile=None
-        )
         driver.close()
 
     @patch('cabu.drivers.webdriver.Firefox', return_value=MagicMock())
@@ -48,6 +45,18 @@ def test_load_firefox_with_binary(self, mock_driver):
         driver = load_driver(self.config)
         driver.close()
 
+    @patch('cabu.drivers.webdriver.Firefox', return_value=MagicMock())
+    def test_load_firefox_with_headers(self, mock_driver):
+        self.config['DRIVER_NAME'] = 'Firefox'
+        self.config['HEADERS'] = {
+            'User-Agent':
+                'Mozilla/7.0 (Macintosh; Intel Mac OS X 10_11_3) Apple'
+                'WebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.10'
+                '3 Safari/537.36'
+        }
+        driver = load_driver(self.config)
+        driver.close()
+
     @patch('cabu.drivers.webdriver.Chrome', return_value=MagicMock())
     def test_load_chrome(self, mock_driver):
         self.config['DRIVER_NAME'] = 'Chrome'
@@ -56,10 +65,12 @@ def test_load_chrome(self, mock_driver):
         driver.close()
 
     @patch('cabu.drivers.webdriver.PhantomJS', return_value=MagicMock())
-    def test_load_phantomjs(self, mock_driver):
+    @patch('cabu.drivers.DesiredCapabilities', return_value='mock_dcap')
+    def test_load_phantomjs(self, mock_dcap, mock_driver):
         self.config['DRIVER_NAME'] = 'PhantomJS'
         driver = load_driver(self.config)
         mock_driver.assert_called_once_with(
+            desired_capabilities={},
             service_args=[
                 '--ignore-ssl-errors=true',
                 '--ssl-protocol=any',
@@ -69,6 +80,19 @@ def test_load_phantomjs(self, mock_driver):
         )
         driver.close()
 
+    @patch('cabu.drivers.webdriver.PhantomJS', return_value=MagicMock())
+    @patch('cabu.drivers.DesiredCapabilities', return_value='mock_dcap')
+    def test_load_phantomjs_with_headers(self, mock_dcap, mock_driver):
+        self.config['DRIVER_NAME'] = 'PhantomJS'
+        self.config['HEADERS'] = {
+            'User-Agent':
+                'Mozilla/7.0 (Macintosh; Intel Mac OS X 10_11_3) Apple'
+                'WebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.10'
+                '3 Safari/537.36'
+        }
+        driver = load_driver(self.config)
+        driver.close()
+
     @patch('cabu.drivers.webdriver.Firefox', return_value=MagicMock())
     def test_load_with_firefox_http_proxy(self, mock_driver):
         self.config['DRIVER_NAME'] = 'Firefox'
@@ -84,16 +108,30 @@ def test_load_with_firefox_https_proxy(self, mock_driver):
         driver.close()
 
     @patch('cabu.drivers.webdriver.PhantomJS', return_value=MagicMock())
-    def test_load_with_phantomjs_http_proxy(self, mock_driver):
+    @patch('cabu.drivers.DesiredCapabilities', return_value='mock_dcap')
+    def test_load_with_phantomjs_http_proxy(self, mock_dcap, mock_driver):
         os.environ['HTTP_PROXY'] = 'http://127.0.0.1:80'
         driver = load_driver(self.config)
+        mock_driver.assert_called_once_with(
+            desired_capabilities={},
+            service_args=[
+                '--ignore-ssl-errors=true',
+                '--ssl-protocol=any',
+                '--web-security=false',
+                '--proxy=127.0.0.1:80',
+                '--proxy-type=http'
+            ],
+            service_log_path='/dev/null'
+        )
         driver.close()
 
     @patch('cabu.drivers.webdriver.PhantomJS', return_value=MagicMock())
-    def test_load_with_phantomjs_https_proxy(self, mock_driver):
+    @patch('cabu.drivers.DesiredCapabilities', return_value='mock_dcap')
+    def test_load_with_phantomjs_https_proxy(self, mock_dcap, mock_driver):
         os.environ['HTTP_PROXY'] = 'http://127.0.0.1:80'
         driver = load_driver(self.config)
         mock_driver.assert_called_once_with(
+            desired_capabilities={},
             service_args=[
                 '--ignore-ssl-errors=true',
                 '--ssl-protocol=any',

cabu/tests/test_headers.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+
+from cabu.tests.test_base import TestBase
+from cabu.utils.headers import Headers
+from mock import patch
+
+from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
+from selenium import webdriver
+
+
+class TestHeaders(TestBase):
+    def setUp(self):
+        super(TestHeaders, self).setUp()
+        self.patcher_phantomjs = patch('cabu.drivers.load_phantomjs', spec=True)
+        self.patcher_phantomjs.start()
+
+        self.config = self.app.config
+        self.config['HEADERS'] = {
+            'User-Agent':
+                'Mozilla/6.0 (Macintosh; Intel Mac OS X 10_11_3) Apple'
+                'WebKit/537.36 (KHTML, like Gecko) Chrome/48.0.2564.103 Safari/537.36'
+        }
+
+    def tearDown(self):
+        super(TestHeaders, self).setUp()
+        self.patcher_phantomjs.stop()
+
+    def test_firefox_headers_loading(self):
+        self.app.config['DRIVER_NAME'] = 'Firefox'
+        profile = webdriver.FirefoxProfile()
+        headers = Headers(self.app.config)
+        profile = headers.set_headers(profile)
+        self.assertEquals(
+            profile.__dict__['default_preferences']['general.useragent.override'],
+            'Mozilla/6.0 (Macintosh; Intel Mac OS X 10_11_3) AppleWebKit/537.36'
+            ' (KHTML, like Gecko) Chrome/48.0.2564.103 Safari/537.36'
+        )
+
+    def test_phantomjs_headers_loading(self):
+        dcap = dict(DesiredCapabilities.PHANTOMJS)
+        headers = Headers(self.config).set_headers(dcap)
+        self.assertEquals(
+            headers['phantomjs.page.customHeaders.User-Agent'],
+            'Mozilla/6.0 (Macintosh; Intel Mac OS X 10_11_3) AppleWebKit/537.36'
+            ' (KHTML, like Gecko) Chrome/48.0.2564.103 Safari/537.36'
+        )
+
+    def test_chrome_headers_loading(self):
+        self.app.config['DRIVER_NAME'] = 'Chrome'
+        dcap = dict(DesiredCapabilities.PHANTOMJS)
+        with self.assertRaises(Exception):
+            Headers(self.config).set_headers(dcap)

cabu/tests/test_settings.py

@@ -22,5 +22,6 @@
 
 SERVER_X = False
 HEADLESS = True  # Run the Browser in Xvfd
+HEADERS = None
 
 IO_CONCURRENCY = False  # If True, Python version should be superior or equal to py34