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)
