More configuration options and support direct invocation (#132)

Author: austinbyers

cli/config.py

@@ -123,14 +123,14 @@ def name_prefix(self, value: str) -> None:
         self._config['name_prefix'] = value
 
     @property
-    def enable_carbon_black_downloader(self) -> int:
+    def enable_carbon_black_downloader(self) -> bool:
         return self._config['enable_carbon_black_downloader']
 
     @enable_carbon_black_downloader.setter
-    def enable_carbon_black_downloader(self, value: int) -> None:
-        if value not in {0, 1}:
+    def enable_carbon_black_downloader(self, value: bool) -> None:
+        if not isinstance(value, bool):
             raise InvalidConfigError(
-                'enable_carbon_black_downloader "{}" must be either 0 or 1.'.format(value)
+                'enable_carbon_black_downloader "{}" must be a boolean.'.format(value)
             )
         self._config['enable_carbon_black_downloader'] = value
 
@@ -252,7 +252,7 @@ def configure(self) -> None:
         get_input('Unique name prefix, e.g. "company_team"', self.name_prefix, self, 'name_prefix')
         enable_downloader = get_input('Enable the CarbonBlack downloader?',
                                       'yes' if self.enable_carbon_black_downloader else 'no')
-        self.enable_carbon_black_downloader = 1 if enable_downloader == 'yes' else 0
+        self.enable_carbon_black_downloader = (enable_downloader == 'yes')
 
         if self.enable_carbon_black_downloader:
             self._configure_carbon_black()

lambda_functions/analyzer/analyzer_aws_lib.py

@@ -70,29 +70,14 @@ def _elide_string_middle(text: str, max_length: int) -> str:
     return '{} ... {}'.format(text[:half_len], text[-half_len:])
 
 
-def publish_alert_to_sns(binary: BinaryInfo, topic_arn: str) -> None:
+def publish_to_sns(binary: BinaryInfo, topic_arn: str, subject: str) -> None:
     """Publish a JSON SNS alert: a binary has matched one or more YARA rules.
 
     Args:
         binary: Instance containing information about the binary.
         topic_arn: Publish to this SNS topic ARN.
+        subject: Message subject (for email subscriptions to the topic)
     """
-    subject = '[BinaryAlert] {} matches a YARA rule'.format(
-        binary.filepath or binary.computed_sha)
-    SNS.Topic(topic_arn).publish(
-        Subject=_elide_string_middle(subject, SNS_PUBLISH_SUBJECT_MAX_SIZE),
-        Message=(json.dumps(binary.summary(), indent=4, sort_keys=True))
-    )
-
-def publish_safe_to_sns(binary: BinaryInfo, topic_arn: str) -> None:
-    """Publish a JSON SNS alert: a binary has matched none and is safe.
-
-    Args:
-        binary: Instance containing information about the binary.
-        topic_arn: Publish to this SNS topic ARN.
-    """
-    subject = '[BinaryAlert] {} is a safe file'.format(
-        binary.filepath or binary.computed_sha)
     SNS.Topic(topic_arn).publish(
         Subject=_elide_string_middle(subject, SNS_PUBLISH_SUBJECT_MAX_SIZE),
         Message=(json.dumps(binary.summary(), indent=4, sort_keys=True))

lambda_functions/analyzer/binary_info.py

@@ -86,31 +86,36 @@ def filepath(self) -> str:
         return self.s3_metadata.get('filepath', '')
 
     def save_matches_and_alert(
-            self, analyzer_version: int, dynamo_table_name: str, sns_topic_arn: str) -> None:
+            self, analyzer_version: int, dynamo_table_name: str, sns_topic_arn: str,
+            sns_enabled: bool = True) -> None:
         """Save match results to Dynamo and publish an alert to SNS if appropriate.
 
         Args:
             analyzer_version: The currently executing version of the Lambda function.
             dynamo_table_name: Save YARA match results to this Dynamo table.
             sns_topic_arn: Publish match alerts to this SNS topic ARN.
+            sns_enabled: If True, match alerts are sent to SNS when applicable.
         """
         table = analyzer_aws_lib.DynamoMatchTable(dynamo_table_name)
         needs_alert = table.save_matches(self, analyzer_version)
 
         # Send alert if appropriate.
-        if needs_alert:
-            LOGGER.info('Publishing an SNS alert')
-            analyzer_aws_lib.publish_alert_to_sns(self, sns_topic_arn)
-
-    # alerts on files that are safe
-    def safe_alert_only(
-            self, sns_topic_arn: str) -> None:
-        """Publish an alert to SNS .
+        if needs_alert and sns_enabled:
+            LOGGER.info('Publishing a YARA match alert to %s', sns_topic_arn)
+            subject = '[BinaryAlert] {} matches a YARA rule'.format(
+                self.filepath or self.computed_sha)
+            analyzer_aws_lib.publish_to_sns(self, sns_topic_arn, subject)
+
+    def publish_negative_match_result(self, sns_topic_arn: str) -> None:
+        """Publish a negative match result (no YARA matches found).
+
         Args:
-            sns_topic_arn: Publish match alerts to this SNS topic ARN.
+            sns_topic_arn: Target topic ARN for negative match alerts.
         """
-        LOGGER.info('Publishing an SNS alert')
-        analyzer_aws_lib.publish_safe_to_sns(self, sns_topic_arn)
+        LOGGER.info('Publishing a negative match result to %s', sns_topic_arn)
+        subject = '[BinaryAlert] {} did not match any YARA rules'.format(
+            self.filepath or self.computed_sha)
+        analyzer_aws_lib.publish_to_sns(self, sns_topic_arn, subject)
 
     def summary(self) -> Dict[str, Any]:
         """Generate a summary dictionary of binary attributes."""

lambda_functions/analyzer/main.py

@@ -1,5 +1,6 @@
 """AWS Lambda function for testing a binary against a list of YARA rules."""
 # Expects the following environment variables:
+#   NO_MATCHES_SNS_TOPIC_ARN: Optional ARN of an SNS topic to notify if there are no YARA matches.
 #   YARA_MATCHES_DYNAMO_TABLE_NAME: Name of the Dynamo table which stores YARA match results.
 #   YARA_ALERTS_SNS_TOPIC_ARN: ARN of the SNS topic which should be alerted on a YARA match.
 # Expects a binary YARA rules file to be at './compiled_yara_rules.bin'
@@ -30,14 +31,21 @@ def _objects_to_analyze(event: Dict[str, Any]) -> Generator[Tuple[str, str], Non
     Yields:
         (bucket_name, object_key) string tuples to analyze
     """
+    if 'BucketName' in event and 'ObjectKeys' in event:
+        # Direct (simple) invocation
+        for key in event['ObjectKeys']:
+            yield event['BucketName'], urllib.parse.unquote_plus(key)
+        return
+
+    # SQS message invocation
     for sqs_message in event['Records']:
         try:
-            msg_body = json.loads(sqs_message['body'])
+            s3_records = json.loads(sqs_message['body'])['Records']
         except (KeyError, TypeError, json.JSONDecodeError):
             LOGGER.exception('Skipping invalid SQS message %s', json.dumps(sqs_message))
             continue
 
-        for s3_message in msg_body['Records']:
+        for s3_message in s3_records:
             yield (
                 s3_message['s3']['bucket']['name'],
                 urllib.parse.unquote_plus(s3_message['s3']['object']['key'])
@@ -59,18 +67,25 @@ def analyze_lambda_handler(event: Dict[str, Any], lambda_context: Any) -> Dict[s
                                         'name': '...'
                                     },
                                     'object': {
-                                        'key': '...'
+                                        'key': '...'  # URL-encoded key
                                     }
                                 },
                                 ...
                             },
                             ...
                         ]
                     }),
-                    'messageId': '...'
+                    ...
                 }
             ]
         }
+
+        Alternatively, direct invocation is supported with the following event - {
+            'BucketName': '...',
+            'EnableSNSAlerts': True,
+            'ObjectKeys': ['key1', 'key2', ...],
+        }
+
         lambda_context: LambdaContext object (with .function_version).
 
     Returns:
@@ -97,6 +112,8 @@ def analyze_lambda_handler(event: Dict[str, Any], lambda_context: Any) -> Dict[s
         LOGGER.warning('Invoked $LATEST instead of a versioned function')
         lambda_version = -1
 
+    alerts_enabled = event.get('EnableSNSAlerts', True)
+
     for bucket_name, object_key in _objects_to_analyze(event):
         LOGGER.info('Analyzing "%s:%s"', bucket_name, object_key)
 
@@ -112,12 +129,12 @@ def analyze_lambda_handler(event: Dict[str, Any], lambda_context: Any) -> Dict[s
             LOGGER.warning('%s matched YARA rules: %s', binary, binary.matched_rule_ids)
             binary.save_matches_and_alert(
                 lambda_version, os.environ['YARA_MATCHES_DYNAMO_TABLE_NAME'],
-                os.environ['YARA_ALERTS_SNS_TOPIC_ARN'])
+                os.environ['YARA_ALERTS_SNS_TOPIC_ARN'],
+                sns_enabled=alerts_enabled)
         else:
             LOGGER.info('%s did not match any YARA rules', binary)
-            if os.environ['SAFE_SNS_TOPIC_ARN']:
-                binary.safe_alert_only(
-                    os.environ['SAFE_SNS_TOPIC_ARN'])
+            if alerts_enabled and os.environ['NO_MATCHES_SNS_TOPIC_ARN']:
+                binary.publish_negative_match_result(os.environ['NO_MATCHES_SNS_TOPIC_ARN'])
 
     # Publish metrics.
     if binaries:

manage.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """Command-line tool for easily managing BinaryAlert."""
 import argparse
+import os
 import sys
 
 from cli import __version__
@@ -23,6 +24,7 @@ def main() -> None:
         '--version', action='version', version='BinaryAlert v{}'.format(__version__))
     args = parser.parse_args()
 
+    os.environ['TF_IN_AUTOMATION'] = '1'
     manager.run(args.command)
 
 

terraform/cloudwatch_dashboard.tf

@@ -179,7 +179,7 @@ EOF
         "FunctionName", "${module.binaryalert_analyzer.function_name}",
         {"label": "Analyzer"}
       ]
-      ${var.enable_carbon_black_downloader == 1 ? local.downloader : ""}
+      ${var.enable_carbon_black_downloader ? local.downloader : ""}
     ]
   }
 }
@@ -199,7 +199,7 @@ EOF
         "FunctionName", "${module.binaryalert_analyzer.function_name}",
         {"label": "Analyzer"}
       ]
-      ${var.enable_carbon_black_downloader == 1 ? local.downloader : ""}
+      ${var.enable_carbon_black_downloader ? local.downloader : ""}
     ],
     "annotations": {
       "horizontal": [
@@ -227,7 +227,7 @@ EOF
         "FunctionName", "${module.binaryalert_analyzer.function_name}",
         {"label": "Analyzer"}
       ]
-      ${var.enable_carbon_black_downloader == 1 ? local.downloader : ""}
+      ${var.enable_carbon_black_downloader ? local.downloader : ""}
     ]
   }
 }
@@ -247,7 +247,7 @@ EOF
         "FunctionName", "${module.binaryalert_analyzer.function_name}",
         {"label": "Analyzer"}
       ]
-      ${var.enable_carbon_black_downloader == 1 ? local.downloader : ""}
+      ${var.enable_carbon_black_downloader ? local.downloader : ""}
     ]
   }
 }
@@ -283,7 +283,7 @@ EOF
         "TopicName", "${aws_sns_topic.yara_match_alerts.name}",
         {"label": "YARA Match Alerts"}
       ],
-      [".", ".", ".", "${aws_sns_topic.metric_alarms.name}", {"label": "Metric Alarms"}]
+      [".", ".", ".", "${element(split(":", local.alarm_target), 5)}", {"label": "Metric Alarms"}]
     ]
   }
 }
@@ -308,7 +308,7 @@ EOF
         "LogGroupName", "${module.binaryalert_analyzer.log_group_name}",
         {"label": "Analyzer"}
       ]
-      ${var.enable_carbon_black_downloader == 1 ? local.downloader_logs : ""}
+      ${var.enable_carbon_black_downloader ? local.downloader_logs : ""}
     ]
   }
 }
@@ -341,7 +341,7 @@ EOF
 }
 EOF
 
-  dashboard_body = "${var.enable_carbon_black_downloader == 1 ? local.dashboard_body_with_downloader : local.dashboard_body_without_downloader}"
+  dashboard_body = "${var.enable_carbon_black_downloader ? local.dashboard_body_with_downloader : local.dashboard_body_without_downloader}"
 }
 
 resource "aws_cloudwatch_dashboard" "binaryalert" {