Merge pull request #11 from psiinon/master

Updated to supply apikey on all reqs and for 2.6.0

Author: kingthorin

setup.py

@@ -12,20 +12,25 @@
     print "You must have setuptools installed to use setup.py. Exiting..."
     raise SystemExit(1)
 
-test_requirements = [
+
+install_dependencies = (
+    'requests'
+)
+test_requirements = (
     'mock',
     'pylama',
-    'pytest'
-]
+    'pytest',
+    'requests_mock'
+)
 setup(
     name="python-owasp-zap-v2.4",
-    version="0.0.9.dev1",
-    description="OWASP ZAP 2.5 API client",
-    long_description="OWASP Zed Attack Proxy 2.5 API python client (the 2.4 package name has been kept to make it easier to upgrade)",
+    version="0.0.9",
+    description="OWASP ZAP 2.6 API client",
+    long_description="OWASP Zed Attack Proxy 2.6 API python client (the 2.4 package name has been kept to make it easier to upgrade)",
     author="ZAP development team",
     author_email='',
     url="https://www.owasp.org/index.php/OWASP_Zed_Attack_Proxy_Project",
-    download_url="https://github.com/zaproxy/zap-api-python/releases/tag/0.0.8",
+    download_url="https://github.com/zaproxy/zap-api-python/releases/tag/0.0.9",
     platforms=['any'],
 
     license="ASL2.0",
@@ -43,7 +48,7 @@
         'Intended Audience :: Developers',
         'Intended Audience :: Information Technology',
         'Programming Language :: Python'],
-
+    install_requires=install_dependencies,
     tests_require=test_requirements,
     extras_require={'tests': test_requirements}
 )

src/zapv2/__init__.py

@@ -21,8 +21,9 @@
 
 __docformat__ = 'restructuredtext'
 
-import json
-import urllib
+import requests
+from requests.packages.urllib3.exceptions import InsecureRequestWarning
+
 from acsrf import acsrf
 from ascan import ascan
 from ajaxSpider import ajaxSpider
@@ -58,8 +59,7 @@ class ZAPv2(object):
     # base OTHER api url
     base_other = 'http://zap/OTHER/'
 
-    def __init__(self, proxies={'http': 'http://127.0.0.1:8080',
-        'https': 'http://127.0.0.1:8080'}):
+    def __init__(self, proxies=None, apikey=None):
         """
         Creates an instance of the ZAP api client.
 
@@ -69,7 +69,11 @@ def __init__(self, proxies={'http': 'http://127.0.0.1:8080',
         Note that all of the other classes in this directory are generated
         new ones will need to be manually added to this file
         """
-        self.__proxies = proxies
+        self.__proxies = proxies or {
+            'http': 'http://127.0.0.1:8080',
+            'https': 'http://127.0.0.1:8080'
+        }
+        self.__apikey = apikey
 
         self.acsrf = acsrf(self)
         self.ajaxSpider = ajaxSpider(self)
@@ -95,6 +99,15 @@ def __init__(self, proxies={'http': 'http://127.0.0.1:8080',
         self.stats = stats(self)
         self.users = users(self)
 
+        # not very nice, but prevents warnings when accessing the ZAP API via https
+        requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
+
+        # Currently create a new session for each request to prevent request failing
+        # e.g. when polling the spider status
+        #self.session = requests.Session()
+        #if apikey is not None:
+        #  self.session.headers['X-ZAP-API-Key'] = apikey
+
     def urlopen(self, *args, **kwargs):
         """
         Opens a url forcing the proxies to be used.
@@ -103,25 +116,50 @@ def urlopen(self, *args, **kwargs):
            - `args`:  all non-keyword arguments.
            - `kwargs`: all other keyword arguments.
         """
-        kwargs['proxies'] = self.__proxies
-        return urllib.urlopen(*args, **kwargs).read()
+        # Must never leak the API key via proxied requests
+        return requests.get(*args, proxies=self.__proxies, verify=False, **kwargs).text
+
+    def _request_api(self, url, query=None):
+        """
+        Shortcut for an API request. Will always add the apikey (if defined)
+
+        :Parameters:
+           - `url`: the url to GET at.
+        """
+        if not url.startswith('http://zap/'):
+          # Only allow requests to the API so that we never leak the apikey
+          raise ValueError('A non ZAP API url was specified ' + url)
+          return;
+
+        # In theory we should be able to reuse the session,
+        # but there have been problems with that
+        self.session = requests.Session()
+        if self.__apikey is not None:
+          self.session.headers['X-ZAP-API-Key'] = self.__apikey
+
+        query = query or {}
+        if self.__apikey is not None:
+          # Add the apikey to get params for backwards compatibility
+          if not query.get('apikey'):
+            query['apikey'] = self.__apikey
+        return self.session.get(url, params=query, proxies=self.__proxies, verify=False)
 
     def _request(self, url, get=None):
         """
         Shortcut for a GET request.
 
         :Parameters:
            - `url`: the url to GET at.
-           - `get`: the disctionary to turn into GET variables.
+           - `get`: the dictionary to turn into GET variables.
         """
-        return json.loads(self.urlopen(url + '?' + urllib.urlencode(get or {})))
+        return self._request_api(url, get).json()
 
-    def _request_other(self, url, get={}):
+    def _request_other(self, url, get=None):
         """
         Shortcut for an API OTHER GET request.
 
         :Parameters:
            - `url`: the url to GET at.
-           - `get`: the disctionary to turn into GET variables.
+           - `get`: the dictionary to turn into GET variables.
         """
-        return self.urlopen(url + '?' + urllib.urlencode(get or {}))
+        return self._request_api(url, get).text

tests/unit/conftest.py

@@ -1,17 +1,19 @@
-from mock import patch
 import pytest
 
+import requests_mock
+
 from zapv2 import ZAPv2
 
+
 @pytest.yield_fixture
 def zap():
     """
     All tests will be able to share the instance of client with the same settings."""
-    yield ZAPv2()
+    yield ZAPv2(apikey='testapikey')
 
 
-@pytest.yield_fixture
-def urllib_mock():
+@pytest.yield_fixture(autouse=True)
+def client_mock():
     """Fixture create a mock for urllib library."""
-    with patch('zapv2.urllib.urlopen') as urllib_mock:
-        yield urllib_mock
+    with requests_mock.mock() as mock:
+        yield mock

tests/unit/test_client.py

@@ -1,39 +1,53 @@
 """
 Tests related to the main Zap Client class
 """
-from mock import call
 
 TEST_PROXIES = {
     'http': 'http://127.0.0.1:8080',
     'https': 'http://127.0.0.1:8080',
 }
 
 
-def test_urlopen_proxies(zap, urllib_mock):
-    """Check if Zap client passes proxy to urllib call."""
-    urllib_mock.return_value.read.return_value = 'contents'
+def assert_api_key(response, apikey='testapikey'):
+    """Some requests should contain valid ZAP api key."""
+    assert response._request.headers['X-ZAP-API-Key'] == apikey
+    assert 'apikey=%s' % apikey in response.query
 
-    assert zap.urlopen() == 'contents'
-    assert urllib_mock.mock_calls[0][2]['proxies'] == TEST_PROXIES
+
+def test_urlopen(zap, client_mock):
+    """Request method should return a python object from parsed output"""
+    api_response ='{"testkey": "testvalue"}'
+    client_mock.get('http://localhost:8080', text=api_response)
+
+    assert zap.urlopen('http://localhost:8080', {'querykey': 'queryvalue'}) == api_response
+
+    response = client_mock.request_history[0]
+
+    assert 'X-ZAP-API-Key' not in response._request.headers
+    assert 'testapikey' not in response.query
+    assert response.proxies == TEST_PROXIES
 
 
-def test_request_response(zap, urllib_mock):
+def test_request_response(zap, client_mock):
     """Request method should return a python object from parsed output"""
-    urllib_mock.return_value.read.return_value = '{"testkey": "testvalue"}'
+    client_mock.get('http://zap/test', text='{"testkey": "testvalue"}')
 
-    assert zap._request('http://allizom.org', {'querykey': 'queryvalue'}) == {'testkey': 'testvalue'}
-    assert urllib_mock.mock_calls == [
-        call('http://allizom.org?querykey=queryvalue', proxies=TEST_PROXIES),
-        call().read()
-    ]
+    assert zap._request('http://zap/test', {'querykey': 'queryvalue'}) == {'testkey': 'testvalue'}
 
+    response = client_mock.request_history[0]
 
-def test_request_other(zap, urllib_mock):
+    assert_api_key(response)
+    assert response.proxies == TEST_PROXIES
+
+
+def test_request_other(zap, client_mock):
     """_request_other should simply return a retrieved content."""
-    urllib_mock.return_value.read.return_value = '{"testkey": "testvalue"}'
+    api_response = '{"testkey": "testvalue"}'
+    client_mock.get('http://zap/test', text=api_response)
+
+    assert zap._request_other('http://zap/test', {'querykey': 'queryvalue'}) == api_response
+
+    response = client_mock.request_history[0]
 
-    assert zap._request('http://allizom.org', {'querykey': 'queryvalue'}) == {'testkey': 'testvalue'}
-    assert urllib_mock.mock_calls == [
-        call('http://allizom.org?querykey=queryvalue', proxies=TEST_PROXIES),
-        call().read()
-    ]
+    assert_api_key(response)
+    assert response.proxies == TEST_PROXIES