def testSuccess(self): # Create two properly-formed events to be returned. file_catalog_1 = bit9_test_utils.CreateFileCatalog( id=100, certificate_id=101) computer_1 = bit9_test_utils.CreateComputer(id=102) cert_1 = bit9_test_utils.CreateCertificate(id=101) event_1 = bit9_test_utils.CreateEvent( id=103, file_catalog_id=100, computer_id=102) event_1 = bit9_test_utils.Expand( event_1, api.Event.file_catalog_id, file_catalog_1) event_1 = bit9_test_utils.Expand( event_1, api.Event.computer_id, computer_1) file_catalog_2 = bit9_test_utils.CreateFileCatalog( id=200, certificate_id=201) computer_2 = bit9_test_utils.CreateComputer(id=202) cert_2 = bit9_test_utils.CreateCertificate(id=201) event_2 = bit9_test_utils.CreateEvent( id=203, file_catalog_id=200, computer_id=202) event_2 = bit9_test_utils.Expand( event_2, api.Event.file_catalog_id, file_catalog_2) event_2 = bit9_test_utils.Expand( event_2, api.Event.computer_id, computer_2) self._AppendMockApiResults([event_1, event_2], cert_2, cert_1) results = sync.GetEvents(0) self.assertEqual(2, len(results)) self.assertListEqual([103, 203], [e.id for e, _ in results]) self.assertListEqual( [[101], [201]], [[c.id for c in sc] for _, sc in results])
def testOtherException(self, mock_get_signing_chain): # Create a properly-formed event that will be returned. file_catalog_1 = bit9_test_utils.CreateFileCatalog( id=100, certificate_id=101) computer_1 = bit9_test_utils.CreateComputer(id=102) cert_1 = bit9_test_utils.CreateCertificate(id=101) signing_chain_1 = [cert_1] event_1 = bit9_test_utils.CreateEvent( id=103, file_catalog_id=100, computer_id=102) event_1 = bit9_test_utils.Expand( event_1, api.Event.file_catalog_id, file_catalog_1) event_1 = bit9_test_utils.Expand( event_1, api.Event.computer_id, computer_1) # Create a second event that will hit an exception. file_catalog_2 = bit9_test_utils.CreateFileCatalog( id=200, certificate_id=201) computer_2 = bit9_test_utils.CreateComputer(id=202) event_2 = bit9_test_utils.CreateEvent( id=203, file_catalog_id=200, computer_id=202) event_2 = bit9_test_utils.Expand( event_2, api.Event.file_catalog_id, file_catalog_2) event_2 = bit9_test_utils.Expand( event_2, api.Event.computer_id, computer_2) # Create another properly-formed event won't be returned. file_catalog_3 = bit9_test_utils.CreateFileCatalog( id=300, certificate_id=301) computer_3 = bit9_test_utils.CreateComputer(id=302) cert_3 = bit9_test_utils.CreateCertificate(id=301) signing_chain_3 = [cert_3] event_3 = bit9_test_utils.CreateEvent( id=303, file_catalog_id=300, computer_id=302) event_3 = bit9_test_utils.Expand( event_3, api.Event.file_catalog_id, file_catalog_3) event_3 = bit9_test_utils.Expand( event_3, api.Event.computer_id, computer_3) self._AppendMockApiResults([event_1, event_2, event_3], cert_3, cert_1) mock_get_signing_chain.side_effect = [ signing_chain_3, Exception, signing_chain_1] results = sync.GetEvents(0) self.assertEqual(2, len(results)) self.assertTrue(sync.monitoring.events_skipped.Increment.called) actual_event_1, actual_signing_chain_1 = results[0] self.assertEqual(1, len(actual_signing_chain_1)) self.assertEqual(103, actual_event_1.id) self.assertEqual(101, actual_signing_chain_1[0].id) actual_event_3, actual_signing_chain_3 = results[1] self.assertEqual(1, len(actual_signing_chain_3)) self.assertEqual(303, actual_event_3.id) self.assertEqual(301, actual_signing_chain_3[0].id)
def testSuccess(self, mock_get_certificate): cert_root = bit9_test_utils.CreateCertificate() cert_intermediate = bit9_test_utils.CreateCertificate( parent_certificate_id=cert_root.id) cert_leaf = bit9_test_utils.CreateCertificate( parent_certificate_id=cert_intermediate.id) expected = [cert_leaf, cert_intermediate, cert_root] mock_get_certificate.side_effect = expected actual = sync._GetSigningChain(cert_leaf.id) self.assertListEqual(expected, actual)
def testWithSigningChain(self): signing_chain = [ bit9_test_utils.CreateCertificate(thumbprint=test_utils.RandomSHA1()) for _ in xrange(4)] bit9_test_utils.LinkSigningChain(*signing_chain) expected_key = ndb.Key( bit9_models.Bit9Certificate, signing_chain[0].thumbprint) self.assertEqual(expected_key, sync._GetCertKey(signing_chain))
def testMalformed_SuccessfulRetry(self, mock_get): bad_cert = bit9_test_utils.CreateCertificate( thumbprint=None, valid_to=None) good_cert = bit9_test_utils.CreateCertificate() mock_get.side_effect = [bad_cert, good_cert] cert_id = 12345 memcache_key = sync._CERT_MEMCACHE_KEY % cert_id self.assertIsNone(memcache.get(memcache_key)) actual_cert = sync._GetCertificate(cert_id) self.assertEqual(good_cert, actual_cert) self.assertEqual(2, mock_get.call_count) # Verify that the cert is present in memcache. cached_cert = memcache.get(memcache_key) self.assertEqual(good_cert, cached_cert)
def testDupeCerts(self): # Create some cert entities, and a matching protobuf signing chain. bit9_certs = test_utils.CreateBit9Certificates(3) thumbprints = [c.key.id() for c in bit9_certs] signing_chain = [ bit9_test_utils.CreateCertificate(thumbprint=t) for t in thumbprints] bit9_test_utils.LinkSigningChain(*signing_chain) self.assertEntityCount(bit9_models.Bit9Certificate, 3) sync._PersistBit9Certificates(signing_chain).wait() self.assertEntityCount(bit9_models.Bit9Certificate, 3)
def _CreateEventsAndCerts( count=1, event_kwargs=None, file_catalog_kwargs=None, computer_kwargs=None): event_kwargs = event_kwargs or {} file_catalog_kwargs = file_catalog_kwargs or {} computer_kwargs = computer_kwargs or {} # Create a generator for each type of ID, with each range starting where the # previous one left off. id_gens = itertools.izip( xrange(100 + (count * 0), 100 + (count * 1)), xrange(100 + (count * 1), 100 + (count * 2)), xrange(100 + (count * 2), 100 + (count * 3)), xrange(100 + (count * 3), 100 + (count * 4))) events = [] certs = [] for event_id, file_catalog_id, computer_id, certificate_id in id_gens: # Construct the Certificate. cert = bit9_test_utils.CreateCertificate(id=certificate_id) # Construct the Computer. computer_id = computer_kwargs.get('id', computer_id) computer_defaults = {'id': computer_id} computer_defaults.update(computer_kwargs.copy()) computer = bit9_test_utils.CreateComputer(**computer_defaults) # Construct the FileCatalog. file_catalog_id = file_catalog_kwargs.get('id', file_catalog_id) file_catalog_defaults = { 'id': file_catalog_id, 'certificate_id': certificate_id} file_catalog_defaults.update(file_catalog_kwargs.copy()) file_catalog = bit9_test_utils.CreateFileCatalog(**file_catalog_defaults) # Construct the Event. event_defaults = { 'id': event_id, 'file_catalog_id': file_catalog_id, 'computer_id': computer_id} event_defaults.update(event_kwargs.copy()) event = bit9_test_utils.CreateEvent(**event_defaults) event = bit9_test_utils.Expand( event, api.Event.file_catalog_id, file_catalog) event = bit9_test_utils.Expand(event, api.Event.computer_id, computer) events.append(event) # Stuff the certs in backwards due to the reverse sorting in GetEvents(). certs.insert(0, cert) return events, certs
def testNewCerts(self): # Create some certs, and an unrelated signing chain. test_utils.CreateBit9Certificates(3) signing_chain = [ bit9_test_utils.CreateCertificate(thumbprint=test_utils.RandomSHA1()) for _ in xrange(4)] bit9_test_utils.LinkSigningChain(*signing_chain) self.assertEntityCount(bit9_db.Bit9Certificate, 3) sync._PersistBit9Certificates(signing_chain).wait() self.assertEntityCount(bit9_db.Bit9Certificate, 7)
def testFileCatalogMissing(self): # Simulate an event with a missing fileCatalog. computer = bit9_test_utils.CreateComputer(id=100) signing_chain = [bit9_test_utils.CreateCertificate(id=101)] event = bit9_test_utils.CreateEvent( id=102, computer_id=100, file_catalog_id=103) event = bit9_test_utils.Expand(event, api.Event.computer_id, computer) self._AppendMockApiResults(event, signing_chain) results = sync.GetEvents(0) self.assertEqual(0, len(results)) self.assertTrue(sync.monitoring.events_skipped.Increment.called)
def testMalformed_UnsuccessfulRetries(self, mock_get): bad_cert = bit9_test_utils.CreateCertificate( thumbprint=None, valid_to=None) mock_get.side_effect = [bad_cert] * sync._GET_CERT_ATTEMPTS cert_id = 12345 memcache_key = sync._CERT_MEMCACHE_KEY % cert_id self.assertIsNone(memcache.get(memcache_key)) with self.assertRaises(sync.MalformedCertificate): sync._GetCertificate(cert_id) self.assertIsNone(memcache.get(memcache_key))
def testNewCerts(self): # Create some certs, and an unrelated signing chain. test_utils.CreateBit9Certificates(3) signing_chain = [ bit9_test_utils.CreateCertificate(thumbprint=test_utils.RandomSHA1()) for _ in xrange(4)] bit9_test_utils.LinkSigningChain(*signing_chain) self.assertEntityCount(bit9_models.Bit9Certificate, 3) sync._PersistBit9Certificates(signing_chain).wait() self.assertEntityCount(bit9_models.Bit9Certificate, 7) self.assertBigQueryInsertions( [constants.BIGQUERY_TABLE.CERTIFICATE] * len(signing_chain))
def testFileCatalogMalformed(self): # Simulate an event with a malformed fileCatalog (in this case, no SHA256). file_catalog = bit9_test_utils.CreateFileCatalog( id=100, certificate_id=101, sha256=None) computer = bit9_test_utils.CreateComputer(id=102) signing_chain = [bit9_test_utils.CreateCertificate(id=101)] event = bit9_test_utils.CreateEvent( id=103, file_catalog_id=100, computer_id=102) event = bit9_test_utils.Expand( event, api.Event.file_catalog_id, file_catalog) event = bit9_test_utils.Expand(event, api.Event.computer_id, computer) self._AppendMockApiResults(event, signing_chain) results = sync.GetEvents(0) self.assertEqual(0, len(results)) self.assertTrue(sync.monitoring.events_skipped.Increment.called)
def _CreateEventTuple(computer=None, file_catalog=None, signing_chain=None, **event_kwargs): if computer is None: computer = bit9_test_utils.CreateComputer() if file_catalog is None: file_catalog = bit9_test_utils.CreateFileCatalog() if signing_chain is None: signing_chain = [bit9_test_utils.CreateCertificate()] event = bit9_test_utils.CreateEvent( computer_id=computer.id, file_catalog_id=file_catalog.id, **event_kwargs) event = bit9_test_utils.Expand(event, api.Event.file_catalog_id, file_catalog) event = bit9_test_utils.Expand(event, api.Event.computer_id, computer) return event, signing_chain
def testSuccess(self, mock_get): expected_cert = bit9_test_utils.CreateCertificate() mock_get.return_value = expected_cert # The key shouldn't initially be in memcache. cert_id = 12345 memcache_key = sync._CERT_MEMCACHE_KEY % cert_id self.assertIsNone(memcache.get(memcache_key)) # The first call should actually trigger an API query. actual_cert = sync._GetCertificate(cert_id) self.assertEqual(expected_cert, actual_cert) self.assertEqual(1, mock_get.call_count) mock_get.reset_mock() # Verify that the cert is present in memcache. cached_cert = memcache.get(memcache_key) self.assertEqual(expected_cert, cached_cert) # Additional calls shouldn't hit the API. actual_cert = sync._GetCertificate(cert_id) self.assertEqual(expected_cert, actual_cert) self.assertEqual(0, mock_get.call_count)