Add error handling to the API blueprint.

src/ocspdash/manager.py

@@ -11,11 +11,11 @@
 from dataclasses import dataclass
 from itertools import groupby
 from operator import attrgetter
-from typing import Iterable, List, Mapping, Optional, Tuple
+from typing import Iterable, List, Mapping, Optional, Tuple, Union
 
 from sqlalchemy import and_, create_engine, func
 from sqlalchemy.engine import Engine
-from sqlalchemy.orm import scoped_session, sessionmaker
+from sqlalchemy.orm import Session, scoped_session, sessionmaker
 
 from ocspdash.constants import OCSPDASH_DEFAULT_CONNECTION, OCSPDASH_USER_AGENT_IDENTIFIER
 from ocspdash.models import Authority, Base, Chain, Location, Responder, Result
@@ -28,6 +28,8 @@
 
 logger = logging.getLogger(__name__)
 
+_SessionHint = Union[Session, scoped_session]
+
 
 def _workaround_pysqlite_transaction_bug():
     """Work around pysqlite transaction bug.
@@ -88,12 +90,13 @@ class ResponderPayload:
 class Manager:
     """Manager for interacting with the database."""
 
-    def __init__(self, engine: Engine, session: scoped_session, server_query: Optional[ServerQuery] = None) -> None:
+    def __init__(self, engine: Engine, session: _SessionHint, server_query: Optional[ServerQuery] = None) -> None:
         """Instantiate a Manager with instances of the objects it needs.
 
         :param engine: The database engine.
         :param session: The database session.
-        :param server_query: The server_query instance. If None, using server_query-related functionality will raise an error.
+        :param server_query: The server_query instance. If None, using server_query-related functionality will raise
+         an error.
         """
         self.engine = engine
         self.session = session
@@ -108,7 +111,8 @@ def from_args(cls, connection: Optional[str] = None, echo: bool = False, api_id:
         :param connection: An SQLAlchemy-compatible connection string.
         :param echo: True to echo SQL emitted by SQLAlchemy.
         :param api_id: The Censys API id. If None, the value will be obtained from configuration or the environment.
-        :param api_secret: The Censys API secret. If none, the value will be obtained from configuration or the environment.
+        :param api_secret: The Censys API secret. If none, the value will be obtained from configuration or the
+         environment.
 
         :returns: An instance of Manager configured according to the arguments provided.
         """
@@ -306,10 +310,7 @@ def ensure_chain(self, responder: Responder) -> Optional[Chain]:
         most_recent_chain = self.get_most_recent_chain_by_responder(responder)
 
         if most_recent_chain and not most_recent_chain.old:
-            if not most_recent_chain.expired:
-                return most_recent_chain
-
-            if not responder.current:
+            if not most_recent_chain.expired or not responder.current:
                 return most_recent_chain
 
         subject, issuer = self.server_query.get_certs_for_issuer_and_url(responder.authority.name, responder.url)
@@ -510,12 +511,13 @@ def get_location_by_selector(self, selector: bytes) -> Optional[Location]:
         return self.session.query(Location).filter(Location.selector == selector).one_or_none()
 
     def process_location(self, invite_token: bytes, public_key: str) -> Optional[Location]:
-        """Given an invite token and public key, check for a valid invite and associate the public key with the corresponding location.
+        """Check for a valid invite and associate the public key with the corresponding location.
 
         :parameter invite_token: a 32-byte string corresponding to an invited Location.
         :parameter public_key: The public key to be associated with the Location.
 
         :returns: The Location if a valid invite was provided, otherwise None.
+        :raises ValueError: with an appropriate message for failures to process the location invite.
         """
         if len(invite_token) != 32:
             raise ValueError('invite_token of wrong length')
@@ -524,11 +526,11 @@ def process_location(self, invite_token: bytes, public_key: str) -> Optional[Loc
 
         location = self.get_location_by_selector(selector)
         if location is None:
-            raise Exception(f'location not found for selector: {selector}')
+            raise ValueError(f'invalid invite token')
         if location.pubkey:  # this invite has already been used
-            return None
+            raise ValueError(f'invite expired')
         if not location.verify(validator):
-            return None
+            raise ValueError(f'invalid invite token')
 
         location.set_public_key(public_key)
 
@@ -577,7 +579,10 @@ def get_most_recent_chains_for_authorities(self, n: Optional[int] = 10) -> List[
         return query.all()
 
     def insert_payload(self, location: Location, results: Iterable[Mapping]):
-        """Take the submitted payload and insert its results into the database."""
+        """Take the submitted payload and insert its results into the database.
+
+        :raises sqlalchemy.IntegrityError: If a duplicate location is trying to be inserted. Rolls back.
+        """
         for prepared_result_dict in results:
             result = Result(**prepared_result_dict)
             result.location = location

src/ocspdash/models.py

@@ -226,6 +226,8 @@ def set_public_key(self, public_key: str):
         """Set the pubkey and key_id for the Location based on an input public key.
 
         :param public_key: The public key for the Location.
+
+        :raises ValueError: if the passed-in public key is not of a valid algorithm as defined in constants.OCSPSCRAPE_PRIVATE_KEY_ALGORITHMS.
         """
         pubkey = b64decode(public_key)
         loaded_pubkey = serialization.load_pem_public_key(pubkey, default_backend())

src/ocspdash/web/app.py

@@ -4,7 +4,7 @@
 
 import logging
 import os
-from typing import Optional
+from typing import Mapping, Optional
 
 from flasgger import Swagger
 from flask import Flask
@@ -23,22 +23,29 @@
 logger = logging.getLogger('web')
 
 
-def create_application(connection: Optional[str] = None, flask_debug: bool = False) -> Flask:
+def create_application(connection: Optional[str] = None, flask_debug: bool = False,
+                       db_session_options: Optional[Mapping] = None, secret_key=None) -> Flask:
     """Create the OCSPdash Flask application.
 
-    :param connection: Database connection string
+    :param connection: Database connection string, overridden by env $OCSPDASH_CONNECTION
     :param flask_debug: Enable Flask debug mode, overridden by env $DEBUG
+    :param db_session_options: Mapping of options passed to the flask-sqlalchemy constructor for eventual passage to
+     sqlalchemy.sessionmaker.
     """
     app = Flask(__name__)
     app.config.update(dict(
         SQLALCHEMY_DATABASE_URI=connection or OCSPDASH_CONNECTION,
         SQLALCHEMY_TRACK_MODIFICATIONS=False,
-        SECRET_KEY=os.environ.get('SECRET_KEY', 'test key'),
+        SECRET_KEY=secret_key or os.environ['SECRET_KEY'],  # keep it secret; keep it safe
         DEBUG=os.environ.get('DEBUG', flask_debug),
         CENSYS_API_ID=os.environ.get('CENSYS_API_ID'),
         CENSYS_API_SECRET=os.environ.get('CENSYS_API_SECRET'),
+        OCSPDASH_API_MANIFEST_DEFAULT_SIZE=int(os.environ.get('OCSPDASH_API_MANIFEST_DEFAULT_SIZE', default=10)),
+        OCSPDASH_API_MANIFEST_MAX_SIZE=int(os.environ.get('OCSPDASH_API_MANIFEST_MAX_SIZE', default=10)),
+
     ))
-    db = OCSPSQLAlchemy(app=app)
+
+    db = OCSPSQLAlchemy(app=app, session_options=db_session_options)
 
     Bootstrap(app)
     Swagger(app)  # Adds Swagger UI

src/ocspdash/web/blueprints/api.py

@@ -2,6 +2,7 @@
 
 """The OCSPdash API blueprint."""
 
+import binascii
 import io
 import logging
 import uuid
@@ -11,11 +12,13 @@
 from http import HTTPStatus
 
 import jsonlines
-from flask import Blueprint, abort, request
+from flask import Blueprint, current_app, jsonify, request
 from jose import jwt
 from jose.exceptions import JWTError
+from sqlalchemy.exc import IntegrityError
 
 from ocspdash.constants import OCSP_JWT_ALGORITHM, OCSP_RESULTS_JWT_CLAIM
+from ocspdash.web.exceptions import InvalidUsage
 from ocspdash.web.proxies import manager
 
 jwt.decode = partial(jwt.decode, algorithms=OCSP_JWT_ALGORITHM)
@@ -25,25 +28,45 @@
 api = Blueprint('api', __name__)
 
 
-@api.route('/register', methods=['POST'])
+@api.route('/register', methods=['POST'])  # noqa: C901
 def register_location_key():
     """Register a public key for an invited location."""
     # TODO: error handling (what if no invite, what if duplicate name, etc.)
-    unverified_claims = jwt.get_unverified_claims(request.data)
-    unverified_public_key = b64decode(unverified_claims['pk']).decode('utf-8')
+    try:
+        unverified_claims = jwt.get_unverified_claims(request.data)
+    except jwt.JWTError as e:
+        raise InvalidUsage(f'failed to decode JWT: {str(e)}')
+
+    try:
+        unverified_public_key = b64decode(unverified_claims['pk']).decode('utf-8')
+    except KeyError as e:
+        raise InvalidUsage(f'missing claim: {str(e)}')
+    except binascii.Error as e:
+        raise InvalidUsage(f"failed to decode 'pk' claim: {str(e)}")
+    except UnicodeError as e:
+        raise InvalidUsage(f"failed to decode 'pk' claim: {e.reason}")
 
     try:
         claims = jwt.decode(request.data, unverified_public_key)
-    except JWTError:
-        return abort(400)  # bad input
+    except JWTError as e:
+        raise InvalidUsage(f'failed to decode JWT: {str(e)}')
 
-    public_key = claims['pk']
-    invite_token = b64decode(claims['token'])
+    try:
+        public_key = claims['pk']
+    except KeyError as e:
+        raise InvalidUsage(f'missing claim: {str(e)}')
 
-    new_location = manager.process_location(invite_token, public_key)
+    try:
+        invite_token = b64decode(claims['token'])
+    except KeyError as e:
+        raise InvalidUsage(f'missing claim: {str(e)}')
+    except binascii.Error as e:
+        raise InvalidUsage(f"failed to decode 'token' claim: {str(e)}")
 
-    if new_location is None:
-        return abort(400)
+    try:
+        manager.process_location(invite_token, public_key)
+    except ValueError as e:
+        raise InvalidUsage(f'failed to process invite: {str(e)}')
 
     return '', HTTPStatus.NO_CONTENT
 
@@ -63,9 +86,9 @@ def get_manifest():
         required: false
         type: integer
     """
-    n = request.args.get('n', type=int, default=10)  # TODO make configurable at app level
-    if n > 10:
-        abort(400, 'n too large, max is 10')  # TODO get the max config value here too
+    n = request.args.get('n', type=int, default=current_app.config['OCSPDASH_API_MANIFEST_DEFAULT_SIZE'])
+    if n > current_app.config['OCSPDASH_API_MANIFEST_MAX_SIZE']:
+        raise InvalidUsage(f'n too large, max is {current_app.config["OCSPDASH_API_MANIFEST_MAX_SIZE"]}: {n}')
     manifest_lines = io.StringIO()
     with jsonlines.Writer(manifest_lines, sort_keys=True) as writer:
         writer.write_all(
@@ -80,7 +103,10 @@ def get_manifest():
 
 def _prepare_result_dictionary(result_data):
     certificate_chain_uuid: uuid.UUID = uuid.UUID(result_data['certificate_chain_uuid'])
+
     chain = manager.get_chain_by_certificate_chain_uuid(certificate_chain_uuid)
+    if not chain:
+        raise ValueError(f'no chain with certificate_chain_uuid: {certificate_chain_uuid}')
 
     retrieved = datetime.strptime(result_data['time'], '%Y-%m-%dT%H:%M:%SZ')
 
@@ -92,26 +118,61 @@ def _prepare_result_dictionary(result_data):
     }
 
 
-@api.route('/submit', methods=['POST'])
+@api.route('/submit', methods=['POST'])  # noqa: C901
 def submit():
     """Submit scrape results.
 
     ---
     tags:
         - ocsp
     """
-    submitted_token_header = jwt.get_unverified_header(request.data)
+    try:
+        submitted_token_header = jwt.get_unverified_header(request.data)
+    except jwt.JWTError as e:
+        raise InvalidUsage(f'failed to decode JWT: {str(e)}')
+
+    try:
+        key_id = uuid.UUID(submitted_token_header['kid'])
+    except KeyError as e:
+        raise InvalidUsage(f'missing header claim: {str(e)}')
 
-    key_id = uuid.UUID(submitted_token_header['kid'])
     submitting_location = manager.get_location_by_key_id(key_id)
+    if not submitting_location:
+        raise InvalidUsage(f'no location for key id', payload={'key_id': key_id})
 
     try:
         claims = jwt.decode(request.data, submitting_location.pubkey.decode('utf-8'))
-    except JWTError:
-        return abort(400)
+    except JWTError as e:
+        raise InvalidUsage(f'failed to decode JWT: {str(e)}')
 
-    prepared_result_dicts = (_prepare_result_dictionary(result_data)
-                             for result_data in claims[OCSP_RESULTS_JWT_CLAIM])
-    manager.insert_payload(submitting_location, prepared_result_dicts)
+    try:
+        results = claims[OCSP_RESULTS_JWT_CLAIM]
+    except KeyError as e:
+        raise InvalidUsage(f'missing claim: {str(e)}')
+
+    prepared_result_dicts = []
+    for result_data in results:
+        try:
+            prepared_dict = _prepare_result_dictionary(result_data)
+            prepared_result_dicts.append(prepared_dict)
+        except (KeyError, ValueError):
+            raise InvalidUsage('invalid result data', payload={'result': result_data})
 
-    return ('', HTTPStatus.NO_CONTENT)
+    try:
+        manager.insert_payload(submitting_location, prepared_result_dicts)
+    except IntegrityError:
+        manager.session.rollback()
+        raise
+
+    return '', HTTPStatus.NO_CONTENT
+
+
+@api.errorhandler(InvalidUsage)
+def handle_invalid_usage(error: InvalidUsage):
+    """Handle InvalidUsage exceptions raised by views in the blueprint.
+
+    :param error: The exception that caused this handler to be called.
+    """
+    response = jsonify(error)
+    response.status_code = error.status_code
+    return response

src/ocspdash/web/exceptions.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+
+"""Exceptions and Errors for the OCSPdash web package."""
+
+from http import HTTPStatus
+from typing import Mapping
+
+
+class InvalidUsage(Exception):
+    """An Exception to be raised by a view for invalid usage of the endpoint."""
+
+    status_code = HTTPStatus.BAD_REQUEST
+
+    def __init__(self, message: str, status_code: HTTPStatus=None, payload: Mapping=None) -> None:
+        """Create an InvalidUsage exception.
+
+        :param message: The message for the exception; will be placed in the `message` key of the JSON returned in the response.
+        :param status_code: An HTTP status code for the response; default is 400 BAD REQUEST. It will also be placed in the `status` key of the returned JSON.
+        :param payload: A mapping that can be jsonified by Flask; will be returned as JSON in the response alongside the message.
+        """
+        super().__init__(self)
+        self.message = message
+        if status_code is not None:
+            self.status_code = status_code
+        self.payload = payload
+
+    def to_json(self) -> Mapping:
+        """Return a representation of the exception suitable to passed for JSON conversion."""
+        rv = dict(self.payload or ())
+        rv['message'] = self.message
+        rv['status'] = self.status_code
+        return rv

src/ocspdash/web/extension.py

@@ -2,6 +2,8 @@
 
 """An extension of Flask-SQLAlchemy to support the OCSPdash manager."""
 
+from __future__ import annotations
+
 from flask import Flask
 from flask_sqlalchemy import SQLAlchemy, get_state
 
@@ -26,6 +28,14 @@ def init_app(self, app: Flask):
 
         self.manager = Manager(engine=self.engine, session=self.session)
 
+    @staticmethod
+    def get_db(app: Flask) -> SQLAlchemy:  # noqa: F821
+        """Get the db from an app.
+
+        :param app: A Flask app
+        """
+        return get_state(app).db
+
     @staticmethod
     def get_manager(app: Flask) -> Manager:
         """Get the manager from an app.