utils: function for rebuilding encrypted fields

* Adds utility function to handle changes of the SECRET_KEY.
  Decrypts entries using the old key and saves using the current one.

Signed-off-by: Dinos Kousidis <konstantinos.kousidis@cern.ch>

invenio_db/shared.py

@@ -46,7 +46,7 @@ class SQLAlchemy(FlaskSQLAlchemy):
     """Implement or overide extension methods."""
 
     def apply_driver_hacks(self, app, info, options):
-        """Called before engine creation."""
+        """Call before engine creation."""
         # Don't forget to apply hacks defined on parent object.
         super(SQLAlchemy, self).apply_driver_hacks(app, info, options)
 

invenio_db/utils.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+#
+# This file is part of Invenio.
+# Copyright (C) 2017 CERN.
+#
+# Invenio is free software; you can redistribute it
+# and/or modify it under the terms of the GNU General Public License as
+# published by the Free Software Foundation; either version 2 of the
+# License, or (at your option) any later version.
+#
+# Invenio is distributed in the hope that it will be
+# useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Invenio; if not, write to the
+# Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
+# MA 02111-1307, USA.
+#
+# In applying this license, CERN does not
+# waive the privileges and immunities granted to it by virtue of its status
+# as an Intergovernmental Organization or submit itself to any jurisdiction.
+# from .signals import secret_key_changed
+
+"""Invenio-DB utility functions."""
+
+from flask import current_app
+from sqlalchemy.engine import reflection
+
+from invenio_db import db
+
+
+def rebuild_encrypted_properties(old_key, model, properties):
+    """Rebuild a model's EncryptedType properties when the SECRET_KEY is changed.
+
+    :param old_key: old SECRET_KEY.
+    :param model: the affected db model.
+    :param properties: list of properties to rebuild.
+    """
+    inspector = reflection.Inspector.from_engine(db.engine)
+    primary_key_names = inspector.get_primary_keys(model.__tablename__)
+
+    new_secret_key = current_app.secret_key
+    db.session.expunge_all()
+    try:
+        with db.session.begin_nested():
+            current_app.secret_key = old_key
+            db_columns = []
+            for primary_key in primary_key_names:
+                db_columns.append(getattr(model, primary_key))
+            for prop in properties:
+                db_columns.append(getattr(model, prop))
+            old_rows = db.session.query(*db_columns).all()
+    except Exception as e:
+        current_app.logger.error(
+            'Exception occurred while reading encrypted properties. '
+            'Try again before starting the server with the new secret key.')
+        raise e
+    finally:
+        current_app.secret_key = new_secret_key
+        db.session.expunge_all()
+
+    for old_row in old_rows:
+        primary_keys, old_entries = old_row[:len(primary_key_names)], \
+                                    old_row[len(primary_key_names):]
+        primary_key_fields = dict(zip(primary_key_names, primary_keys))
+        update_values = dict(zip(properties, old_entries))
+        model.query.filter_by(**primary_key_fields).\
+            update(update_values)
+    db.session.commit()

setup.py

@@ -25,7 +25,6 @@
 """Database management for Invenio."""
 
 import os
-import sys
 
 from setuptools import find_packages, setup
 
@@ -35,6 +34,7 @@
 tests_require = [
     'check-manifest>=0.25',
     'coverage>=4.0',
+    'cryptography>=1.5',
     'isort>=4.2.2',
     'mock>=1.3.0',
     'pydocstyle>=1.0.0',

tests/test_utils.py

@@ -0,0 +1,75 @@
+# -*- coding: utf-8 -*-
+#
+# This file is part of Invenio.
+# Copyright (C) 2017 CERN.
+#
+# Invenio is free software; you can redistribute it
+# and/or modify it under the terms of the GNU General Public License as
+# published by the Free Software Foundation; either version 2 of the
+# License, or (at your option) any later version.
+#
+# Invenio is distributed in the hope that it will be
+# useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with Invenio; if not, write to the
+# Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
+# MA 02111-1307, USA.
+#
+# In applying this license, CERN does not
+# waive the privileges and immunities granted to it by virtue of its status
+# as an Intergovernmental Organization or submit itself to any jurisdiction.
+
+"""Test DB utilities."""
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy_utils.types.encrypted import EncryptedType
+
+from invenio_db import InvenioDB
+from invenio_db.utils import rebuild_encrypted_properties
+
+
+def test_rebuild_encrypted_properties(db, app):
+    old_secret_key = "SECRET_KEY_1"
+    new_secret_key = "SECRET_KEY_2"
+    app.secret_key = old_secret_key
+
+    def _secret_key():
+        return app.config.get('SECRET_KEY').encode('utf-8')
+
+    class Demo(db.Model):
+        __tablename__ = 'demo'
+        pk = db.Column(sa.Integer, primary_key=True)
+        et = db.Column(
+            EncryptedType(type_in=db.Unicode, key=_secret_key), nullable=False
+        )
+
+    InvenioDB(app, entry_point_group=False, db=db)
+
+    with app.app_context():
+        db.create_all()
+        d1 = Demo(et="something")
+        db.session.add(d1)
+        db.session.commit()
+
+    app.secret_key = new_secret_key
+
+    with app.app_context():
+        with pytest.raises(UnicodeDecodeError):
+            db.session.query(Demo).all()
+        with pytest.raises(AttributeError):
+            rebuild_encrypted_properties(old_secret_key, Demo, ['nonexistent'])
+        assert app.secret_key == new_secret_key
+
+    with app.app_context():
+        with pytest.raises(UnicodeDecodeError):
+            db.session.query(Demo).all()
+        rebuild_encrypted_properties(old_secret_key, Demo, ['et'])
+        d1_after = db.session.query(Demo).first()
+        assert d1_after.et == "something"
+
+    with app.app_context():
+        db.drop_all()