def testChangeModeForGroup_MultiBatch(self, mock_ctor):
users = [
test_utils.CreateUser()
for _ in xrange(role_syncing.BATCH_SIZE + 1)
]
hosts = [
test_utils.CreateSantaHost(primary_user=user_utils.EmailToUsername(
user.key.id()),
client_mode=MONITOR) for user in users
]
mock_ctor.return_value.AllMembers.return_value = [
user.key.id() for user in users
]
response = self.testapp.get('', headers={'X-AppEngine-Cron': 'true'})
self.assertEqual(httplib.OK, response.status_int)
self.assertTaskCount(constants.TASK_QUEUE.DEFAULT, 2)
self.DrainTaskQueue(constants.TASK_QUEUE.DEFAULT)
new_hosts = ndb.get_multi(host.key for host in hosts)
self.assertTrue(all(host.client_mode == LOCKDOWN
for host in new_hosts))
def _ChangeModeForHosts(mode, user_keys, honor_lock=True):
"""Performs a client mode change for the specified users' hosts.
Args:
mode: The new client_mode to set.
user_keys: The users whose host modes are to be changed.
honor_lock: bool, whether the client_mode_lock property will be honored.
"""
predicates = [
host_models.SantaHost.primary_user == user_utils.EmailToUsername(
key.id()) for key in user_keys
]
query = host_models.SantaHost.query(ndb.OR(*predicates))
hosts = query.fetch()
updated_hosts = []
for host in hosts:
# If lock is honored, skip locked users.
if honor_lock and host.client_mode_lock:
continue
# Ignore non-changes also.
if host.client_mode == mode:
continue
# Proceed with the mode change.
host.client_mode = mode
host.client_mode_lock = False
updated_hosts.append(host)
ndb.put_multi(updated_hosts)
logging.info('Client mode changed to %s for %d host(s)', mode,
len(updated_hosts))
def testGetEventKeysToInsert_HostOwner(self):
self.PatchSetting('EVENT_CREATION', constants.EVENT_CREATION.HOST_OWNER)
keys = model_utils.GetEventKeysToInsert(self.event, [], ['foo'])
self.assertLen(keys, 1)
key_usernames = [user_utils.EmailToUsername(key.flat()[1]) for key in keys]
self.assertSameElements(['foo'], key_usernames)
def testModeChange(self):
host = test_utils.CreateSantaHost(
primary_user=user_utils.EmailToUsername(self.user.key.id()),
client_mode=MONITOR)
role_syncing._ChangeModeForHosts(LOCKDOWN, [self.user.key])
host = host.key.get()
self.assertEqual(LOCKDOWN, host.client_mode)
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 testClientModeLockOnHonored(self):
host = test_utils.CreateSantaHost(
primary_user=user_utils.EmailToUsername(self.user.key.id()),
client_mode=MONITOR,
client_mode_lock=True)
role_syncing._ChangeModeForHosts(LOCKDOWN, [self.user.key])
host = host.key.get()
self.assertTrue(host.client_mode_lock)
self.assertEqual(MONITOR, host.client_mode)
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 CreateTestEntities(email_addr):
"""Create some test Datastore data if specified, but only if running locally.
Note that this code doesn't (and shouldn't) delete any existing entities.
The risk of such code being accidentally triggered in prod is too great, so
if local entities need to be deleted, use the local Datastore viewer (e.g.
http://127.0.0.1:8000/datastore).
Args:
email_addr: Email address of the local users for whom test data should
be created.
Raises:
NotRunningLocallyError: if called anywhere other than a local deployment.
"""
if not env_utils.RunningLocally():
raise NotRunningLocallyError
# Create a user entity with all available roles.
user = user_models.User.GetOrInsert(email_addr=email_addr)
user_models.User.SetRoles(email_addr, constants.USER_ROLE.SET_ALL)
username = user_utils.EmailToUsername(email_addr)
# Create associated SantaHosts for the user.
santa_hosts = CreateSantaHosts(2, primary_user=username)
# For each SantaHost, create some SantaEvents.
for santa_host in santa_hosts:
for santa_blockable in CreateSantaBlockables(5):
parent_key = datastore_utils.ConcatenateKeys(
user.key, santa_host.key, santa_blockable.key)
CreateSantaEvent(santa_blockable,
executing_user=username,
event_type=constants.EVENT_TYPE.BLOCK_BINARY,
host_id=santa_host.key.id(),
parent=parent_key)
# Create associated Bit9Hosts for the user.
bit9_hosts = CreateBit9Hosts(2, users=[username])
# For each Bit9Host, create some Bit9Events.
for bit9_host in bit9_hosts:
for bit9_binary in CreateBit9Binaries(5):
parent_key = datastore_utils.ConcatenateKeys(
user.key, bit9_host.key, bit9_binary.key)
CreateBit9Event(bit9_binary,
executing_user=username,
event_type=constants.EVENT_TYPE.BLOCK_BINARY,
host_id=bit9_host.key.id(),
parent=parent_key)
def get(self, feature):
# If requesting an unknown feature, fail closed.
if feature not in _SUPPORTED_FEATURES:
logging.error('Unsupported feature: %s', feature)
self.abort(httplib.FORBIDDEN)
# See if memcache already has an entry for this feature.
memcache_key = 'feature_%s' % feature
csv_string = memcache.get(memcache_key)
# If it's already in memcache, construct a set of approved usernames.
if csv_string:
approved_users = set(csv_string.split(','))
# Otherwise there was a cache miss, so retrieve the list of all approved
# users for this feature.
else:
try:
approved_users = set()
group_manager = group_utils.GroupManager()
for group in _SUPPORTED_FEATURES[feature]:
# If a group isn't found, fail closed.
if not group_manager.DoesGroupExist(group):
logging.error('Unknown group: %s', group)
self.abort(httplib.FORBIDDEN)
approved_users |= set(
user_utils.EmailToUsername(member)
for member in group_manager.AllMembers(group))
# Build a CSV string and stuff it in memcache for future use.
csv_string = ','.join(sorted(list(approved_users)))
memcache.set(memcache_key, csv_string)
# If anything fails while interacting with the GroupManager, just fail
# closed. Odds are the user doesn't have access to this feature anyway,
# and should be able to go about their business without seeing an error
# message.
except Exception: # pylint: disable=broad-except
logging.exception('Unexpected error while retrieving group members')
self.abort(httplib.FORBIDDEN)
approved = self.user.nickname in approved_users
self.response.status = httplib.OK if approved else httplib.FORBIDDEN
def InsertBigQueryRow(self, **kwargs):
user = None
if self.user_key:
user = user_utils.EmailToUsername(self.user_key.id())
defaults = {
'sha256': self.key.parent().id(),
'timestamp': datetime.datetime.utcnow(),
'scope': LOCAL if self.host_id or self.user_key else GLOBAL,
'policy': self.policy,
'target_type': self.rule_type,
'device_id': self.host_id if self.host_id else None,
'user': user}
defaults.update(kwargs.copy())
tables.RULE.InsertRow(**defaults)
def _CopyLocalRules(user_key, dest_host_id):
"""Creates copies of all local rules for the new host."""
logging.info('Copying rules for user %s to host %s', user_key.id(),
dest_host_id)
# 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.
username = user_utils.EmailToUsername(user_key.id())
query = host_models.SantaHost.query(
host_models.SantaHost.primary_user == username,
host_models.SantaHost.last_postflight_dt != None) # pylint: disable=g-equals-none
src_host = query.get()
if src_host is None:
logging.warning('User %s has no hosts to copy from', username)
return datastore_utils.GetNoOpFuture()
else:
logging.info('Copying local rules from %s', src_host.key.id())
# Query for all SantaRules for the given user on the chosen host.
query = rule_models.SantaRule.query(
rule_models.SantaRule.host_id == src_host.key.id(),
rule_models.SantaRule.user_key == user_key)
# Copy the local rules to the new host.
new_rules = []
for src_rules in datastore_utils.Paginate(query):
for src_rule in src_rules:
logging.info('Copying local rule for %s',
src_rule.key.parent().id())
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()
logging.info('Copying %d rule(s) to host %s', len(new_rules), dest_host_id)
futures = ndb.put_multi_async(new_rules)
return datastore_utils.GetMultiFuture(futures)
def testUserGetListOwnEventsWithBlockableAndContext(self):
"""Normal user getting list of their events with context, by blockable."""
test_utils.CreateVote(self.santa_blockable1,
user_email=self.user_1.email)
test_utils.CreateVote(self.santa_blockable1,
user_email=self.user_2.email)
params = {
'blockableKey': self.santa_blockable1.key.urlsafe(),
'withContext': 'true'
}
with self.LoggedInUser(user=self.user_1):
response = self.testapp.get(self.ROUTE, params)
output = response.json
self.assertIn('application/json', response.headers['Content-type'])
self.assertIsInstance(output, dict)
self.assertLen(output['content'], 2)
event_with_context = output['content'][0]
self.assertLen(event_with_context.keys(), 5)
self.assertEqual(event_with_context['host']['id'],
event_with_context['event']['hostId'])
blockable_key = ndb.Key(
urlsafe=event_with_context['event']['blockableKey'])
self.assertEqual(event_with_context['blockable']['id'],
blockable_key.id())
self.assertEqual(event_with_context['blockable']['fileName'],
event_with_context['event']['fileName'])
self.assertEqual(event_with_context['cert']['id'],
event_with_context['blockable']['certId'])
self.assertEqual(
user_utils.EmailToUsername(
event_with_context['vote']['userEmail']),
event_with_context['event']['executingUser'])
def nickname(self):
return user_utils.EmailToUsername(self.key.string_id())
def testEmailToUsername(self):
self.assertEqual('user', user_utils.EmailToUsername('*****@*****.**'))
self.assertEqual('user', user_utils.EmailToUsername('user'))
