def testDedupeMultiple(self):
keys = self.event_1.GetKeysToInsert(['foo'], [])
event1 = utils.CopyEntity(
self.event_1,
new_key=keys[0],
first_blocked_dt=self.middle,
last_blocked_dt=self.middle)
event2 = utils.CopyEntity(
self.event_1,
new_key=keys[0],
first_blocked_dt=self.earlier,
last_blocked_dt=self.later)
event3 = utils.CopyEntity(
self.event_1,
new_key=keys[0],
first_blocked_dt=self.middle,
last_blocked_dt=self.later)
events = base.Event.DedupeMultiple([event1, event2, event3])
self.assertEqual(1, len(events))
self.assertEqual(self.earlier, events[0].first_blocked_dt)
self.assertEqual(self.later, events[0].last_blocked_dt)
self.assertEqual(3, events[0].count)
def testGetAssociatedUsers(self):
host_id_1 = 'kiwi'
host_id_2 = 'narwal'
event_1 = utils.CopyEntity(self.santa_event,
host_id=host_id_1,
executing_user='******')
event_2 = utils.CopyEntity(self.santa_event,
host_id=host_id_2,
executing_user='******')
event_1.put()
event_2.put()
expected_users_1 = [event_1.executing_user]
expected_users_2 = [event_2.executing_user]
actual_users_1 = santa.SantaHost.GetAssociatedUsers(host_id_1)
actual_users_2 = santa.SantaHost.GetAssociatedUsers(host_id_2)
event_2.host_id = host_id_1
event_2.put()
users_combined = set(santa.SantaHost.GetAssociatedUsers(host_id_1))
self.assertEqual(expected_users_1, actual_users_1)
self.assertEqual(expected_users_2, actual_users_2)
self.assertTrue(
set(expected_users_1 + expected_users_2).issubset(users_combined))
def testDedupeSantaEvent_AddNewerQuarantineData(self):
new_quarantine = utils.CopyEntity(self.santa_event.quarantine,
data_url='http://3vil.com')
later_dt = (self.santa_event.last_blocked_dt +
datetime.timedelta(seconds=1))
new_event = utils.CopyEntity(self.santa_event,
quarantine=new_quarantine,
last_blocked_dt=later_dt)
self.santa_event.Dedupe(new_event)
self.assertEqual('http://3vil.com',
self.santa_event.quarantine.data_url)
def testIdWithKey(self):
inst = self.default_model(a='abc')
inst.put()
with self.assertRaises(datastore_errors.BadArgumentError):
utils.CopyEntity(
inst, new_key=ndb.Key('A', 'a_key'), id='an_id')
def _CreateNewLocalRules(self, uuid, user_key):
"""Creates copies of all local rules for the new host."""
# Pick any host owned by the user to copy rules from. Exclude hosts that
# haven't completed a full sync because they won't have a complete rule set.
# NOTE: Because we expect all hosts owned by a user to have the same local
# rules, we should get the same rules set with any one of the user's hosts.
username = user_map.EmailToUsername(user_key.id())
host_query = santa_db.SantaHost.query(
santa_db.SantaHost.primary_user == username,
santa_db.SantaHost.last_postflight_dt != None) # pylint: disable=g-equals-none
a_host = host_query.get()
if a_host is None:
return utils.GetNoOpFuture()
# Get all local rules from that host.
rules_query = santa_db.SantaRule.query(
santa_db.SantaRule.host_id == a_host.key.id(),
santa_db.SantaRule.in_effect == True) # pylint: disable=g-explicit-bool-comparison
# Copy the local rules to the new host.
new_rules = []
for batch in query_utils.Paginate(rules_query):
for rule in batch:
new_rule = utils.CopyEntity(
rule, new_parent=rule.key.parent(), host_id=uuid, user_key=user_key)
new_rules.append(new_rule)
futures = ndb.put_multi_async(new_rules)
return utils.GetMultiFuture(futures)
def testParentWithKey(self):
inst = self.default_model(a='abc')
inst.put()
parent = ndb.Key('C', 'c', 'B', 'b')
with self.assertRaises(datastore_errors.BadArgumentError):
utils.CopyEntity(inst, new_key=ndb.Key('A', 'a_key'), new_parent=parent)
def testFailToSet_AutoNowProperty(self):
class A(ndb.Model):
a = ndb.DateTimeProperty(auto_now=True)
inst = A()
inst.put()
with self.assertRaises(datastore_utils.PropertyError):
datastore_utils.CopyEntity(inst, a=datetime.datetime.utcnow())
def testDedupe_OutOfOrder(self):
earlier_dt = self.bit9_event.last_blocked_dt - datetime.timedelta(hours=1)
earlier_bit9_event = datastore_utils.CopyEntity(
self.bit9_event,
first_blocked_dt=earlier_dt,
last_blocked_dt=earlier_dt,
bit9_id=self.bit9_event.bit9_id + 1, # Earlier event has larger ID
)
# Always choose the larger ID.
more_recent_deduped = datastore_utils.CopyEntity(earlier_bit9_event)
more_recent_deduped.Dedupe(self.bit9_event)
self.assertEquals(self.bit9_event.bit9_id + 1, more_recent_deduped.bit9_id)
earlier_deduped = datastore_utils.CopyEntity(self.bit9_event)
earlier_deduped.Dedupe(earlier_bit9_event)
self.assertEquals(self.bit9_event.bit9_id + 1, earlier_deduped.bit9_id)
def testNewId(self):
inst = self.default_model(a='abc')
inst.put()
new = utils.CopyEntity(inst, id='an_id')
new.put()
self.assertEqual('abc', new.a)
self.assertEqual('an_id', new.key.id())
def testUpdateProperties(self):
inst = self.default_model(a='abc')
inst.put()
new = utils.CopyEntity(inst, a='xyz')
new.put()
self.assertEqual('abc', inst.a)
self.assertEqual('xyz', new.a)
self.assertNotEqual(new.key, inst.key)
def testDedupe(self):
earlier_dt = self.bit9_event.last_blocked_dt - datetime.timedelta(
hours=1)
earlier_bit9_event = utils.CopyEntity(
self.bit9_event,
first_blocked_dt=earlier_dt,
last_blocked_dt=earlier_dt,
bit9_id=self.bit9_event.bit9_id - 1,
)
# Always choose the larger ID.
more_recent_deduped = utils.CopyEntity(earlier_bit9_event)
more_recent_deduped.Dedupe(self.bit9_event)
self.assertEquals(self.bit9_event.bit9_id, more_recent_deduped.bit9_id)
earlier_deduped = utils.CopyEntity(self.bit9_event)
earlier_deduped.Dedupe(earlier_bit9_event)
self.assertEquals(self.bit9_event.bit9_id, earlier_deduped.bit9_id)
def testGetKeysToInsert_Admin(self):
usernames = ['foo', 'bar']
with mock.patch.object(
base.Event, 'run_by_local_admin', return_value=True):
event = utils.CopyEntity(self.event_1)
keys = event.GetKeysToInsert(usernames, [])
self.assertEqual(2, len(keys))
key_usernames = [user_map.EmailToUsername(key.flat()[1]) for key in keys]
self.assertSameElements(usernames, key_usernames)
def testGetEventKeysToInsert_Admin(self):
usernames = ['foo', 'bar']
with mock.patch.object(
base_models.Event, 'run_by_local_admin', return_value=True):
event = datastore_utils.CopyEntity(self.event)
keys = model_utils.GetEventKeysToInsert(event, usernames, [])
self.assertLen(keys, 2)
key_usernames = [user_utils.EmailToUsername(key.flat()[1]) for key in keys]
self.assertSameElements(usernames, key_usernames)
def testNewIdWithParent(self):
inst = self.default_model(a='abc')
inst.put()
parent = ndb.Key('C', 'c', 'B', 'b')
expected = ndb.Key('C', 'c', 'B', 'b', 'A', 'an_id')
new = datastore_utils.CopyEntity(inst, new_parent=parent, id='an_id')
new.put()
self.assertEqual(expected, new.key)
def _CopyLocalRules(user_key, dest_host_id):
"""Copy over a user's local rules to a newly-associated host.
NOTE: Because of the implementation of local whitelisting on Bit9, many of
these new copied local rules will likely be initially unfulfilled, that is,
held in Upvote and not saved to Bit9.
Args:
user_key: str, The user for whom the rules will be copied.
dest_host_id: str, The ID of the host for which the new rules will be
created.
"""
logging.info('Copying rules for %s to host %s', user_key.id(),
dest_host_id)
username = user_map.EmailToUsername(user_key.id())
host_query = bit9.Bit9Host.query(bit9.Bit9Host.users == username)
src_host = yield host_query.get_async()
if src_host is None:
raise ndb.Return()
assert src_host.key.id() != dest_host_id, (
'User already associated with target host')
# Get all local rules from that host.
rules_query = bit9.Bit9Rule.query(
bit9.Bit9Rule.host_id == src_host.key.id(),
bit9.Bit9Rule.in_effect == True) # pylint: disable=g-explicit-bool-comparison
# Get a rough idea of how many rules we're in for. Since this is a
# non-critical query, we limit the max number to a fairly low bound.
rule_count = yield rules_query.count_async(limit=250)
logging.info('Retrieved %s%s rules to copy',
'>' if rule_count == 250 else '', rule_count)
# Copy the local rules to the new host.
new_rules = []
for batch in query_utils.Paginate(rules_query):
for rule in batch:
new_rule = model_utils.CopyEntity(rule,
new_parent=rule.key.parent(),
host_id=dest_host_id,
user_key=user_key)
new_rules.append(new_rule)
logging.info('Copying %s rules to new host', len(new_rules))
yield ndb.put_multi_async(new_rules)
# Create the change sets necessary to submit the new rules to Bit9.
changes = []
for new_rule in new_rules:
change = bit9.RuleChangeSet(rule_keys=[new_rule.key],
change_type=new_rule.policy,
parent=new_rule.key.parent())
changes.append(change)
logging.info('Creating %s RuleChangeSet', len(changes))
yield ndb.put_multi_async(changes)
def _CopyLocalRules(user_key, dest_host_id):
"""Copy over a user's local rules to a newly-associated host.
NOTE: Because of the implementation of local whitelisting on Bit9, many of
these new copied local rules will likely be initially unfulfilled, that is,
held in Upvote and not saved to Bit9.
Args:
user_key: str, The user for whom the rules will be copied.
dest_host_id: str, The ID of the host for which the new rules will be
created.
"""
logging.info('Copying rules for user %s to host %s', user_key.id(),
dest_host_id)
# Query for a host belonging to the user.
username = user_utils.EmailToUsername(user_key.id())
query = host_models.Bit9Host.query(host_models.Bit9Host.users == username)
src_host = yield query.get_async()
if src_host is None:
logging.warning('User %s has no hosts to copy from', username)
raise ndb.Return()
src_host_id = src_host.key.id()
# Query for all the Bit9Rules in effect for the given user on the chosen host.
query = rule_models.Bit9Rule.query(
rule_models.Bit9Rule.host_id == src_host_id,
rule_models.Bit9Rule.user_key == user_key,
rule_models.Bit9Rule.in_effect == True) # pylint: disable=g-explicit-bool-comparison, singleton-comparison
src_rules = yield query.fetch_async()
logging.info('Found a total of %d rule(s) for user %s', len(src_rules),
user_key.id())
# Copy the local rules to the new host.
logging.info('Copying %d rule(s) to host %s', len(src_rules), dest_host_id)
new_rules = []
for src_rule in src_rules:
new_rule = datastore_utils.CopyEntity(src_rule,
new_parent=src_rule.key.parent(),
host_id=dest_host_id,
user_key=user_key)
new_rules.append(new_rule)
new_rule.InsertBigQueryRow()
yield ndb.put_multi_async(new_rules)
# Create the change sets necessary to submit the new rules to Bit9.
changes = []
for new_rule in new_rules:
change = bit9.RuleChangeSet(rule_keys=[new_rule.key],
change_type=new_rule.policy,
parent=new_rule.key.parent())
changes.append(change)
logging.info('Creating %d RuleChangeSet(s)', len(changes))
yield ndb.put_multi_async(changes)
def testDeletedProperty(self):
inst = self.default_model(a='abc')
inst.put()
class A(ndb.Model): # pylint: disable=unused-variable
b = ndb.StringProperty()
inst = inst.key.get(use_cache=False)
copy = utils.CopyEntity(inst)
self.assertFalse(hasattr(copy, 'a'))
def testAdminGetList_OnlyInEffect(self):
inactive_key = ndb.Key(flat=self.vote_1.key.flat()[:-1] + (None, ))
inactive_vote = utils.CopyEntity(self.vote_1, new_key=inactive_key)
inactive_vote.put()
self.assertEqual(4, base_db.Vote.query().count())
with self.LoggedInUser(admin=True):
response = self.testapp.get('')
self.assertEqual(len(response.json['content']), 3)
def testFailToSet_ComputedProperty(self):
class A(ndb.Model):
a = ndb.StringProperty()
b = ndb.ComputedProperty(lambda self: self.a[0])
inst = A(a='xyz')
inst.put()
self.assertEqual('x', inst.b)
with self.assertRaises(utils.PropertyError):
utils.CopyEntity(inst, b='a')
def testUserGet(self):
"""Normal user reads a vote."""
inactive_key = ndb.Key(flat=self.vote_2.key.flat()[:-1] + (None, ))
inactive_vote = utils.CopyEntity(self.vote_2, new_key=inactive_key)
inactive_vote.put()
self.assertFalse(inactive_vote.in_effect)
with self.LoggedInUser(email_addr=self.user_1.email):
response = self.testapp.get('/%s' % self.other_blockable.key.id())
self.assertEqual(self.vote_2.key.urlsafe(), response.json['key'])
def testDedupeMultiple(self):
keys = model_utils.GetEventKeysToInsert(self.event_1, ['foo'], [])
event1 = datastore_utils.CopyEntity(self.event_1,
new_key=keys[0],
first_blocked_dt=self.middle,
last_blocked_dt=self.middle)
event2 = datastore_utils.CopyEntity(self.event_1,
new_key=keys[0],
first_blocked_dt=self.earlier,
last_blocked_dt=self.later)
event3 = datastore_utils.CopyEntity(self.event_1,
new_key=keys[0],
first_blocked_dt=self.middle,
last_blocked_dt=self.later)
events = event_models.Event.DedupeMultiple([event1, event2, event3])
self.assertLen(events, 1)
self.assertEqual(self.earlier, events[0].first_blocked_dt)
self.assertEqual(self.later, events[0].last_blocked_dt)
self.assertEqual(3, events[0].count)
def testDedupe_AddNewerQuarantineData(self):
blockable = test_utils.CreateSantaBlockable()
quarantine = event_models.QuarantineMetadata(
data_url='http://notbad.com',
referer_url='http://sourceforge.com',
downloaded_dt=datetime.datetime.utcnow(),
agent_bundle_id='123456')
now = datetime.datetime.utcnow()
event = test_utils.CreateSantaEvent(blockable,
quarantine=quarantine,
last_blocked_dt=now)
new_quarantine = datastore_utils.CopyEntity(event.quarantine,
data_url='http://3vil.com')
later_dt = event.last_blocked_dt + datetime.timedelta(seconds=1)
later_event = datastore_utils.CopyEntity(event,
quarantine=new_quarantine,
last_blocked_dt=later_dt)
event.Dedupe(later_event)
self.assertEqual('http://3vil.com', event.quarantine.data_url)
def testModelWithComputedProperty(self):
class A(ndb.Model):
a = ndb.StringProperty()
b = ndb.ComputedProperty(lambda self: self.a[0])
inst = A(a='xyz')
inst.put()
self.assertEqual('x', inst.b)
new = utils.CopyEntity(inst, a='abc')
new.put()
self.assertEqual('a', new.b)
def _DedupeExistingAndPut(cls, events):
"""Dedupes a list of new-style Events with existing Events and puts them."""
# NOTE: We copy each entity in the input list because this function
# is transactional and, consequently, may be retried with the same
# parameters in the event of a failure. If we modify the event objects in
# place, subsequent retries will see the changes made by previous attempts.
event_copies = [
utils.CopyEntity(event, new_key=event.key) for event in events]
existing_events = yield ndb.get_multi_async(event.key for event in events)
for event, existing_event in zip(event_copies, existing_events):
if existing_event:
event.Dedupe(existing_event)
yield ndb.put_multi_async(event_copies)
def testPolyModel_NoClass(self):
class A(utils.polymodel.PolyModel):
a = ndb.StringProperty()
class B(A):
pass
inst = B(a='abc')
a_copy = utils.CopyEntity(inst, a='xyz')
a_copy.put()
inst.put()
self.assertEqual('xyz', a_copy.a)
self.assertEqual('abc', inst.a)
def testDedupeSantaEvent_AddOldQuarantineData(self):
quarantine = self.santa_event.quarantine
self.santa_event.quarantine = None
self.santa_event.put()
earlier_dt = (self.santa_event.first_blocked_dt -
datetime.timedelta(seconds=1))
new_event = utils.CopyEntity(self.santa_event,
quarantine=quarantine,
first_blocked_dt=earlier_dt)
self.santa_event.Dedupe(new_event)
self.assertIsNotNone(self.santa_event.quarantine)
def testDedupeSantaEvent(self):
later_dt = (self.santa_event.last_blocked_dt +
datetime.timedelta(seconds=1))
new_event = utils.CopyEntity(
self.santa_event,
quarantine=None,
event_type=constants.EVENT_TYPE.BLOCK_CERTIFICATE,
last_blocked_dt=later_dt)
self.santa_event.Dedupe(new_event)
self.assertEqual(constants.EVENT_TYPE.BLOCK_CERTIFICATE,
self.santa_event.event_type)
self.assertIsNotNone(self.santa_event.quarantine)
def testPolyModel(self):
class A(utils.polymodel.PolyModel):
a = ndb.StringProperty()
class B(A):
pass
inst = B(a='abc')
inst.put()
new = utils.CopyEntity(inst, a='xyz')
new.put()
self.assertEqual('xyz', new.a)
self.assertIsInstance(new, B)
def testGetVotes_Inactive(self):
self.assertLen(self.blockable_1.GetVotes(), 0)
test_utils.CreateVotes(self.blockable_1, 2)
self.assertLen(self.blockable_1.GetVotes(), 2)
votes = vote_models.Vote.query().fetch()
new_votes = []
for vote in votes:
new_key = ndb.Key(flat=vote.key.flat()[:-1] + (None,))
new_votes.append(datastore_utils.CopyEntity(vote, new_key=new_key))
ndb.delete_multi(vote.key for vote in votes)
ndb.put_multi(new_votes)
self.assertLen(self.blockable_1.GetVotes(), 0)
def testDedupeSantaEvent_AddOldQuarantineData(self):
quarantine = self.santa_event.quarantine
self.santa_event.quarantine = None
self.santa_event.put()
earlier_dt = (
self.santa_event.first_blocked_dt - datetime.timedelta(seconds=1))
earlier_event = datastore_utils.CopyEntity(
self.santa_event,
quarantine=quarantine,
event_type=constants.EVENT_TYPE.BLOCK_CERTIFICATE,
first_blocked_dt=earlier_dt)
self.santa_event.Dedupe(earlier_event)
self.assertNotEqual(
constants.EVENT_TYPE.BLOCK_CERTIFICATE, self.santa_event.event_type)
self.assertIsNotNone(self.santa_event.quarantine)
