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 get(self, user_id):
user = base_db.User.GetById(user_id)
if user is None:
logging.warning('Unknown user ID: %s', user_id)
self.abort(httplib.NOT_FOUND)
bit9_hosts = _GetAssociatedHosts(user.nickname)
ids = [str(host.id) for host in bit9_hosts]
# Get Host datastore entities corresponding to the IDs retrieved.
gae_hosts = ndb.get_multi(
ndb.Key(bit9_db.Bit9Host, id_) for id_ in ids)
hosts_to_put = []
for gae_host, bit9_host in zip(gae_hosts, bit9_hosts):
policy_key = (ndb.Key(bit9_db.Bit9Policy, str(bit9_host.policy_id))
if bit9_host.policy_id is not None else None)
changed = False
if gae_host is None:
# If the host doesn't exist, create it.
hostname = utils.ExpandHostname(
rest_utils.StripDownLevelDomain(bit9_host.name))
gae_host = bit9_db.Bit9Host(id=str(bit9_host.id),
hostname=hostname,
last_event_dt=None,
policy_key=policy_key,
users=rest_utils.ExtractHostUsers(
bit9_host.users))
changed = True
else:
# If the host does exist, update any changed fields.
if gae_host.policy_key != policy_key:
gae_host.policy_key = policy_key
changed = True
if set(gae_host.users) != set(bit9_host.users):
gae_host.users = bit9_host.users
changed = True
if changed:
hosts_to_put.append(gae_host)
if hosts_to_put:
logging.info('Updating %s Bit9Hosts...', len(hosts_to_put))
ndb.put_multi(hosts_to_put)
self.respond_json(ids)
def testSuccess(self):
user = test_utils.CreateUser(
email='foobar@' + settings.USER_EMAIL_DOMAIN)
computer = api.Computer(
id=123, name='foo', policy_id=456,
users=settings.AD_DOMAIN + '\\foobar')
self._PatchApiRequests([computer])
response = self.testapp.get('/%s' % user.key.id())
self.assertSameElements(['123'], response.json)
hosts = bit9_db.Bit9Host.query().fetch()
self.assertEqual(1, len(hosts))
host = hosts[0]
self.assertEqual('123', host.key.id())
self.assertEqual(utils.ExpandHostname('foo'), host.hostname)
self.assertEqual(ndb.Key(bit9_db.Bit9Policy, '456'), host.policy_key)
self.assertEqual(['foobar'], host.users)
def _PersistBit9Host(computer, occurred_dt):
"""Creates a Bit9Host from the Event protobuf if one does not already exist.
NOTE: This function could be transactional but, at least for now, host puts in
multiple requests don't really need to be processed in a fixed order.
last_event_dt is the only frequently modified property and there's currently
no need for it to be perfectly accurate.
Args:
computer: api.Computer object associated with the event.
occurred_dt: datetime object corresponding to the time of the event.
Returns:
ndb.Future that resolves when the host is updated.
"""
host_id = str(computer.id)
policy = computer.policy_id
policy_key = (ndb.Key(bit9.Bit9Policy, str(policy))
if policy is not None else None)
hostname = utils.ExpandHostname(
rest_utils.StripDownLevelDomain(computer.name))
policy_entity = policy_key.get()
mode = (policy_entity.enforcement_level
if policy_entity is not None else constants.HOST_MODE.UNKNOWN)
# Grab the corresponding Bit9Host.
bit9_host = yield bit9.Bit9Host.get_by_id_async(host_id)
host_users = list(rest_utils.ExtractHostUsers(computer.users))
# Perform initialization for users new to this host.
existing_users = set(bit9_host.users if bit9_host is not None else [])
new_host_users = set(host_users) - existing_users
for username in new_host_users:
# Create User if we haven't seen this user before.
email = user_map.UsernameToEmail(username)
user = base.User.GetOrInsert(email_addr=email)
# Copy the user's local rules over from a pre-existing host.
yield _CopyLocalRules(user.key, host_id)
# List of all row action that need to be persisted.
row_actions = []
# Doesn't exist? Guess we better fix that.
if bit9_host is None:
logging.info('Creating new Bit9Host')
bit9_host = bit9.Bit9Host(id=host_id,
hostname=hostname,
last_event_dt=occurred_dt,
policy_key=policy_key,
users=host_users)
row_actions.append(constants.HOST_ACTION.FIRST_SEEN)
else:
changed = False
if not bit9_host.last_event_dt or bit9_host.last_event_dt < occurred_dt:
bit9_host.last_event_dt = occurred_dt
changed = True
if bit9_host.hostname != hostname:
bit9_host.hostname = hostname
changed = True
if bit9_host.policy_key != policy_key:
bit9_host.policy_key = policy_key
changed = True
row_actions.append(constants.HOST_ACTION.MODE_CHANGE)
if set(bit9_host.users) != set(host_users):
bit9_host.users = host_users
changed = True
row_actions.append(constants.HOST_ACTION.USERS_CHANGE)
if not changed:
raise ndb.Return()
logging.info('Attempting to put Bit9Host...')
yield bit9_host.put_async()
for action in row_actions:
bigquery.HostRow.DeferCreate(
device_id=host_id,
timestamp=(bit9_host.recorded_dt
if action == constants.HOST_ACTION.FIRST_SEEN else
bit9_host.last_event_dt),
action=action,
hostname=hostname,
platform=constants.PLATFORM.WINDOWS,
users=host_users,
mode=mode)
def testExpandHostname_AlreadyFullyQualified(self):
hostname = 'im-a-computer.' + settings.AD_HOSTNAME.lower()
self.assertEqual(hostname, utils.ExpandHostname(hostname))
def testExpandHostname(self):
partial_hostname = 'im-a-computer'
expected = partial_hostname + '.' + settings.AD_HOSTNAME.lower()
self.assertEqual(expected, utils.ExpandHostname(partial_hostname))