Fix/tweak resigning (#1868)

app/dao/api_key_dao.py

@@ -13,21 +13,24 @@
 
 
 @transactional
-def resign_api_keys(unsafe: bool = False):
+def resign_api_keys(resign: bool, unsafe: bool = False):
     """Resign the _secret column of the api_keys table with (potentially) a new key.
 
     Args:
+        resign (bool): whether to resign the api keys
         unsafe (bool, optional): resign regardless of whether the unsign step fails with a BadSignature.
         Defaults to False.
 
     Raises:
         e: BadSignature if the unsign step fails and unsafe is False.
     """
     rows = ApiKey.query.all()  # noqa
-    current_app.logger.info(f"Resigning {len(rows)} api keys")
+    current_app.logger.info(f"Total of {len(rows)} api keys")
+    rows_to_update = []
 
     for row in rows:
         try:
+            old_signature = row._secret
             unsigned_secret = getattr(row, "secret")  # unsign the secret
         except BadSignature as e:
             if unsafe:
@@ -36,7 +39,16 @@ def resign_api_keys(unsafe: bool = False):
                 current_app.logger.error(f"BadSignature for api_key {row.id}, using verify_unsafe instead")
                 raise e
         setattr(row, "secret", unsigned_secret)  # resigns the api key secret with (potentially) a new signing secret
-    db.session.bulk_save_objects(rows)
+        if old_signature != row._secret:
+            rows_to_update.append(row)
+        if not resign:
+            row._secret = old_signature  # reset the signature to the old value
+
+    if resign:
+        current_app.logger.info(f"Resigning {len(rows_to_update)} api keys")
+        db.session.bulk_save_objects(rows)
+    elif not resign:
+        current_app.logger.info(f"{len(rows_to_update)} api keys need resigning")
 
 
 @transactional

app/dao/inbound_sms_dao.py

@@ -11,21 +11,24 @@
 
 
 @transactional
-def resign_inbound_sms(unsafe: bool = False):
+def resign_inbound_sms(resign: bool, unsafe: bool = False):
     """Resign the _content column of the inbound_sms table with (potentially) a new key.
 
     Args:
+        resign (bool): whether to resign the inbound sms
         unsafe (bool, optional): resign regardless of whether the unsign step fails with a BadSignature.
         Defaults to False.
 
     Raises:
         e: BadSignature if the unsign step fails and unsafe is False.
     """
     rows = InboundSms.query.all()  # noqa
-    current_app.logger.info(f"Resigning {len(rows)} inbound sms")
+    current_app.logger.info(f"Total of {len(rows)} inbound sms")
+    rows_to_update = []
 
     for row in rows:
         try:
+            old_signature = row._content
             unsigned_content = getattr(row, "content")  # unsign the content
         except BadSignature as e:
             if unsafe:
@@ -34,7 +37,16 @@ def resign_inbound_sms(unsafe: bool = False):
                 current_app.logger.error(f"BadSignature for inbound_sms {row.id}")
                 raise e
         setattr(row, "content", unsigned_content)  # resigns the content with (potentially) a new signing secret
-    db.session.bulk_save_objects(rows)
+        if old_signature != row._content:
+            rows_to_update.append(row)
+        if not resign:
+            row._content = old_signature  # reset the signature to the old value
+
+    if resign:
+        current_app.logger.info(f"Resigning {len(rows_to_update)} inbound sms")
+        db.session.bulk_save_objects(rows)
+    elif not resign:
+        current_app.logger.info(f"{len(rows_to_update)} inbound sms need resigning")
 
 
 @transactional

app/dao/notifications_dao.py

@@ -59,21 +59,28 @@
 
 
 @transactional
-def resign_notifications(unsafe: bool = False):
-    """Resign the _personalisation column of the notifications table with (potentially) a new key.
+def _resign_notifications_chunk(chunk_offset: int, chunk_size: int, resign: bool, unsafe: bool) -> int:
+    """Resign the _personalisation column of the notifications in a chunk of notifications with (potentially) a new key.
 
     Args:
-        unsafe (bool, optional): resign regardless of whether the unsign step fails with a BadSignature.
-        Defaults to False.
+        chunk_offset (int): start index of the chunk
+        chunk_size (int): size of the chunk
+        resign (bool): resign the personalisation
+        unsafe (bool): ignore bad signatures
 
     Raises:
         e: BadSignature if the unsign step fails and unsafe is False.
+
+    Returns:
+        int: number of notifications resigned or needing to be resigned
     """
-    rows = Notification.query.all()  # noqa
-    current_app.logger.info(f"Resigning {len(rows)} notifications")
+    rows = Notification.query.order_by(Notification.created_at).slice(chunk_offset, chunk_offset + chunk_size).all()
+    current_app.logger.info(f"Processing chunk {chunk_offset} to {chunk_offset + len(rows) - 1}")
 
+    rows_to_update = []
     for row in rows:
-        if row._personalisation:
+        old_signature = row._personalisation
+        if old_signature:
             try:
                 unsigned_personalisation = getattr(row, "personalisation")  # unsign the personalisation
             except BadSignature as e:
@@ -85,7 +92,49 @@ def resign_notifications(unsafe: bool = False):
         setattr(
             row, "personalisation", unsigned_personalisation
         )  # resigns the personalisation with (potentially) a new signing secret
-    db.session.bulk_save_objects(rows)
+        if old_signature != row._personalisation:
+            rows_to_update.append(row)
+        if not resign:
+            row._personalisation = old_signature  # reset the signature to the old value
+
+    if resign and len(rows_to_update) > 0:
+        current_app.logger.info(f"Resigning {len(rows_to_update)} notifications")
+        db.session.bulk_save_objects(rows)
+    elif len(rows_to_update) > 0:
+        current_app.logger.info(f"{len(rows_to_update)} notifications need resigning")
+
+    return len(rows_to_update)
+
+
+def resign_notifications(chunk_size: int, resign: bool, unsafe: bool = False) -> int:
+    """Resign the _personalisation column of the notifications table with (potentially) a new key.
+
+    Args:
+        chunk_size (int): number of rows to update at once.
+        resign (bool): resign the notifications.
+        unsafe (bool, optional): resign regardless of whether the unsign step fails with a BadSignature. Defaults to False.
+        max_update_size(int, -1): max number of rows to update at once, -1 for no limit. Defautls to -1.
+
+    Returns:
+        int: number of notifications that were resigned or need to be resigned.
+
+    Raises:
+        e: BadSignature if the unsign step fails and unsafe is False.
+    """
+
+    total_notifications = Notification.query.count()
+    current_app.logger.info(f"Total of {total_notifications} notifications")
+    num_old_signatures = 0
+
+    for chunk_offset in range(0, total_notifications, chunk_size):
+        num_old_signatures_in_chunk = _resign_notifications_chunk(chunk_offset, chunk_size, resign, unsafe)
+        num_old_signatures += num_old_signatures_in_chunk
+
+    if resign:
+        current_app.logger.info(f"Overall, {num_old_signatures} notifications were resigned")
+    else:
+        current_app.logger.info(f"Overall, {num_old_signatures} notifications need resigning")
+    return num_old_signatures
 
 
 @statsd(namespace="dao")

app/dao/service_callback_api_dao.py

@@ -13,22 +13,25 @@
 
 
 @transactional
-def resign_service_callbacks(unsafe: bool = False):
+def resign_service_callbacks(resign: bool, unsafe: bool = False):
     """Resign the _bearer_token column of the service_callbacks table with (potentially) a new key.
 
     Args:
+        resign (bool): whether to resign the service_callbacks
         unsafe (bool, optional): resign regardless of whether the unsign step fails with a BadSignature.
         Defaults to False.
 
     Raises:
         e: BadSignature if the unsign step fails and unsafe is False.
     """
     rows = ServiceCallbackApi.query.all()  # noqa
-    current_app.logger.info(f"Resigning {len(rows)} service_callbacks")
+    current_app.logger.info(f"Total of {len(rows)} service callbacks")
+    rows_to_update = []
 
     for row in rows:
         if row._bearer_token:
             try:
+                old_signature = row._bearer_token
                 unsigned_token = getattr(row, "bearer_token")  # unsign the token
             except BadSignature as e:
                 if unsafe:
@@ -37,7 +40,16 @@ def resign_service_callbacks(unsafe: bool = False):
                     current_app.logger.error(f"BadSignature for service_callback {row.id}")
                     raise e
             setattr(row, "bearer_token", unsigned_token)  # resigns the token with (potentially) a new signing secret
-    db.session.bulk_save_objects(rows)
+            if old_signature != row._bearer_token:
+                rows_to_update.append(row)
+            if not resign:
+                row._bearer_token = old_signature  # reset the signature to the old value
+
+    if resign:
+        current_app.logger.info(f"Resigning {len(rows_to_update)} service callbacks")
+        db.session.bulk_save_objects(rows)
+    elif not resign:
+        current_app.logger.info(f"{len(rows_to_update)} service callbacks need resigning")
 
 
 @transactional

scripts/resign_database.py

@@ -19,28 +19,34 @@
 
 sys.path.append('..')  # needed so we can find app (as run from scripts/ folder)
 
+from flask import current_app  # noqa: E402
+
 from app import create_app  # noqa: E402
 from app.dao.api_key_dao import resign_api_keys  # noqa: E402
 from app.dao.inbound_sms_dao import resign_inbound_sms  # noqa: E402
 from app.dao.notifications_dao import resign_notifications  # noqa: E402
 from app.dao.service_callback_api_dao import resign_service_callbacks  # noqa: E402
 
 
-def resign_all(unsafe: bool = False):
-    resign_api_keys(unsafe)
-    resign_inbound_sms(unsafe)
-    resign_service_callbacks(unsafe)
-    resign_notifications(unsafe)
+def resign_all(chunk: int, resign: bool, unsafe: bool):
+    resign_api_keys(resign, unsafe)
+    resign_inbound_sms(resign, unsafe)
+    resign_service_callbacks(resign, unsafe)
+    resign_notifications(chunk, resign, unsafe)
+    if not resign:
+        current_app.logger.info("NOTE: this is a preview, fields have not been changed. To resign fields, run with --resign flag")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--unsafe", default=False, action='store_true', help="resign notifications that have a bad signature")
+    parser.add_argument("--chunk", default=25000, type=int, help="size of chunks of notifications to resign at a time (default 25000)")
+    parser.add_argument("--resign", default=False, action='store_true', help="resign columns (default false)")
+    parser.add_argument("--unsafe", default=False, action='store_true', help="ignore bad signatures (default false)")
     args = parser.parse_args()
 
     load_dotenv()
     application = Flask("resign_database")
     create_app(application)
     application.app_context().push()
 
-    resign_all(args.unsafe)
+    resign_all(args.chunk, args.resign, args.unsafe)

tests/app/dao/notification_dao/test_notification_dao.py

@@ -1803,20 +1803,31 @@ def test_bulk_insert_notification_duplicate_ids(self, sample_template):
 
 
 class TestResigning:
-    def test_resign_notifications_resigns_with_new_key(self, sample_template_with_placeholders):
+    @pytest.mark.parametrize("resign,chunk_size", [(True, 2), (False, 2), (True, 10), (False, 10)])
+    def test_resign_notifications_resigns_or_previews(self, resign, chunk_size, sample_template_with_placeholders):
         from app import signer_personalisation
 
         with set_signer_secret_key(signer_personalisation, ["k1", "k2"]):
-            initial_notification = create_notification(sample_template_with_placeholders, personalisation={"Name": "test"})
-            save_notification(initial_notification)
-            personalisation = initial_notification.personalisation
-            _personalisation = initial_notification._personalisation
+            initial_notifications = [
+                create_notification(sample_template_with_placeholders, personalisation={"Name": "test"}) for _ in range(5)
+            ]
+            personalisations = [n.personalisation for n in initial_notifications]
+            _personalisations = [n._personalisation for n in initial_notifications]
+            for notification in initial_notifications:
+                save_notification(notification)
 
         with set_signer_secret_key(signer_personalisation, ["k2", "k3"]):
-            resign_notifications()
-            notification = Notification.query.get(initial_notification.id)
-            assert notification.personalisation == personalisation  # unsigned value is the same
-            assert notification._personalisation != _personalisation  # signature is different
+            resign_notifications(chunk_size=chunk_size, resign=resign)
+            notifications = [Notification.query.get(n.id) for n in initial_notifications]
+            assert [n.personalisation for n in notifications] == personalisations  # unsigned values are the same
+            if resign:
+                for (
+                    notification,
+                    _personalisation,
+                ) in zip(notifications, _personalisations):
+                    assert notification._personalisation != _personalisation  # signature is different
+            else:
+                assert [n._personalisation for n in notifications] == _personalisations  # signatures are the same
 
     def test_resign_notifications_fails_if_cannot_verify_signatures(self, sample_template_with_placeholders):
         from app import signer_personalisation
@@ -1827,7 +1838,7 @@ def test_resign_notifications_fails_if_cannot_verify_signatures(self, sample_tem
 
         with set_signer_secret_key(signer_personalisation, ["k3"]):
             with pytest.raises(BadSignature):
-                resign_notifications()
+                resign_notifications(chunk_size=10, resign=True)
 
     def test_resign_notifications_unsafe_resigns_with_new_key(self, sample_template_with_placeholders):
         from app import signer_personalisation
@@ -1839,7 +1850,7 @@ def test_resign_notifications_unsafe_resigns_with_new_key(self, sample_template_
             _personalisation = initial_notification._personalisation
 
         with set_signer_secret_key(signer_personalisation, ["k3"]):
-            resign_notifications(unsafe=True)
+            resign_notifications(chunk_size=10, resign=True, unsafe=True)
             notification = Notification.query.get(initial_notification.id)
             assert notification.personalisation == personalisation  # unsigned value is the same
             assert notification._personalisation != _personalisation  # signature is different

tests/app/dao/test_api_key_dao.py

@@ -166,7 +166,8 @@ def test_should_not_return_revoked_api_keys_older_than_7_days(sample_service, da
 
 
 class TestResigning:
-    def test_resign_api_keys_resigns_with_new_key(self, sample_service):
+    @pytest.mark.parametrize("resign", [True, False])
+    def test_resign_api_keys_resigns_or_previews(self, resign, sample_service):
         from app import signer_api_key
 
         with set_signer_secret_key(signer_api_key, ["k1", "k2"]):
@@ -175,10 +176,13 @@ def test_resign_api_keys_resigns_with_new_key(self, sample_service):
             _secret = initial_key._secret
 
         with set_signer_secret_key(signer_api_key, ["k2", "k3"]):
-            resign_api_keys()
+            resign_api_keys(resign=resign)
             api_key = ApiKey.query.get(initial_key.id)
             assert api_key.secret == secret  # unsigned value is the same
-            assert api_key._secret != _secret  # signature is different
+            if resign:
+                assert api_key._secret != _secret  # signature is different
+            else:
+                assert api_key._secret == _secret  # signature is the same
 
     def test_resign_api_keys_fails_if_cannot_verify_signatures(self, sample_service):
         from app import signer_api_key
@@ -188,7 +192,7 @@ def test_resign_api_keys_fails_if_cannot_verify_signatures(self, sample_service)
 
         with set_signer_secret_key(signer_api_key, "k3"):
             with pytest.raises(BadSignature):
-                resign_api_keys()
+                resign_api_keys(resign=True)
 
     def test_resign_api_keys_unsafe_resigns_with_new_key(self, sample_service):
         from app import signer_api_key
@@ -199,7 +203,7 @@ def test_resign_api_keys_unsafe_resigns_with_new_key(self, sample_service):
             _secret = initial_key._secret
 
         with set_signer_secret_key(signer_api_key, ["k3"]):
-            resign_api_keys(unsafe=True)
+            resign_api_keys(resign=True, unsafe=True)
             api_key = ApiKey.query.get(initial_key.id)
             assert api_key.secret == secret  # unsigned value is the same
             assert api_key._secret != _secret  # signature is different