def test_get_utc_now_microtimestamp(self):
# 2018/01/01 00:00:00.123456
mock_date = datetime(2018, 1, 1, 0, 0, 0, 123456)
result = date_time.get_utc_now_microtimestamp(mock_date)
# Expected timestamp = 1514764800 * 1000000 + 123456
expected_result = 1514764800123456
self.assertEqual(expected_result, result)
def test_scanner_index_create(self, mock_date_time):
"""`ScannerIndex` create() works as expected."""
utc_now = datetime.utcnow()
mock_date_time.return_value = utc_now
expected_id = date_time.get_utc_now_microtimestamp(utc_now)
db_row = scanner_dao.ScannerIndex.create(expected_id)
self.assertEquals(expected_id, db_row.id)
self.assertEquals(utc_now, db_row.created_at_datetime)
self.assertEquals(IndexState.CREATED, db_row.scanner_status)
def test_init_scanner_index(self, mock_date_time):
utc_now = datetime.utcnow()
mock_date_time.return_value = utc_now
scanner.init_scanner_index(self.session, self.inv_index_id3)
expected_id = date_time.get_utc_now_microtimestamp(utc_now)
db_row = (self.session.query(scanner_dao.ScannerIndex).filter(
scanner_dao.ScannerIndex.id == expected_id).one())
self.assertEquals(IndexState.RUNNING, db_row.scanner_status)
self.assertEquals(utc_now, db_row.created_at_datetime)
def test_scanner_index_complete(self, mock_date_time):
"""`ScannerIndex` complete() works as expected."""
start = datetime.utcnow()
end = start + timedelta(minutes=5)
# ScannerIndex.create() calls get_utc_now_datetime() twice.
mock_date_time.side_effect = [start, end]
expected_id = date_time.get_utc_now_microtimestamp(start)
db_row = scanner_dao.ScannerIndex.create(expected_id)
self.assertEquals(expected_id, db_row.id)
db_row.complete()
self.assertEquals(end, db_row.completed_at_datetime)
self.assertEquals(IndexState.SUCCESS, db_row.scanner_status)
def create(cls):
"""Create a new inventory index row.
Returns:
InventoryIndex: InventoryIndex row object.
"""
utc_now = date_time.get_utc_now_datetime()
micro_timestamp = date_time.get_utc_now_microtimestamp(utc_now)
return InventoryIndex(id=micro_timestamp,
created_at_datetime=utc_now,
completed_at_datetime=None,
inventory_status=IndexState.CREATED,
schema_version=CURRENT_SCHEMA,
counter=0)
def _enforce_projects(self,
project_policies,
prechange_callback=None,
new_result_callback=None,
add_rule_callback=None):
"""Do a single enforcement run on the projects.
Args:
project_policies (iterable): An iterable of
(project_id, firewall_policy) tuples to enforce.
prechange_callback (Callable): See docstring for self.Run().
new_result_callback (Callable): See docstring for self.Run().
add_rule_callback (Callable): See docstring for self.Run().
Returns:
int: The number of projects that were enforced.
"""
# Get a 64 bit int to use as the unique batch ID for this run.
batch_id = date_time.get_utc_now_microtimestamp()
self.enforcement_log.summary.batch_id = batch_id
projects_enforced_count = 0
future_to_key = {}
with (concurrent.futures.ThreadPoolExecutor(
max_workers=self._concurrent_workers)) as executor:
for (project_id, firewall_policy) in project_policies:
future = executor.submit(self._enforce_project, project_id,
firewall_policy, prechange_callback,
add_rule_callback)
future_to_key[future] = project_id
for future in concurrent.futures.as_completed(future_to_key):
project_id = future_to_key[future]
LOGGER.debug('Project %s finished enforcement run.',
project_id)
projects_enforced_count += 1
result = self.enforcement_log.results.add()
result.CopyFrom(future.result())
# Make sure all results have the current batch_id set
result.batch_id = batch_id
result.run_context = enforcer_log_pb2.ENFORCER_BATCH
if new_result_callback:
new_result_callback(result)
return projects_enforced_count
def create(cls, inv_index_id):
"""Create a new scanner index row.
Args:
inv_index_id (str): Id of the inventory index.
Returns:
object: ScannerIndex row object.
"""
utc_now = date_time.get_utc_now_datetime()
micro_timestamp = date_time.get_utc_now_microtimestamp(utc_now)
return ScannerIndex(id=micro_timestamp,
inventory_index_id=inv_index_id,
created_at_datetime=utc_now,
scanner_status=IndexState.CREATED,
schema_version=CURRENT_SCHEMA)
def __init__(self,
project_id,
global_configs=None,
compute_client=None,
dry_run=False,
project_sema=None,
max_running_operations=0):
"""Initialize.
Args:
project_id (str): The project id for the project to enforce.
global_configs (dict): Global configurations.
compute_client (ComputeClient): A Compute API client.
If not provided, one will be created using the default
credentials.
dry_run (bool): Set to true to ensure no actual changes are made to
the project. EnforcePolicy will still return a ProjectResult
proto showing the changes that would have been made.
project_sema (threading.BoundedSemaphore): An optional semaphore
object, used to limit the number of concurrent projects getting
written to.
max_running_operations (int): [DEPRECATED] Used to limit the number
of concurrent running operations on an API.
"""
self.project_id = project_id
if not compute_client:
compute_client = compute.ComputeClient(global_configs,
dry_run=dry_run)
self.compute_client = compute_client
self.result = enforcer_log_pb2.ProjectResult()
self.result.status = STATUS_UNSPECIFIED
self.result.project_id = self.project_id
self.result.timestamp_sec = date_time.get_utc_now_microtimestamp()
self._dry_run = dry_run
self._project_sema = project_sema
if max_running_operations:
LOGGER.warning(
'Max running operations is deprecated. Argument ignored.')
self._operation_sema = None
def test_get_latest_scanner_index_id(self, mock_date_time):
"""The method under test returns the newest `ScannerIndex` row."""
time1 = datetime.utcnow()
time2 = time1 + timedelta(minutes=5)
time3 = time1 + timedelta(minutes=7)
mock_date_time.side_effect = [time1, time2, time3]
expected_id = date_time.get_utc_now_microtimestamp(time2)
self.session.add(scanner_dao.ScannerIndex.create(expected_id))
index2 = scanner_dao.ScannerIndex.create(expected_id)
index2.scanner_status = IndexState.SUCCESS
self.session.add(index2)
self.session.add(scanner_dao.ScannerIndex.create(expected_id))
self.session.flush()
self.assertEquals(
expected_id,
scanner_dao.get_latest_scanner_index_id(self.session, expected_id))
def run(self,
project_policies,
prechange_callback=None,
new_result_callback=None,
add_rule_callback=None):
"""Runs the enforcer over all projects passed in to the function.
Args:
project_policies (iterable): An iterable of
(project_id, firewall_policy) tuples to enforce or a dictionary
in the format {project_id: firewall_policy}
prechange_callback (Callable): A callback function that will get
called if the firewall policy for a project does not match the
expected policy, before any changes are actually applied. If the
callback returns False then no changes will be made to the
project. If it returns True then the changes will be pushed.
See FirewallEnforcer.apply_firewall() docstring for more details.
new_result_callback (Callable): An optional function to call with each
new result proto message as they are returned from a
ProjectEnforcer thread.
add_rule_callback (Callable): A callback function that checks whether
a firewall rule should be applied. If the callback returns False,
that rule will not be modified.
Returns:
enforcer_log_pb2.EnforcerLog: The EnforcerLog proto for the last run,
including individual results for each project, and a summary of all
results.
"""
if self._dry_run:
LOGGER.info('Simulating changes')
if isinstance(project_policies, dict):
project_policies = list(project_policies.items())
self.enforcement_log.Clear()
self.enforcement_log.summary.projects_total = len(project_policies)
started_time = date_time.get_utc_now_datetime()
LOGGER.info('starting enforcement wave: %s', started_time)
projects_enforced_count = self._enforce_projects(
project_policies, prechange_callback, new_result_callback,
add_rule_callback)
finished_time = date_time.get_utc_now_datetime()
started_timestamp = date_time.get_utc_now_unix_timestamp(started_time)
finished_timestamp = date_time.get_utc_now_unix_timestamp(
finished_time)
total_time = finished_timestamp - started_timestamp
LOGGER.info('finished wave in %i seconds', total_time)
self.enforcement_log.summary.timestamp_start_msec = (
date_time.get_utc_now_microtimestamp(started_time))
self.enforcement_log.summary.timestamp_end_msec = (
date_time.get_utc_now_microtimestamp(finished_time))
self._summarize_results()
if not projects_enforced_count:
LOGGER.warning('No projects enforced on the last run, exiting.')
return self.enforcement_log
import shutil
import tempfile
import unittest
from datetime import datetime
from tests.unittest_utils import ForsetiTestCase
from google.cloud.forseti.common.util import date_time
from google.cloud.forseti.common.util import logger
from google.cloud.forseti.services.dao import create_engine
from google.cloud.forseti.services.dao import ModelManager
from google.cloud.forseti.services.model.importer import importer
from google.cloud.forseti.services.model.importer.importer import InventoryImporter
FAKE_DATETIME = datetime(2018, 1, 28, 10, 20, 30, 0)
FAKE_DATETIME_TIMESTAMP = date_time.get_utc_now_microtimestamp(FAKE_DATETIME)
class ServiceConfig(object):
"""Helper class to implement dependency injection to Forseti Server services.
"""
def __init__(self, db_connect_string):
engine = create_engine(db_connect_string, echo=False)
self.model_manager = ModelManager(engine)
def run_in_background(self, function):
"""Runs a function in a thread pool in the background."""
return function()
def get_db_file_path(db_name):
