diff --git a/README.md b/README.md index 23a60b59..6d911919 100644 --- a/README.md +++ b/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 diff --git a/clamav.py b/clamav.py index 8042e739..a44ab3a2 100644 --- a/clamav.py +++ b/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( diff --git a/common.py b/common.py index 9e95af96..bb953fca 100644 --- a/common.py +++ b/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): diff --git a/display_infected.py b/display_infected.py index 0c40bc98..b80e1347 100755 --- a/display_infected.py +++ b/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: diff --git a/scan.py b/scan.py index c55716d4..48545a06 100644 --- a/scan.py +++ b/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", "") diff --git a/scan_bucket.py b/scan_bucket.py index 6043ffb0..7c6367fb 100755 --- a/scan_bucket.py +++ b/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,7 +88,7 @@ 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: @@ -95,7 +96,7 @@ def main(lambda_function_name, s3_bucket_name, limit): 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: diff --git a/update.py b/update.py index 9730f230..80aa46d1 100644 --- a/update.py +++ b/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(