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)