automatically cull oldest 25% of signatures

when cache size is reached

IPython/nbformat/sign.py

@@ -1,411 +1,424 b''
 """Utilities for signing notebooks"""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import base64
 from contextlib import contextmanager
 from datetime import datetime
 import hashlib
 from hmac import HMAC
 import io
 import os
 
 try:
     import sqlite3
 except ImportError:
     try:
         from pysqlite2 import dbapi2 as sqlite3
     except ImportError:
         sqlite3 = None
 
 from IPython.utils.io import atomic_writing
 from IPython.utils.py3compat import unicode_type, cast_bytes
 from IPython.utils.traitlets import Instance, Bytes, Enum, Any, Unicode, Bool, Integer
 from IPython.config import LoggingConfigurable, MultipleInstanceError
 from IPython.core.application import BaseIPythonApplication, base_flags
 
 from . import read, write, NO_CONVERT
 
 try:
     # Python 3
     algorithms = hashlib.algorithms_guaranteed
 except AttributeError:
     algorithms = hashlib.algorithms
 
 
 def yield_everything(obj):
     """Yield every item in a container as bytes
     
     Allows any JSONable object to be passed to an HMAC digester
     without having to serialize the whole thing.
     """
     if isinstance(obj, dict):
         for key in sorted(obj):
             value = obj[key]
             yield cast_bytes(key)
             for b in yield_everything(value):
                 yield b
     elif isinstance(obj, (list, tuple)):
         for element in obj:
             for b in yield_everything(element):
                 yield b
     elif isinstance(obj, unicode_type):
         yield obj.encode('utf8')
     else:
         yield unicode_type(obj).encode('utf8')
 
 def yield_code_cells(nb):
     """Iterator that yields all cells in a notebook
     
     nbformat version independent
     """
     if nb.nbformat >= 4:
         for cell in nb['cells']:
             if cell['cell_type'] == 'code':
                 yield cell
     elif nb.nbformat == 3:
         for ws in nb['worksheets']:
             for cell in ws['cells']:
                 if cell['cell_type'] == 'code':
                     yield cell
 
 @contextmanager
 def signature_removed(nb):
     """Context manager for operating on a notebook with its signature removed
     
     Used for excluding the previous signature when computing a notebook's signature.
     """
     save_signature = nb['metadata'].pop('signature', None)
     try:
         yield
     finally:
         if save_signature is not None:
             nb['metadata']['signature'] = save_signature
 
 
 class NotebookNotary(LoggingConfigurable):
     """A class for computing and verifying notebook signatures."""
     
     profile_dir = Instance("IPython.core.profiledir.ProfileDir")
     def _profile_dir_default(self):
         from IPython.core.application import BaseIPythonApplication
         app = None
         try:
             if BaseIPythonApplication.initialized():
                 app = BaseIPythonApplication.instance()
         except MultipleInstanceError:
             pass
         if app is None:
             # create an app, without the global instance
             app = BaseIPythonApplication()
             app.initialize(argv=[])
         return app.profile_dir
     
-    db_file = Unicode(config=True)
+    db_file = Unicode(config=True,
+        help="""The sqlite file in which to store notebook signatures.
+        By default, this will be in your IPython profile.
+        You can set it to ':memory:' to disable sqlite writing to the filesystem.
+        """)
     def _db_file_default(self):
         if self.profile_dir is None:
             return ':memory:'
         return os.path.join(self.profile_dir.security_dir, u'nbsignatures.db')
     
     # 64k entries ~ 12MB
-    db_size_limit = Integer(65535, config=True)
+    cache_size = Integer(65535, config=True,
+        help="""The number of notebook signatures to cache.
+        When the number of signatures exceeds this value,
+        the oldest 25% of signatures will be culled.
+        """
+    )
     db = Any()
     def _db_default(self):
         if sqlite3 is None:
             self.log.warn("Missing SQLite3, all notebooks will be untrusted!")
             return
         kwargs = dict(detect_types=sqlite3.PARSE_DECLTYPES|sqlite3.PARSE_COLNAMES)
         db = sqlite3.connect(self.db_file, **kwargs)
         self.init_db(db)
         return db
     
     def init_db(self, db):
         db.execute("""
         CREATE TABLE IF NOT EXISTS nbsignatures
         (
             id integer PRIMARY KEY AUTOINCREMENT,
             algorithm text,
             signature text,
             path text,
             last_seen timestamp
         )""")
         db.execute("""
         CREATE INDEX IF NOT EXISTS algosig ON nbsignatures(algorithm, signature)
         """)
         db.commit()
     
     algorithm = Enum(algorithms, default_value='sha256', config=True,
         help="""The hashing algorithm used to sign notebooks."""
     )
     def _algorithm_changed(self, name, old, new):
         self.digestmod = getattr(hashlib, self.algorithm)
     
     digestmod = Any()
     def _digestmod_default(self):
         return getattr(hashlib, self.algorithm)
     
     secret_file = Unicode(config=True,
         help="""The file where the secret key is stored."""
     )
     def _secret_file_default(self):
         if self.profile_dir is None:
             return ''
         return os.path.join(self.profile_dir.security_dir, 'notebook_secret')
     
     secret = Bytes(config=True,
         help="""The secret key with which notebooks are signed."""
     )
     def _secret_default(self):
         # note : this assumes an Application is running
         if os.path.exists(self.secret_file):
             with io.open(self.secret_file, 'rb') as f:
                 return f.read()
         else:
             secret = base64.encodestring(os.urandom(1024))
             self._write_secret_file(secret)
             return secret
     
     def _write_secret_file(self, secret):
         """write my secret to my secret_file"""
         self.log.info("Writing notebook-signing key to %s", self.secret_file)
         with io.open(self.secret_file, 'wb') as f:
             f.write(secret)
         try:
             os.chmod(self.secret_file, 0o600)
         except OSError:
             self.log.warn(
                 "Could not set permissions on %s",
                 self.secret_file
             )
         return secret
     
     def compute_signature(self, nb):
         """Compute a notebook's signature
         
         by hashing the entire contents of the notebook via HMAC digest.
         """
         hmac = HMAC(self.secret, digestmod=self.digestmod)
         # don't include the previous hash in the content to hash
         with signature_removed(nb):
             # sign the whole thing
             for b in yield_everything(nb):
                 hmac.update(b)
         
         return hmac.hexdigest()
     
     def check_signature(self, nb):
         """Check a notebook's stored signature
         
         If a signature is stored in the notebook's metadata,
         a new signature is computed and compared with the stored value.
         
         Returns True if the signature is found and matches, False otherwise.
         
         The following conditions must all be met for a notebook to be trusted:
         - a signature is stored in the form 'scheme:hexdigest'
         - the stored scheme matches the requested scheme
         - the requested scheme is available from hashlib
         - the computed hash from notebook_signature matches the stored hash
         """
         if nb.nbformat < 3:
             return False
         if self.db is None:
             return False
         signature = self.compute_signature(nb)
         r = self.db.execute("""SELECT id FROM nbsignatures WHERE
             algorithm = ? AND
             signature = ?;
             """, (self.algorithm, signature)).fetchone()
         if r is None:
             return False
         self.db.execute("""UPDATE nbsignatures SET last_seen = ? WHERE
             algorithm = ? AND
             signature = ?;
             """,
             (datetime.utcnow(), self.algorithm, signature),
         )
         self.db.commit()
         return True
     
     def sign(self, nb):
         """Sign a notebook, indicating that its output is trusted on this machine
         
         Stores hash algorithm and hmac digest in a local database of trusted notebooks.
         """
         if nb.nbformat < 3:
             return
         signature = self.compute_signature(nb)
         self.store_signature(signature, nb)
 
     def store_signature(self, signature, nb):
         if self.db is None:
             return
         self.db.execute("""INSERT OR IGNORE INTO nbsignatures
             (algorithm, signature, last_seen) VALUES (?, ?, ?)""",
             (self.algorithm, signature, datetime.utcnow())
         )
         self.db.execute("""UPDATE nbsignatures SET last_seen = ? WHERE
             algorithm = ? AND
             signature = ?;
             """,
             (datetime.utcnow(), self.algorithm, signature),
         )
         self.db.commit()
+        n, = self.db.execute("SELECT Count(*) FROM nbsignatures").fetchone()
+        if n > self.cache_size:
+            self.cull_db()
     
     def unsign(self, nb):
         """Ensure that a notebook is untrusted
         
         by removing its signature from the trusted database, if present.
         """
         signature = self.compute_signature(nb)
         self.db.execute("""DELETE FROM nbsignatures WHERE
                 algorithm = ? AND
                 signature = ?;
             """,
             (self.algorithm, signature)
         )
         self.db.commit()
     
     def cull_db(self):
+        """Cull oldest 25% of the trusted signatures when the size limit is reached"""
         self.db.execute("""DELETE FROM nbsignatures WHERE id IN (
             SELECT id FROM nbsignatures ORDER BY last_seen DESC LIMIT -1 OFFSET ?
         );
-        """, (self.db_size_limit,))
+        """, (max(int(0.75 * self.cache_size), 1),))
     
     def mark_cells(self, nb, trusted):
         """Mark cells as trusted if the notebook's signature can be verified
         
         Sets ``cell.metadata.trusted = True | False`` on all code cells,
         depending on whether the stored signature can be verified.
         
         This function is the inverse of check_cells
         """
         if nb.nbformat < 3:
             return
         
         for cell in yield_code_cells(nb):
             cell['metadata']['trusted'] = trusted
     
     def _check_cell(self, cell, nbformat_version):
         """Do we trust an individual cell?
         
         Return True if:
         
         - cell is explicitly trusted
         - cell has no potentially unsafe rich output
         
         If a cell has no output, or only simple print statements,
         it will always be trusted.
         """
         # explicitly trusted
         if cell['metadata'].pop("trusted", False):
             return True
         
         # explicitly safe output
         if nbformat_version >= 4:
             unsafe_output_types = ['execute_result', 'display_data']
             safe_keys = {"output_type", "execution_count", "metadata"}
         else: # v3
             unsafe_output_types = ['pyout', 'display_data']
             safe_keys = {"output_type", "prompt_number", "metadata"}
         
         for output in cell['outputs']:
             output_type = output['output_type']
             if output_type in unsafe_output_types:
                 # if there are any data keys not in the safe whitelist
                 output_keys = set(output)
                 if output_keys.difference(safe_keys):
                     return False
         
         return True
     
     def check_cells(self, nb):
         """Return whether all code cells are trusted
         
         If there are no code cells, return True.
         
         This function is the inverse of mark_cells.
         """
         if nb.nbformat < 3:
             return False
         trusted = True
         for cell in yield_code_cells(nb):
             # only distrust a cell if it actually has some output to distrust
             if not self._check_cell(cell, nb.nbformat):
                 trusted = False
 
         return trusted
 
 
 trust_flags = {
     'reset' : (
         {'TrustNotebookApp' : { 'reset' : True}},
         """Generate a new key for notebook signature.
         All previously signed notebooks will become untrusted.
         """
     ),
 }
 trust_flags.update(base_flags)
 trust_flags.pop('init')
 
 
 class TrustNotebookApp(BaseIPythonApplication):
     
     description="""Sign one or more IPython notebooks with your key,
     to trust their dynamic (HTML, Javascript) output.
     
     Trusting a notebook only applies to the current IPython profile.
     To trust a notebook for use with a profile other than default,
     add `--profile [profile name]`.
     
     Otherwise, you will have to re-execute the notebook to see output.
     """
     
     examples = """
     ipython trust mynotebook.ipynb and_this_one.ipynb
     ipython trust --profile myprofile mynotebook.ipynb
     """
     
     flags = trust_flags
     
     reset = Bool(False, config=True,
         help="""If True, generate a new key for notebook signature.
         After reset, all previously signed notebooks will become untrusted.
         """
     )
     
     notary = Instance(NotebookNotary)
     def _notary_default(self):
         return NotebookNotary(parent=self, profile_dir=self.profile_dir)
     
     def sign_notebook(self, notebook_path):
         if not os.path.exists(notebook_path):
             self.log.error("Notebook missing: %s" % notebook_path)
             self.exit(1)
         with io.open(notebook_path, encoding='utf8') as f:
             nb = read(f, NO_CONVERT)
         if self.notary.check_signature(nb):
             print("Notebook already signed: %s" % notebook_path)
         else:
             print("Signing notebook: %s" % notebook_path)
             self.notary.sign(nb)
             with atomic_writing(notebook_path) as f:
                 write(nb, f, NO_CONVERT)
     
     def generate_new_key(self):
         """Generate a new notebook signature key"""
         print("Generating new notebook key: %s" % self.notary.secret_file)
         self.notary._write_secret_file(os.urandom(1024))
     
     def start(self):
         if self.reset:
             self.generate_new_key()
             return
         if not self.extra_args:
             self.log.critical("Specify at least one notebook to sign.")
             self.exit(1)
         
         for notebook_path in self.extra_args:
             self.sign_notebook(notebook_path)
 

IPython/nbformat/tests/test_sign.py

@@ -1,200 +1,191 b''
 """Test Notebook signing"""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import copy
 import time
 
 from .base import TestsBase
 
 from IPython.nbformat import read, sign
 from IPython.core.getipython import get_ipython
 
 
 class TestNotary(TestsBase):
     
     def setUp(self):
         self.notary = sign.NotebookNotary(
             secret=b'secret',
             profile_dir=get_ipython().profile_dir,
-            db_url=':memory:'
+            db_file=':memory:'
         )
         with self.fopen(u'test3.ipynb', u'r') as f:
             self.nb = read(f, as_version=4)
         with self.fopen(u'test3.ipynb', u'r') as f:
             self.nb3 = read(f, as_version=3)
     
     def test_algorithms(self):
         last_sig = ''
         for algo in sign.algorithms:
             self.notary.algorithm = algo
             sig = self.notary.compute_signature(self.nb)
             self.assertNotEqual(last_sig, sig)
             last_sig = sig
     
     def test_sign_same(self):
         """Multiple signatures of the same notebook are the same"""
         sig1 = self.notary.compute_signature(self.nb)
         sig2 = self.notary.compute_signature(self.nb)
         self.assertEqual(sig1, sig2)
     
     def test_change_secret(self):
         """Changing the secret changes the signature"""
         sig1 = self.notary.compute_signature(self.nb)
         self.notary.secret = b'different'
         sig2 = self.notary.compute_signature(self.nb)
         self.assertNotEqual(sig1, sig2)
     
     def test_sign(self):
         self.assertFalse(self.notary.check_signature(self.nb))
         self.notary.sign(self.nb)
         self.assertTrue(self.notary.check_signature(self.nb))
     
     def test_unsign(self):
         self.notary.sign(self.nb)
         self.assertTrue(self.notary.check_signature(self.nb))
         self.notary.unsign(self.nb)
         self.assertFalse(self.notary.check_signature(self.nb))
         self.notary.unsign(self.nb)
         self.assertFalse(self.notary.check_signature(self.nb))
     
     def test_cull_db(self):
         # this test has various sleeps of 2ms
         # to ensure low resolution timestamps compare as expected
         dt = 2e-3
         nbs = [
-            copy.deepcopy(self.nb) for i in range(5)
+            copy.deepcopy(self.nb) for i in range(10)
         ]
-        for i, nb in enumerate(nbs):
+        for row in self.notary.db.execute("SELECT * FROM nbsignatures"):
+            print(row)
+        self.notary.cache_size = 8
+        for i, nb in enumerate(nbs[:8]):
             nb.metadata.dirty = i
             self.notary.sign(nb)
         
-        for i, nb in enumerate(nbs):
+        for i, nb in enumerate(nbs[:8]):
             time.sleep(dt)
             self.assertTrue(self.notary.check_signature(nb), 'nb %i is trusted' % i)
         
-        self.notary.db_size_limit = 2
-        self.notary.cull_db()
-        
-        # expect all but last two signatures to be culled
-        self.assertEqual(
-            [self.notary.check_signature(nb) for nb in nbs],
-            [False] * (len(nbs) - 2) + [True] * 2
-        )
-        
-        # sign them all again
-        for nb in nbs:
-            time.sleep(dt)
-            self.notary.sign(nb)
-        
-        # checking front two marks them as newest for next cull instead of oldest
-        time.sleep(dt)
-        self.notary.check_signature(nbs[0])
-        self.notary.check_signature(nbs[1])
-        self.notary.cull_db()
-        
-        self.assertEqual(
-            [self.notary.check_signature(nb) for nb in nbs],
-            [True] * 2 + [False] * (len(nbs) - 2)
-        )
-        
+        # signing the 9th triggers culling of first 3
+        # (75% of 8 = 6, 9 - 6 = 3 culled)
+        self.notary.sign(nbs[8])
+        self.assertFalse(self.notary.check_signature(nbs[0]))
+        self.assertFalse(self.notary.check_signature(nbs[1]))
+        self.assertFalse(self.notary.check_signature(nbs[2]))
+        self.assertTrue(self.notary.check_signature(nbs[3]))
+        # checking nb3 should keep it from being culled:
+        self.notary.sign(nbs[0])
+        self.notary.sign(nbs[1])
+        self.notary.sign(nbs[2])
+        self.assertTrue(self.notary.check_signature(nbs[3]))
+        self.assertFalse(self.notary.check_signature(nbs[4]))
     
     def test_check_signature(self):
         nb = self.nb
         md = nb.metadata
         notary = self.notary
         check_signature = notary.check_signature
         # no signature:
         md.pop('signature', None)
         self.assertFalse(check_signature(nb))
         # hash only, no algo
         md.signature = notary.compute_signature(nb)
         self.assertFalse(check_signature(nb))
         # proper signature, algo mismatch
         notary.algorithm = 'sha224'
         notary.sign(nb)
         notary.algorithm = 'sha256'
         self.assertFalse(check_signature(nb))
         # check correctly signed notebook
         notary.sign(nb)
         self.assertTrue(check_signature(nb))
     
     def test_mark_cells_untrusted(self):
         cells = self.nb.cells
         self.notary.mark_cells(self.nb, False)
         for cell in cells:
             self.assertNotIn('trusted', cell)
             if cell.cell_type == 'code':
                 self.assertIn('trusted', cell.metadata)
                 self.assertFalse(cell.metadata.trusted)
             else:
                 self.assertNotIn('trusted', cell.metadata)
     
     def test_mark_cells_trusted(self):
         cells = self.nb.cells
         self.notary.mark_cells(self.nb, True)
         for cell in cells:
             self.assertNotIn('trusted', cell)
             if cell.cell_type == 'code':
                 self.assertIn('trusted', cell.metadata)
                 self.assertTrue(cell.metadata.trusted)
             else:
                 self.assertNotIn('trusted', cell.metadata)
     
     def test_check_cells(self):
         nb = self.nb
         self.notary.mark_cells(nb, True)
         self.assertTrue(self.notary.check_cells(nb))
         for cell in nb.cells:
             self.assertNotIn('trusted', cell)
         self.notary.mark_cells(nb, False)
         self.assertFalse(self.notary.check_cells(nb))
         for cell in nb.cells:
             self.assertNotIn('trusted', cell)
     
     def test_trust_no_output(self):
         nb = self.nb
         self.notary.mark_cells(nb, False)
         for cell in nb.cells:
             if cell.cell_type == 'code':
                 cell.outputs = []
         self.assertTrue(self.notary.check_cells(nb))
     
     def test_mark_cells_untrusted_v3(self):
         nb = self.nb3
         cells = nb.worksheets[0].cells
         self.notary.mark_cells(nb, False)
         for cell in cells:
             self.assertNotIn('trusted', cell)
             if cell.cell_type == 'code':
                 self.assertIn('trusted', cell.metadata)
                 self.assertFalse(cell.metadata.trusted)
             else:
                 self.assertNotIn('trusted', cell.metadata)
     
     def test_mark_cells_trusted_v3(self):
         nb = self.nb3
         cells = nb.worksheets[0].cells
         self.notary.mark_cells(nb, True)
         for cell in cells:
             self.assertNotIn('trusted', cell)
             if cell.cell_type == 'code':
                 self.assertIn('trusted', cell.metadata)
                 self.assertTrue(cell.metadata.trusted)
             else:
                 self.assertNotIn('trusted', cell.metadata)
     
     def test_check_cells_v3(self):
         nb = self.nb3
         cells = nb.worksheets[0].cells
         self.notary.mark_cells(nb, True)
         self.assertTrue(self.notary.check_cells(nb))
         for cell in cells:
             self.assertNotIn('trusted', cell)
         self.notary.mark_cells(nb, False)
         self.assertFalse(self.notary.check_cells(nb))
         for cell in cells:
             self.assertNotIn('trusted', cell)