Improve the inventory loading email. (#131)

* Improve the inventory loading email. * tweak * tweak * address comments
forseti-security · Mar 21, 2017 · 454ae06 · 454ae06
1 parent 751e385
commit 454ae06
Show file tree

Hide file tree

Showing 6 changed files with 121 additions and 17 deletions.
diff --git a/google/cloud/security/common/data_access/dao.py b/google/cloud/security/common/data_access/dao.py
@@ -97,6 +97,30 @@ def load_data(self, resource_name, timestamp, data):
                 OperationalError, ProgrammingError) as e:
             raise MySQLError(resource_name, e)
 
+    def select_record_count(self, resource_name, timestamp):
+        """Select the record count from a snapshot table.
+
+        Args:
+            resource_name: String of the resource name, which is embedded in
+                the table name.
+            timestamp: String of timestamp, formatted as YYYYMMDDTHHMMSSZ.
+
+        Returns:
+             Integer of the record count in a snapshot table.
+
+        Raises:
+            MySQLError: An error with MySQL has occurred.
+        """
+        try:
+            record_count_sql = select_data.RECORD_COUNT.format(
+                resource_name, timestamp)
+            cursor = self.conn.cursor()
+            cursor.execute(record_count_sql)
+            return cursor.fetchone()[0]
+        except (DataError, IntegrityError, InternalError, NotSupportedError,
+                OperationalError, ProgrammingError) as e:
+            raise MySQLError(resource_name, e)
+
     def select_project_numbers(self, resource_name, timestamp):
         """Select the project numbers from a snapshot table.
 

diff --git a/google/cloud/security/common/data_access/sql_queries/select_data.py b/google/cloud/security/common/data_access/sql_queries/select_data.py
@@ -14,6 +14,10 @@
 
 """SQL queries to select data from snapshot tables."""
 
+RECORD_COUNT = """
+    SELECT COUNT(*) FROM {0}_{1};
+"""
+
 PROJECT_NUMBERS = """
     SELECT project_number from projects_{0};
 """

diff --git a/google/cloud/security/common/email_templates/inventory_snapshot_summary.jinja b/google/cloud/security/common/email_templates/inventory_snapshot_summary.jinja
@@ -0,0 +1,42 @@
+<!--
+Copyright 2017 Google Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+<!doctype html>
+<html>
+<body>
+    Organization ID: {{ organization_id }}<br>
+    Timestamp: {{ cycle_timestamp }}; {{ cycle_time }}<br>
+    Status Summary: {{ status_summary }}<br>
+    <br>
+    <table>
+        <col width="250">
+        <col width="150">
+        <col width="150">
+        <tr>
+            <th align="left">Resource Name</th>
+            <th align="left">Status</th>
+            <th align="left"># in Snapshot</th>
+        </tr>
+        {% for pipeline in pipelines %}
+        <tr>
+            <td>{{ pipeline['pipeline'].RESOURCE_NAME }}</td>
+            <td>{{ pipeline['status'] }}</td>
+            <td>{{ pipeline['count'] }}</td>
+        </tr>
+        {% endfor %}
+    </table>
+</body>
+</html>
diff --git a/google/cloud/security/common/util/email_util.py b/google/cloud/security/common/util/email_util.py
@@ -77,7 +77,7 @@ def _execute_send(self, email):
         return self.sendgrid.client.mail.send.post(request_body=email.get())
 
     def send(self, email_sender=None, email_recipient=None,
-             email_subject=None, email_content=None, content_type='text/plain',
+             email_subject=None, email_content=None, content_type=None,
              attachment=None):
         """Send an email.
 

diff --git a/google/cloud/security/inventory/inventory_loader.py b/google/cloud/security/inventory/inventory_loader.py
@@ -112,20 +112,21 @@ def _create_snapshot_cycles_table(dao):
         LOGGER.error('Unable to create snapshot cycles table: %s', e)
         sys.exit()
 
-def _start_snapshot_cycle(dao):
+def _start_snapshot_cycle(cycle_time, cycle_timestamp, dao):
     """Start snapshot cycle.
 
     Args:
+        cycle_time: Datetime object of the cycle, in UTC.
+        cycle_timestamp: String of timestamp, formatted as YYYYMMDDTHHMMSSZ.
         dao: Data access object.
 
     Returns:
-        cycle_timestamp: String of timestamp, formatted as YYYYMMDDTHHMMSSZ.
+        None
 
     Raises:
         MySQLError: An error with MySQL has occurred.
     """
-    cycle_time = datetime.utcnow()
-    cycle_timestamp = cycle_time.strftime(CYCLE_TIMESTAMP_FORMAT)
+
 
     if not _exists_snapshot_cycles_table(dao):
         LOGGER.info('snapshot_cycles is not created yet.')
@@ -141,7 +142,6 @@ def _start_snapshot_cycle(dao):
         sys.exit()
 
     LOGGER.info('Inventory snapshot cycle started: %s', cycle_timestamp)
-    return cycle_timestamp
 
 def _complete_snapshot_cycle(dao, cycle_timestamp, status):
     """Complete the snapshot cycle.
@@ -171,13 +171,18 @@ def _complete_snapshot_cycle(dao, cycle_timestamp, status):
     LOGGER.info('Inventory load cycle completed with %s: %s',
                 status, cycle_timestamp)
 
-def _send_email(cycle_timestamp, status, sendgrid_api_key,
-                email_sender, email_recipient, email_content=None):
+def _send_email(organization_id, cycle_time, cycle_timestamp, status, pipelines,
+                dao, sendgrid_api_key, email_sender, email_recipient,
+                email_content=None):
     """Send an email.
 
     Args:
+        organization_id: String of the organization id
+        cycle_time: Datetime object of the cycle, in UTC.
         cycle_timestamp: String of timestamp, formatted as YYYYMMDDTHHMMSSZ.
-        status: String of the current snapshot cycle.
+        status: String of the overall status of current snapshot cycle.
+        pipelines: List of pipelines and their statuses.
+        dao: Data access object.
         sendgrid_api_key: String of the sendgrid api key to auth email service.
         email_sender: String of the sender of the email.
         email_recipient: String of the recipient of the email.
@@ -186,15 +191,36 @@ def _send_email(cycle_timestamp, status, sendgrid_api_key,
     Returns:
          None
     """
-    email_subject = 'Inventory loading {0}: {1}'.format(cycle_timestamp, status)
 
-    if email_content is None:
-        email_content = email_subject
+    for pipeline in pipelines:
+        try:
+            pipeline['count'] = dao.select_record_count(
+                pipeline['pipeline'].RESOURCE_NAME,
+                cycle_timestamp)
+        except MySQLError as e:
+            LOGGER.error('Unable to retrieve record count for %s_%s:\n%s',
+                         pipeline['pipeline'].RESOURCE_NAME,
+                         cycle_timestamp,
+                         e)
+            pipeline['count'] = 'N/A'
+
+    email_subject = 'Inventory Snapshot Complete: {0} {1}'.format(
+        cycle_timestamp, status)
+
+    email_content = EmailUtil.render_from_template(
+        'inventory_snapshot_summary.jinja', {
+            'organization_id': organization_id,
+            'cycle_time': cycle_time.strftime('%Y %b %d, %H:%M:%S (UTC)'),
+            'cycle_timestamp': cycle_timestamp,
+            'status_summary': status,
+            'pipelines': pipelines,
+        })
 
     try:
         email_util = EmailUtil(sendgrid_api_key)
         email_util.send(email_sender, email_recipient,
-                        email_subject, email_content)
+                        email_subject, email_content,
+                        content_type='text/html')
     except EmailSendError:
         LOGGER.error('Unable to send email that inventory snapshot completed.')
 
@@ -209,7 +235,9 @@ def main(argv):
         LOGGER.error('Encountered error with Cloud SQL. Abort.\n%s', e)
         sys.exit()
 
-    cycle_timestamp = _start_snapshot_cycle(dao)
+    cycle_time = datetime.utcnow()
+    cycle_timestamp = cycle_time.strftime(CYCLE_TIMESTAMP_FORMAT)
+    _start_snapshot_cycle(cycle_time, cycle_timestamp, dao)
 
     configs = FLAGS.FlagValuesDict()
 
@@ -250,9 +278,15 @@ def main(argv):
         snapshot_cycle_status = 'FAILURE'
 
     _complete_snapshot_cycle(dao, cycle_timestamp, snapshot_cycle_status)
-    _send_email(cycle_timestamp, snapshot_cycle_status,
+    _send_email(configs.get('organization_id'),
+                cycle_time,
+                cycle_timestamp,
+                snapshot_cycle_status,
+                pipelines,
+                dao,
                 configs.get('sendgrid_api_key'),
-                configs.get('email_sender'), configs.get('email_recipient'))
+                configs.get('email_sender'),
+                configs.get('email_recipient'))
 
 
 if __name__ == '__main__':

diff --git a/google/cloud/security/inventory/pipelines/load_projects_iam_policies_pipeline.py b/google/cloud/security/inventory/pipelines/load_projects_iam_policies_pipeline.py
@@ -27,11 +27,11 @@
 from google.cloud.security.inventory.errors import LoadDataPipelineError
 
 
-
 LOGGER = LogUtil.setup_logging(__name__)
 RESOURCE_NAME = 'project_iam_policies'
 RAW_PROJECT_IAM_POLICIES = 'raw_project_iam_policies'
 
+
 def run(dao=None, cycle_timestamp=None, configs=None, crm_rate_limiter=None):
     """Runs the load IAM policies data pipeline.