Harden MARTA API requests and credential handling

marta/api.py

@@ -1,21 +1,29 @@
+import json
+from functools import wraps
+from os import getenv
 from typing import List, Union
 
-import requests
 import requests_cache
-import json
-from os import getenv
-from functools import wraps
 
 from .exceptions import APIKeyError, InvalidDirectionError
 from .vehicles import Bus, Train
 
 _CACHE_EXPIRE = int(getenv('MARTA_CACHE_EXPIRE', 30))
-_BASE_URL = 'http://developer.itsmarta.com'
+_BASE_URL = 'https://developer.itsmarta.com'
 _TRAIN_PATH = '/RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals'
 _BUS_PATH = '/BRDRestService/RestBusRealTimeService/GetAllBus'
 _BUS_ROUTE_PATH = '/BRDRestService/RestBusRealTimeService/GetBusByRoute/'
-
-requests_cache.install_cache('marta_api_cache', backend='sqlite', expire_after=_CACHE_EXPIRE)
+_REQUEST_TIMEOUT = (3.05, 10)
+_MAX_RESPONSE_BYTES = 5 * 1024 * 1024
+_JSON_CONTENT_TYPES = {'application/json', 'text/json'}
+
+# Keep caching private to this library and in memory. Importing marta must not
+# patch requests globally or persist request URLs containing credentials.
+CACHE = requests_cache.CachedSession(
+    backend='memory',
+    expire_after=_CACHE_EXPIRE,
+    ignored_parameters=['apikey'],
+)
 
 def require_api_key(func):
     """
@@ -31,41 +39,84 @@ def with_key(self, *args, **kwargs):
 
     return with_key
 
+def get_bus_direction(user_direction) -> str:
+    direction = user_direction.lower()
+    if direction.startswith('n'):
+        return 'Northbound'
+    elif direction.startswith('s'):
+        return 'Southbound'
+    elif direction.startswith('e'):
+        return 'Eastbound'
+    elif direction.startswith('w'):
+        return 'Westbound'
+    else:
+        return None
+
+def get_train_direction(user_direction) -> str:
+    train_direction = user_direction.lower()
+    if train_direction.startswith('n'):
+        return 'N'
+    elif train_direction.startswith('s'):
+        return 'S'
+    elif train_direction.startswith('e'):
+        return 'E'
+    elif train_direction.startswith('w'):
+        return 'W'
+    else:
+        return None
+
+
 def _convert_direction(user_direction: str, vehicle_type: str = 'bus') -> Union[str, None]:
     if not user_direction:
         return None
     if vehicle_type == 'bus':
-        if user_direction.lower().startswith('n'):
-            return 'Northbound'
-        elif user_direction.lower().startswith('s'):
-            return 'Southbound'
-        elif user_direction.lower().startswith('e'):
-            return 'Eastbound'
-        elif user_direction.lower().startswith('w'):
-            return 'Westbound'
+        bus_direction = get_bus_direction(user_direction)
+        if bus_direction is not None:
+            return bus_direction
         else:
             raise InvalidDirectionError(direction_provided=user_direction)
     elif vehicle_type == 'train':
-        if user_direction.lower().startswith('n'):
-            return 'N'
-        elif user_direction.lower().startswith('s'):
-            return 'S'
-        elif user_direction.lower().startswith('e'):
-            return 'E'
-        elif user_direction.lower().startswith('w'):
-            return 'W'
+        train_direction = get_train_direction(user_direction)
+        if train_direction is not None:
+            return train_direction
         else:
             raise InvalidDirectionError(direction_provided=user_direction)
     else:
         return user_direction
 
 
-def _get_data(endpoint: str, api_key: str) -> dict:
-    url = f'{_BASE_URL}{endpoint}?apikey={api_key}'
-    response = requests.get(url)
-    if response.status_code == 401 or response.status_code == 403:
-        raise APIKeyError(f'Your API key seems to be invalid. Try visiting {url}.')
-    return json.loads(response.text)
+def _get_data(endpoint: str, api_key: str) -> List[dict]:
+    url = f'{_BASE_URL}{endpoint}'
+    response = CACHE.get(
+        url,
+        params={'apikey': api_key},
+        timeout=_REQUEST_TIMEOUT,
+        stream=True,
+    )
+    if response.status_code in (401, 403):
+        raise APIKeyError('Your MARTA API key was rejected.')
+    response.raise_for_status()
+
+    content_type = response.headers.get('Content-Type', '').split(';', 1)[0].lower()
+    if content_type and content_type not in _JSON_CONTENT_TYPES:
+        raise ValueError(f'Unexpected MARTA API content type: {content_type}')
+
+    content_length = response.headers.get('Content-Length')
+    if content_length and int(content_length) > _MAX_RESPONSE_BYTES:
+        raise ValueError('MARTA API response exceeds the maximum allowed size')
+
+    chunks = []
+    received = 0
+    for chunk in response.iter_content(chunk_size=64 * 1024):
+        received += len(chunk)
+        if received > _MAX_RESPONSE_BYTES:
+            raise ValueError('MARTA API response exceeds the maximum allowed size')
+        chunks.append(chunk)
+
+    data = json.loads(b''.join(chunks).decode(response.encoding or 'utf-8'))
+    if not isinstance(data, list) or any(not isinstance(item, dict) for item in data):
+        raise ValueError('MARTA API response must be a list of objects')
+    return data
 
 def _filter_response(response: dict, filters: dict) -> List[dict]:
     valid_items = []

marta/exceptions.py

@@ -4,12 +4,12 @@ def __init__(self, message: str = None):
 
         if not message:
             message = 'API Key is missing. Please set MARTA_API_KEY or use api_key kwarg.'
-        super(Exception, self).__init__(message)
+        super().__init__(message)
 
 class InvalidDirectionError(Exception):
     """Exception thrown for an invalid bus/train direction"""
     def __init__(self, direction_provided: str, message: str = None):
 
         if not message:
             message = f'{direction_provided} is an invalid direction.'
-        super(Exception, self).__init__(message)
+        super().__init__(message)

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+pythonpath = .
+testpaths = tests

requirements.txt

@@ -1,2 +1,2 @@
-requests
-requests-cache
+requests>=2.32.4,<3
+requests-cache>=1.2,<2

setup.py

@@ -3,11 +3,11 @@
 setup(
     name='marta',
     description='Python library for accessing MARTA real-time API',
-    url='http://www.itsmarta.com/app-developer-resources.aspx',
+    url='https://www.itsmarta.com/app-developer-resources.aspx',
     packages=['marta'],
     install_requires=[
         'pytest-runner',
-        'requests',
-        'requests-cache'
+        'requests>=2.32.4,<3',
+        'requests-cache>=1.2,<2'
     ],
 )