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_LocalAdmin(self):
computer = bit9_test_utils.CreateComputer(
users='{0}\\foobar,{0}\\bazqux'.format(settings.AD_DOMAIN))
event, signing_chain = _CreateEventTuple(
user_name=constants.LOCAL_ADMIN.WINDOWS, computer=computer)
file_catalog = event.get_expand(api.Event.file_catalog_id)
self.assertEntityCount(bit9_db.Bit9Event, 0)
sync._PersistBit9Events(
event, file_catalog, computer, signing_chain).wait()
self.assertEntityCount(bit9_db.Bit9Event, 2)
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 testNewHost(self):
users = test_utils.RandomStrings(2)
host = bit9_test_utils.CreateComputer(
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
occurred_dt = datetime.datetime.utcnow()
self.assertEntityCount(bit9_models.Bit9Host, 0)
sync._PersistBit9Host(host, occurred_dt).wait()
self.assertEntityCount(bit9_models.Bit9Host, 1)
self.assertBigQueryInsertions(
[constants.BIGQUERY_TABLE.HOST] + [constants.BIGQUERY_TABLE.USER] * 2)
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 testSuccess_LocalAdmin(self):
users = '{0}\\foobar,{0}\\bazqux'.format(settings.AD_DOMAIN)
computer = bit9_test_utils.CreateComputer(users=users)
event_kwargs = {'user_name': constants.LOCAL_ADMIN.WINDOWS}
computer_kwargs = {'users': users}
event, cert = _CreateEventAndCert(
event_kwargs=event_kwargs, computer_kwargs=computer_kwargs)
file_catalog = event.get_expand(api.Event.file_catalog_id)
self.assertEntityCount(bit9_models.Bit9Event, 0)
sync._PersistBit9Events(
event, file_catalog, computer, [cert]).wait()
self.assertEntityCount(bit9_models.Bit9Event, 2)
self.assertBigQueryInsertions([constants.BIGQUERY_TABLE.EXECUTION])
def testNewHost(self):
users = test_utils.RandomStrings(2)
host = bit9_test_utils.CreateComputer(
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
occurred_dt = datetime.datetime.utcnow()
self.assertEntityCount(bit9_db.Bit9Host, 0)
self.assertEntityCount(bigquery_db.HostRow, 0)
self.assertEntityCount(bigquery_db.UserRow, 0)
sync._PersistBit9Host(host, occurred_dt).wait()
self.assertEntityCount(bit9_db.Bit9Host, 1)
self.assertTaskCount(constants.TASK_QUEUE.BQ_PERSISTENCE, 3)
self.DrainTaskQueue(constants.TASK_QUEUE.BQ_PERSISTENCE)
self.assertEntityCount(bigquery_db.HostRow, 1)
self.assertEntityCount(bigquery_db.UserRow, 2)
def testUpdateHostUsers_NoIncomingUsers(self):
old_users = test_utils.RandomStrings(2)
policy_key = ndb.Key(bit9_models.Bit9Policy, '22222')
test_utils.CreateBit9Host(
id='12345', users=old_users, policy_key=policy_key)
host = bit9_test_utils.CreateComputer(
id=12345, policy_id=22222, users='')
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
# Verify that the users weren't updated in Datastore.
host = bit9_models.Bit9Host.get_by_id('12345')
self.assertSameElements(old_users, host.users)
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 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 testUpdatePolicyKey(self):
users = test_utils.RandomStrings(2)
old_policy_key = ndb.Key(bit9_models.Bit9Policy, '22222')
test_utils.CreateBit9Host(
id='11111', policy_key=old_policy_key, users=users)
host = bit9_test_utils.CreateComputer(
id=11111,
policy_id=33333,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
bit9_host = bit9_models.Bit9Host.get_by_id('11111')
new_policy_key = ndb.Key(bit9_models.Bit9Policy, '33333')
self.assertEqual(new_policy_key, bit9_host.policy_key)
self.assertBigQueryInsertions([constants.BIGQUERY_TABLE.HOST])
def testUpdateHostname(self):
users = test_utils.RandomStrings(2)
policy_key = ndb.Key(bit9_models.Bit9Policy, '100')
test_utils.CreateBit9Host(
id='12345', hostname=bit9_utils.ExpandHostname('hostname1'),
users=users, policy_key=policy_key)
host = bit9_test_utils.CreateComputer(
id=12345,
name='hostname2',
policy_id=100,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
bit9_host = bit9_models.Bit9Host.get_by_id('12345')
self.assertEqual(bit9_utils.ExpandHostname('hostname2'), bit9_host.hostname)
self.assertNoBigQueryInsertions()
def testUpdateHostname(self):
users = test_utils.RandomStrings(2)
policy_key = ndb.Key(bit9_db.Bit9Policy, '100')
test_utils.CreateBit9Host(
id='12345', hostname=utils.ExpandHostname('hostname1'), users=users,
policy_key=policy_key)
host = bit9_test_utils.CreateComputer(
id=12345,
name='hostname2',
policy_id=100,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
bit9_host = bit9_db.Bit9Host.get_by_id('12345')
self.assertEqual(utils.ExpandHostname('hostname2'), bit9_host.hostname)
self.assertTaskCount(constants.TASK_QUEUE.BQ_PERSISTENCE, 0)
def testUpdateLastEventDt(self):
now_dt = datetime.datetime.utcnow()
earlier_dt = now_dt - datetime.timedelta(days=7)
users = test_utils.RandomStrings(2)
policy_key = ndb.Key(bit9_models.Bit9Policy, '100')
test_utils.CreateBit9Host(
id='12345', last_event_dt=earlier_dt, users=users,
policy_key=policy_key)
host = bit9_test_utils.CreateComputer(
id=12345,
policy_id=100,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
sync._PersistBit9Host(host, now_dt).wait()
bit9_host = bit9_models.Bit9Host.get_by_id('12345')
self.assertEqual(now_dt, bit9_host.last_event_dt)
self.assertNoBigQueryInsertions()
def testUpdatePolicyKey(self):
users = test_utils.RandomStrings(2)
old_policy_key = ndb.Key(bit9_db.Bit9Policy, '22222')
test_utils.CreateBit9Host(
id='11111', policy_key=old_policy_key, users=users)
host = bit9_test_utils.CreateComputer(
id=11111,
policy_id=33333,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *users))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
bit9_host = bit9_db.Bit9Host.get_by_id('11111')
new_policy_key = ndb.Key(bit9_db.Bit9Policy, '33333')
self.assertEqual(new_policy_key, bit9_host.policy_key)
self.assertTaskCount(constants.TASK_QUEUE.BQ_PERSISTENCE, 1)
self.RunDeferredTasks(constants.TASK_QUEUE.BQ_PERSISTENCE)
self.assertEntityCount(bigquery_db.HostRow, 1)
def testUpdateHostUsers(self):
old_users = test_utils.RandomStrings(2)
new_users = test_utils.RandomStrings(2)
policy_key = ndb.Key(bit9_db.Bit9Policy, '22222')
test_utils.CreateBit9Host(
id='12345', users=old_users, policy_key=policy_key)
host = bit9_test_utils.CreateComputer(
id=12345,
policy_id=22222,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *new_users))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
bit9_db.Bit9Host.get_by_id('12345')
self.assertTaskCount(constants.TASK_QUEUE.BQ_PERSISTENCE, 3)
self.DrainTaskQueue(constants.TASK_QUEUE.BQ_PERSISTENCE)
self.assertEntityCount(bigquery_db.HostRow, 1)
self.assertEntityCount(bigquery_db.UserRow, 2)
def testUpdateHostUsers_WindowManager(self):
test_utils.CreateBit9Host(
id='12345',
users=['a_real_person'],
policy_key=ndb.Key(bit9_models.Bit9Policy, '22222'))
host = bit9_test_utils.CreateComputer(
id=12345,
policy_id=22222,
users=r'Window Manager\DWM-999')
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
# Verify that the user didn't get updated.
host = bit9_models.Bit9Host.get_by_id('12345')
self.assertSameElements(['a_real_person'], host.users)
# Verify no BigQuery row persistence.
self.assertNoBigQueryInsertions()
def testCopyLocalRules_NoPreviousHosts(self):
old_user = test_utils.CreateUser(email=user_map.UsernameToEmail('foo'))
new_user = test_utils.CreateUser(email=user_map.UsernameToEmail('bar'))
test_utils.CreateBit9Host(
id='12345', users=[old_user.nickname],
policy_key=ndb.Key(bit9_db.Bit9Policy, '22222'))
host = bit9_test_utils.CreateComputer(
id=12345,
policy_id=22222,
users='{0}\\{1},{0}\\{2}'.format(
settings.AD_DOMAIN, old_user.nickname, new_user.nickname))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
self.assertEntityCount(bit9_db.Bit9Rule, 0)
self.assertEntityCount(bit9_db.RuleChangeSet, 0)
self.assertTaskCount(constants.TASK_QUEUE.BQ_PERSISTENCE, 1)
self.DrainTaskQueue(constants.TASK_QUEUE.BQ_PERSISTENCE)
self.assertEntityCount(bigquery_db.HostRow, 1)
def _CreateUnsyncedEvents(self, host_count=1, events_per_host=-1):
"""Creates a bunch of _UnsycnedEvents across a number of Windows hosts.
Args:
host_count: The number of hosts to create _UnsyncedEvents for.
events_per_host: The number of _UnsyncedEvents to create per host. If set
to -1 (default), creates a random number of _UnsyncedEvents.
Returns:
A sorted list of the randomly generated host IDs.
"""
hosts = [
bit9_test_utils.CreateComputer(id=host_id)
for host_id in xrange(host_count)]
for host in hosts:
if events_per_host == -1:
events_per_host = random.randint(1, 5)
for _ in xrange(events_per_host):
event, _ = _CreateEventTuple(computer=host)
sync._UnsyncedEvent.Generate(event, []).put()
return sorted(host.id for host in hosts)
def testUpdateHostUsers_People(self):
old_users = test_utils.RandomStrings(2)
new_users = test_utils.RandomStrings(2)
policy_key = ndb.Key(bit9_models.Bit9Policy, '22222')
test_utils.CreateBit9Host(
id='12345', users=old_users, policy_key=policy_key)
host = bit9_test_utils.CreateComputer(
id=12345,
policy_id=22222,
users='{0}\\{1},{0}\\{2}'.format(settings.AD_DOMAIN, *new_users))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
# Verify that the users were updated in Datastore.
host = bit9_models.Bit9Host.get_by_id('12345')
self.assertSameElements(new_users, host.users)
# Verify all BigQuery row persistence.
self.assertBigQueryInsertions(
[constants.BIGQUERY_TABLE.HOST] + [constants.BIGQUERY_TABLE.USER] * 2)
def testCopyLocalRules_Success(self):
old_user = test_utils.CreateUser(email=user_map.UsernameToEmail('foo'))
new_user = test_utils.CreateUser(email=user_map.UsernameToEmail('bar'))
policy_key = ndb.Key(bit9_db.Bit9Policy, '22222')
host1 = test_utils.CreateBit9Host(
id='12345', users=[old_user.nickname], policy_key=policy_key)
test_utils.CreateBit9Host(
id='67890', users=[new_user.nickname], policy_key=policy_key)
blockable1 = test_utils.CreateBit9Binary()
test_utils.CreateBit9Rule(
blockable1.key, host_id=host1.key.id(), user_key=old_user.key)
blockable2 = test_utils.CreateBit9Binary()
test_utils.CreateBit9Rule(
blockable2.key, host_id=host1.key.id(), user_key=old_user.key)
host = bit9_test_utils.CreateComputer(
id=67890,
policy_id=22222,
users='{0}\\{1},{0}\\{2}'.format(
settings.AD_DOMAIN, old_user.nickname, new_user.nickname))
occurred_dt = datetime.datetime.utcnow()
sync._PersistBit9Host(host, occurred_dt).wait()
self.assertEntityCount(bit9_db.Bit9Rule, 4) # 2 New + 2 Old
self.assertEntityCount(bit9_db.RuleChangeSet, 2)
rules_for_host1 = bit9_db.Bit9Rule.query(
bit9_db.Bit9Rule.host_id == host1.key.id()).fetch()
self.assertEqual(2, len(rules_for_host1))
self.assertSameElements(
[blockable1.key, blockable2.key],
[rule.key.parent() for rule in rules_for_host1])
self.assertTaskCount(constants.TASK_QUEUE.BQ_PERSISTENCE, 1)
self.DrainTaskQueue(constants.TASK_QUEUE.BQ_PERSISTENCE)
self.assertEntityCount(bigquery_db.HostRow, 1)