def testTimeRemains(self, mock_now):
mock_now.side_effect = [
datetime.datetime(2017, 1, 1, 1, 1, s) for s in xrange(3)
]
start_time = datetime.datetime(2017, 1, 1, 1, 1, 0)
delta = datetime.timedelta(seconds=1)
self.assertTrue(time_utils.TimeRemains(start_time, delta))
self.assertFalse(time_utils.TimeRemains(start_time, delta))
self.assertFalse(time_utils.TimeRemains(start_time, delta))
def Process(host_id):
"""Processes _UnsyncedEvents for a single Windows host.
Args:
host_id: The integer ID of this host in Bit9.
"""
try:
with datastore_locks.DatastoreLock(
'bit9-process-%d' % host_id, default_timeout=_PROCESS_LOCK_TIMEOUT,
default_max_acquire_attempts=_PROCESS_LOCK_MAX_ACQUIRE_ATTEMPTS):
total_process_count = 0
start_time = _Now()
logging.info('Starting a new processing task for %d', host_id)
# Query for all _UnsyncedEvents that belong to the given host, in batches,
# and process them until we run out, or the task nears its deadline.
query = (_UnsyncedEvent.query(_UnsyncedEvent.host_id == host_id)
.order(_UnsyncedEvent.bit9_id))
event_pages = datastore_utils.Paginate(query, page_size=25)
event_page = next(event_pages, None)
while time_utils.TimeRemains(start_time, _TASK_DURATION) and event_page:
for unsynced_event in event_page:
event = api.Event.from_dict(unsynced_event.event)
signing_chain = [
api.Certificate.from_dict(cert)
for cert in unsynced_event.signing_chain
]
file_catalog = event.get_expand(api.Event.file_catalog_id)
computer = event.get_expand(api.Event.computer_id)
# Persist the event data.
persist_futures = [
_PersistBit9Certificates(signing_chain),
_PersistBit9Binary(
event, file_catalog, signing_chain,
datetime.datetime.utcnow()),
_PersistBanNote(file_catalog),
_PersistBit9Host(computer, event.timestamp),
_PersistBit9Events(event, file_catalog, computer, signing_chain)
]
ndb.Future.wait_all(persist_futures)
for persist_future in persist_futures:
persist_future.check_success()
# Now that the event sync has completed successfully, remove the
# intermediate proto entity.
unsynced_event.key.delete()
monitoring.events_processed.Increment()
total_process_count += 1
event_page = next(event_pages, None)
logging.info('Processed %d event(s)', total_process_count)
except datastore_locks.AcquireLockError:
logging.info('Unable to acquire datastore lock')
def Pull(batch_size=_PULL_BATCH_SIZE):
"""Retrieve events to sync from Bit9.
Args:
batch_size: int, The number of events to retrieve in each batch.
"""
total_pull_count = 0
start_time = _Now()
logging.info('Starting a new pull task')
try:
with datastore_locks.DatastoreLock(
_PULL_LOCK_ID,
default_timeout=_PULL_LOCK_TIMEOUT,
default_max_acquire_attempts=_PULL_LOCK_MAX_ACQUIRE_ATTEMPTS):
while time_utils.TimeRemains(start_time, _TASK_DURATION):
last_synced_id = GetLastSyncedId()
logging.info('Syncing from ID=%s', last_synced_id)
# Make an API call for a batch of events. If it fails, just log it and
# try again.
try:
event_tuples = GetEvents(last_synced_id, batch_size)
except Exception as e: # pylint: disable=broad-except
logging.warning('Event retrieval failed: %s', e)
continue
pull_count = len(event_tuples)
total_pull_count += pull_count
logging.info('Retrieved %d events (%d events total)',
pull_count, total_pull_count)
monitoring.events_pulled.IncrementBy(pull_count)
# Persist an _UnsyncedEvent for each retrieved Event proto.
ndb.put_multi(
_UnsyncedEvent.Generate(event, signing_chain)
for event, signing_chain in event_tuples)
# Briefly pause between requests in order to avoid hammering the Bit9
# server too hard.
time.sleep(0.25)
except datastore_locks.AcquireLockError:
logging.debug('Unable to acquire datastore lock')