Merge pull request #104 from victor-torres/shub-jobauth

fallback to SHUB_JOBAUTH environment variable if SH_APIKEY is not set

Author: vshlapakov

scrapinghub/client/utils.py

@@ -4,6 +4,7 @@
 import json
 import logging
 import binascii
+import warnings
 from codecs import decode
 
 import six
@@ -101,10 +102,17 @@ def parse_auth(auth):
     """
     if auth is None:
         apikey = os.environ.get('SH_APIKEY')
-        if apikey is None:
-            raise RuntimeError("No API key provided and SH_APIKEY "
-                               "environment variable not set")
-        return (apikey, '')
+        if apikey:
+            return (apikey, '')
+
+        jobauth = os.environ.get('SHUB_JOBAUTH')
+        if jobauth:
+            warnings.warn("You are using the SHUB_JOBAUTH environment "
+                          "variable which may not work for some API endpoints")
+            return _search_for_jwt_credentials(jobauth)
+
+        raise RuntimeError("No API key provided and neither SH_APIKEY "
+                           "nor SHUB_JOBAUTH environment variables is set")
 
     if isinstance(auth, tuple):
         all_strings = all(isinstance(k, six.string_types) for k in auth)

tests/client/test_utils.py

@@ -17,6 +17,11 @@ def test_parse_auth_none_with_env():
     assert parse_auth(None) == ('testkey', '')
 
 
+@mock.patch.dict(os.environ, {'SH_APIKEY': 'testkey', 'SHUB_JOBAUTH': 'jwt'})
+def test_parse_auth_none_with_multiple_env():
+    assert parse_auth(None) == ('testkey', '')
+
+
 def test_parse_auth_tuple():
     assert parse_auth(('test', 'test')) == ('test', 'test')
     assert parse_auth(('apikey', '')) == ('apikey', '')
@@ -52,6 +57,28 @@ def test_parse_auth_jwt_token():
     assert parse_auth(encoded_token) == (test_job, test_token)
 
 
+def test_parse_auth_jwt_token_with_jwt_token_env():
+    dummy_test_job, dummy_test_token = '1/2/3', 'some.dummy.jwt.token'
+    raw_token = (dummy_test_job + ':' + dummy_test_token).encode('utf8')
+    dummy_encoded_token = encode(raw_token, 'hex_codec').decode('ascii')
+
+    test_job, test_token = '1/2/3', 'some.jwt.token'
+    raw_token = (test_job + ':' + test_token).encode('utf8')
+    encoded_token = encode(raw_token, 'hex_codec').decode('ascii')
+
+    with mock.patch.dict(os.environ, {'SHUB_JOBAUTH': dummy_encoded_token}):
+        assert parse_auth(encoded_token) == (test_job, test_token)
+
+
+def test_parse_auth_none_with_jwt_token_env():
+    test_job, test_token = '1/2/3', 'some.jwt.token'
+    raw_token = (test_job + ':' + test_token).encode('utf8')
+    encoded_token = encode(raw_token, 'hex_codec').decode('ascii')
+
+    with mock.patch.dict(os.environ, {'SHUB_JOBAUTH': encoded_token}):
+        assert parse_auth(None) == (test_job, test_token)
+
+
 def test_parse_job_key():
     job_key = parse_job_key('123/10/11')
     assert job_key.project_id == '123'