Merge pull request bluesentry#167 from dmarkey/master

Add configurable endpoints for AWS services.

README.md

@@ -103,6 +103,9 @@ the table below for reference.
 | AV_PROCESS_ORIGINAL_VERSION_ONLY | Controls that only original version of an S3 key is processed (if bucket versioning is enabled) | False | No |
 | AV_DELETE_INFECTED_FILES | Controls whether infected files should be automatically deleted | False | No |
 | EVENT_SOURCE | The source of antivirus scan event "S3" or "SNS" (optional) | S3 | No |
+| S3_ENDPOINT | The Endpoint to use when interacting wth S3 | None | No |
+| SNS_ENDPOINT | The Endpoint to use when interacting wth SNS | None | No |
+| LAMBDA_ENDPOINT | The Endpoint to use when interacting wth Lambda | None | No |
 
 ## S3 Bucket Policy Examples
 

clamav.py

@@ -24,7 +24,7 @@
 import botocore
 from pytz import utc
 
-from common import AV_DEFINITION_S3_PREFIX
+from common import AV_DEFINITION_S3_PREFIX, S3_ENDPOINT
 from common import AV_DEFINITION_PATH
 from common import AV_DEFINITION_FILE_PREFIXES
 from common import AV_DEFINITION_FILE_SUFFIXES
@@ -90,7 +90,7 @@ def upload_defs_to_s3(s3_client, bucket, prefix, local_path):
                         "Uploading %s to s3://%s"
                         % (local_file_path, os.path.join(bucket, prefix, filename))
                     )
-                    s3 = boto3.resource("s3")
+                    s3 = boto3.resource("s3", endpoint_url=S3_ENDPOINT)
                     s3_object = s3.Object(bucket, os.path.join(prefix, filename))
                     s3_object.upload_file(os.path.join(local_path, filename))
                     s3_client.put_object_tagging(

common.py

@@ -43,6 +43,9 @@
 
 AV_DEFINITION_FILE_PREFIXES = ["main", "daily", "bytecode"]
 AV_DEFINITION_FILE_SUFFIXES = ["cld", "cvd"]
+SNS_ENDPOINT = os.getenv("SNS_ENDPOINT", None)
+S3_ENDPOINT = os.getenv("S3_ENDPOINT", None)
+LAMBDA_ENDPOINT = os.getenv("LAMBDA_ENDPOINT", None)
 
 
 def create_dir(path):

display_infected.py

@@ -20,7 +20,7 @@
 
 import boto3
 
-from common import AV_SIGNATURE_METADATA
+from common import AV_SIGNATURE_METADATA, S3_ENDPOINT
 from common import AV_SIGNATURE_OK
 from common import AV_SIGNATURE_UNKNOWN
 from common import AV_STATUS_METADATA
@@ -78,7 +78,7 @@ def object_infected(s3_client, s3_bucket_name, key_name):
 def main(s3_bucket_name):
 
     # Verify the S3 bucket exists
-    s3_client = boto3.client("s3")
+    s3_client = boto3.client("s3", endpoint_url=S3_ENDPOINT)
     try:
         s3_client.head_bucket(Bucket=s3_bucket_name)
     except Exception:

scan.py

@@ -37,6 +37,8 @@
 from common import AV_STATUS_SNS_PUBLISH_CLEAN
 from common import AV_STATUS_SNS_PUBLISH_INFECTED
 from common import AV_TIMESTAMP_METADATA
+from common import SNS_ENDPOINT
+from common import S3_ENDPOINT
 from common import create_dir
 from common import get_timestamp
 
@@ -73,7 +75,7 @@ def event_object(event, event_source="s3"):
         raise Exception("Unable to retrieve object from event.\n{}".format(event))
 
     # Create and return the object
-    s3 = boto3.resource("s3")
+    s3 = boto3.resource("s3", endpoint_url=S3_ENDPOINT)
     return s3.Object(bucket_name, key_name)
 
 
@@ -199,9 +201,9 @@ def sns_scan_results(
 
 
 def lambda_handler(event, context):
-    s3 = boto3.resource("s3")
-    s3_client = boto3.client("s3")
-    sns_client = boto3.client("sns")
+    s3 = boto3.resource("s3", endpoint_url=S3_ENDPOINT)
+    s3_client = boto3.client("s3", endpoint_url=S3_ENDPOINT)
+    sns_client = boto3.client("sns", endpoint_url=SNS_ENDPOINT)
 
     # Get some environment variables
     ENV = os.getenv("ENV", "")

scan_bucket.py

@@ -21,8 +21,9 @@
 
 import boto3
 
-from common import AV_STATUS_METADATA
+from common import AV_STATUS_METADATA, LAMBDA_ENDPOINT
 from common import AV_TIMESTAMP_METADATA
+from common import S3_ENDPOINT
 
 
 # Get all objects in an S3 bucket that have not been previously scanned
@@ -87,15 +88,15 @@ def format_s3_event(s3_bucket_name, key_name):
 
 def main(lambda_function_name, s3_bucket_name, limit):
     # Verify the lambda exists
-    lambda_client = boto3.client("lambda")
+    lambda_client = boto3.client("lambda", endpoint_url=LAMBDA_ENDPOINT)
     try:
         lambda_client.get_function(FunctionName=lambda_function_name)
     except Exception:
         print("Lambda Function '{}' does not exist".format(lambda_function_name))
         sys.exit(1)
 
     # Verify the S3 bucket exists
-    s3_client = boto3.client("s3")
+    s3_client = boto3.client("s3", endpoint_url=S3_ENDPOINT)
     try:
         s3_client.head_bucket(Bucket=s3_bucket_name)
     except Exception:

update.py

@@ -22,12 +22,13 @@
 from common import AV_DEFINITION_S3_BUCKET
 from common import AV_DEFINITION_S3_PREFIX
 from common import CLAMAVLIB_PATH
+from common import S3_ENDPOINT
 from common import get_timestamp
 
 
 def lambda_handler(event, context):
-    s3 = boto3.resource("s3")
-    s3_client = boto3.client("s3")
+    s3 = boto3.resource("s3", endpoint_url=S3_ENDPOINT)
+    s3_client = boto3.client("s3", endpoint_url=S3_ENDPOINT)
 
     print("Script starting at %s\n" % (get_timestamp()))
     to_download = clamav.update_defs_from_s3(