Merge pull request #1 from alannix-lw/awoj/add-pdf-compliance

add PDF output functionality for compliance library
lacework · Jun 24, 2020 · 3146061 · 3146061
2 parents 3207926 + 146a06e
commit 3146061
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 12 deletions.
diff --git a/examples/example_compliance.py b/examples/example_compliance.py
@@ -24,7 +24,10 @@
     # Compliance API
 
     # Get latest compliance report in JSON format for AWS account
-    lacework_client.compliance.get_latest_aws_report("123456789")
+    lacework_client.compliance.get_latest_aws_report(aws_account_id="123456789", file_format="json")
+
+    # Get latest compliance report in PDF format for AWS account
+    lacework_client.compliance.get_latest_aws_report(aws_account_id="123456789", file_format="pdf", pdf_path='<PATH_TO_PDF_OUTPUT>')
 
     # Get a list of subscriptions for an Azure Tenant
     lacework_client.compliance.list_azure_subscriptions("xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx")
diff --git a/laceworksdk/api/compliance.py b/laceworksdk/api/compliance.py
@@ -29,15 +29,17 @@ def __init__(self, session):
 
     def get_latest_aws_report(self,
                               aws_account_id,
-                              file_format="json",
-                              report_type=None):
+                              file_format=None,
+                              report_type=None,
+                              pdf_path=None):
         """
         A method to get the latest compiance report for an AWS account.
 
         :param aws_account_id: A string representing which AWS Account to query.
         :param file_format: A string representing the desired file format. ("pdf" or "json")
         :param report_type: A string representing the desired report type.
             ("AWS_CIS_S3", "NIST_800-53_Rev4", "ISO_2700", "HIPAA", "SOC", or "PCI")
+        :param pdf_path: An absolute path for writing PDF compliance reports
 
         :return response json
         """
@@ -54,15 +56,22 @@ def get_latest_aws_report(self,
 
         response = self._session.get(api_uri)
 
-        logger.debug(json.dumps(response.json(), indent=2))
+        #logger.debug(json.dumps(response.json(), indent=2))
 
-        return response.json()
+        if file_format == "json":
+            logger.debug(json.dumps(response.json(), indent=2))
+            return response.json()
+        elif file_format == "pdf":
+            logger.debug('creating pdf at {}'.format(pdf_path))
+            with open(pdf_path, 'wb') as f:
+                f.write(response.content)
 
     def get_latest_azure_report(self,
                                 azure_tenant_id,
                                 azure_subscription_id,
-                                file_format="json",
-                                report_type=None):
+                                file_format=None,
+                                report_type=None,
+                                pdf_path=None):
         """
         A method to get the latest compiance report for an Azure tenant.
 
@@ -71,6 +80,7 @@ def get_latest_azure_report(self,
         :param file_format: A string representing the desired file format. ("pdf" or "json")
         :param report_type: A string representing the desired report type.
             ("AZURE_CIS", "AZURE_SOC", or "AZURE_PCI")
+        :param pdf_path: An absolute path for writing PDF compliance reports
 
         :return response json
         """
@@ -87,15 +97,22 @@ def get_latest_azure_report(self,
 
         response = self._session.get(api_uri)
 
-        logger.debug(json.dumps(response.json(), indent=2))
+        #logger.debug(json.dumps(response.json(), indent=2))
 
-        return response.json()
+        if file_format == "json":
+            logger.debug(json.dumps(response.json(), indent=2))
+            return response.json()
+        elif file_format == "pdf":
+            logger.debug('creating pdf at {}'.format(pdf_path))
+            with open(pdf_path, 'wb') as f:
+                f.write(response.content)
 
     def get_latest_gcp_report(self,
                               gcp_organization_id,
                               gcp_project_id,
-                              file_format="json",
-                              report_type=None):
+                              file_format=None,
+                              report_type=None,
+                              pdf_path=None):
         """
         A method to get the latest compiance report for a Google Cloud organization.
 
@@ -104,6 +121,7 @@ def get_latest_gcp_report(self,
         :param file_format: A string representing the desired file format. ("pdf" or "json")
         :param report_type: A string representing the desired report type.
             ("GCP_CIS", "GCP_SOC", or "GCP_PCI")
+        :param pdf_path: An absolute path for writing PDF compliance reports
 
         :return response json
         """
@@ -121,8 +139,15 @@ def get_latest_gcp_report(self,
         response = self._session.get(api_uri)
 
         logger.debug(json.dumps(response.json(), indent=2))
+
+        if file_format == "json":
+            logger.debug(json.dumps(response.json(), indent=2))
+            return response.json()
+        elif file_format == "pdf":
+            logger.debug('creating pdf at {}'.format(pdf_path))
+            with open(pdf_path, 'wb') as f:
+                f.write(response.content)
 
-        return response.json()
 
     def list_azure_subscriptions(self, azure_tenant_id):
         """