text: # -*- coding: utf-8 -*-
  # Upside Travel, 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.
  
  import datetime
  import json
  import unittest
  
  import boto3
  import botocore.session
  from botocore.stub import Stubber
  
  from common import AV_SCAN_START_METADATA
  from common import AV_SIGNATURE_METADATA
  from common import AV_SIGNATURE_OK
  from common import AV_STATUS_METADATA
  from common import AV_TIMESTAMP_METADATA
  from common import get_timestamp
  from scan import delete_s3_object
  from scan import event_object
  from scan import get_local_path
  from scan import set_av_metadata
  from scan import set_av_tags
  from scan import sns_start_scan
  from scan import sns_scan_results
  from scan import verify_s3_object_version
  
  
  class TestScan(unittest.TestCase):
      def setUp(self):
          # Common data
          self.s3_bucket_name = "test_bucket"
          self.s3_key_name = "test_key"
  
          # Clients and Resources
          self.s3 = boto3.resource("s3")
          self.s3_client = botocore.session.get_session().create_client("s3")
          self.sns_client = botocore.session.get_session().create_client(
              "sns", region_name="us-west-2"
          )
  
      def test_sns_event_object(self):
          event = {
              "Records": [
                  {
                      "s3": {
                          "bucket": {"name": self.s3_bucket_name},
                          "object": {"key": self.s3_key_name},
                      }
                  }
              ]
          }
          sns_event = {"Records": [{"Sns": {"Message": json.dumps(event)}}]}
          s3_obj = event_object(sns_event, self.s3)
          expected_s3_object = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
          self.assertEquals(s3_obj, expected_s3_object)
  
      def test_s3_event_object_missing_bucket(self):
          event = {"Records": [{"s3": {"object": {"key": self.s3_key_name}}}]}
          with self.assertRaises(Exception) as cm:
              event_object(event)
              self.assertEquals(cm.exception.message, "No bucket found in event!")
  
      def test_s3_event_object_missing_key(self):
          event = {"Records": [{"s3": {"bucket": {"name": self.s3_bucket_name}}}]}
          with self.assertRaises(Exception) as cm:
              event_object(event)
              self.assertEquals(cm.exception.message, "No key found in event!")
  
      def test_s3_event_object_bucket_key_missing(self):
          event = {"Records": [{"s3": {"bucket": {}, "object": {}}}]}
          with self.assertRaises(Exception) as cm:
              event_object(event)
              self.assertEquals(
                  cm.exception.message,
                  "Unable to retrieve object from event.\n{}".format(event),
              )
  
      def test_s3_event_object_no_records(self):
          event = {"Records": []}
          with self.assertRaises(Exception) as cm:
              event_object(event)
              self.assertEquals(cm.exception.message, "No records found in event!")
  
      def test_verify_s3_object_version(self):
          s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
  
          # Set up responses
          get_bucket_versioning_response = {"Status": "Enabled"}
          get_bucket_versioning_expected_params = {"Bucket": self.s3_bucket_name}
          s3_stubber_resource = Stubber(self.s3.meta.client)
          s3_stubber_resource.add_response(
              "get_bucket_versioning",
              get_bucket_versioning_response,
              get_bucket_versioning_expected_params,
          )
          list_object_versions_response = {
              "Versions": [
                  {
                      "ETag": "string",
                      "Size": 123,
                      "StorageClass": "STANDARD",
                      "Key": "string",
                      "VersionId": "string",
                      "IsLatest": True,
                      "LastModified": datetime.datetime(2015, 1, 1),
                      "Owner": {"DisplayName": "string", "ID": "string"},
                  }
              ]
          }
          list_object_versions_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Prefix": self.s3_key_name,
          }
          s3_stubber_resource.add_response(
              "list_object_versions",
              list_object_versions_response,
              list_object_versions_expected_params,
          )
          try:
              with s3_stubber_resource:
                  verify_s3_object_version(self.s3, s3_obj)
          except Exception as e:
              self.fail("verify_s3_object_version() raised Exception unexpectedly!")
              raise e
  
      def test_verify_s3_object_versioning_not_enabled(self):
          s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
  
          # Set up responses
          get_bucket_versioning_response = {"Status": "Disabled"}
          get_bucket_versioning_expected_params = {"Bucket": self.s3_bucket_name}
          s3_stubber_resource = Stubber(self.s3.meta.client)
          s3_stubber_resource.add_response(
              "get_bucket_versioning",
              get_bucket_versioning_response,
              get_bucket_versioning_expected_params,
          )
          with self.assertRaises(Exception) as cm:
              with s3_stubber_resource:
                  verify_s3_object_version(self.s3, s3_obj)
              self.assertEquals(
                  cm.exception.message,
                  "Object versioning is not enabled in bucket {}".format(
                      self.s3_bucket_name
                  ),
              )
  
      def test_verify_s3_object_version_multiple_versions(self):
          s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
  
          # Set up responses
          get_bucket_versioning_response = {"Status": "Enabled"}
          get_bucket_versioning_expected_params = {"Bucket": self.s3_bucket_name}
          s3_stubber_resource = Stubber(self.s3.meta.client)
          s3_stubber_resource.add_response(
              "get_bucket_versioning",
              get_bucket_versioning_response,
              get_bucket_versioning_expected_params,
          )
          list_object_versions_response = {
              "Versions": [
                  {
                      "ETag": "string",
                      "Size": 123,
                      "StorageClass": "STANDARD",
                      "Key": "string",
                      "VersionId": "string",
                      "IsLatest": True,
                      "LastModified": datetime.datetime(2015, 1, 1),
                      "Owner": {"DisplayName": "string", "ID": "string"},
                  },
                  {
                      "ETag": "string",
                      "Size": 123,
                      "StorageClass": "STANDARD",
                      "Key": "string",
                      "VersionId": "string",
                      "IsLatest": True,
                      "LastModified": datetime.datetime(2015, 1, 1),
                      "Owner": {"DisplayName": "string", "ID": "string"},
                  },
              ]
          }
          list_object_versions_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Prefix": self.s3_key_name,
          }
          s3_stubber_resource.add_response(
              "list_object_versions",
              list_object_versions_response,
              list_object_versions_expected_params,
          )
          with self.assertRaises(Exception) as cm:
              with s3_stubber_resource:
                  verify_s3_object_version(self.s3, s3_obj)
              self.assertEquals(
                  cm.exception.message,
                  "Detected multiple object versions in {}.{}, aborting processing".format(
                      self.s3_bucket_name, self.s3_key_name
                  ),
              )
  
      def test_sns_start_scan(self):
          sns_stubber = Stubber(self.sns_client)
          s3_stubber_resource = Stubber(self.s3.meta.client)
  
          sns_arn = "some_arn"
          version_id = "version-id"
          timestamp = get_timestamp()
          message = {
              "bucket": self.s3_bucket_name,
              "key": self.s3_key_name,
              "version": version_id,
              AV_SCAN_START_METADATA: True,
              AV_TIMESTAMP_METADATA: timestamp,
          }
          publish_response = {"MessageId": "message_id"}
          publish_expected_params = {
              "TargetArn": sns_arn,
              "Message": json.dumps({"default": json.dumps(message)}),
              "MessageStructure": "json",
          }
          sns_stubber.add_response("publish", publish_response, publish_expected_params)
  
          head_object_response = {"VersionId": version_id}
          head_object_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
          }
          s3_stubber_resource.add_response(
              "head_object", head_object_response, head_object_expected_params
          )
          with sns_stubber, s3_stubber_resource:
              s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
              sns_start_scan(self.sns_client, s3_obj, sns_arn, timestamp)
  
      def test_get_local_path(self):
          local_prefix = "/tmp"
  
          s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
          file_path = get_local_path(s3_obj, local_prefix)
          expected_file_path = "/tmp/test_bucket/test_key"
          self.assertEquals(file_path, expected_file_path)
  
      def test_set_av_metadata(self):
          scan_result = "CLEAN"
          scan_signature = AV_SIGNATURE_OK
          timestamp = get_timestamp()
  
          s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
          s3_stubber_resource = Stubber(self.s3.meta.client)
  
          # First head call is done to get content type and meta data
          head_object_response = {"ContentType": "content", "Metadata": {}}
          head_object_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
          }
          s3_stubber_resource.add_response(
              "head_object", head_object_response, head_object_expected_params
          )
  
          # Next two calls are done when copy() is called
          head_object_response_2 = {
              "ContentType": "content",
              "Metadata": {},
              "ContentLength": 200,
          }
          head_object_expected_params_2 = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
          }
          s3_stubber_resource.add_response(
              "head_object", head_object_response_2, head_object_expected_params_2
          )
          copy_object_response = {"VersionId": "version_id"}
          copy_object_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
              "ContentType": "content",
              "CopySource": {"Bucket": self.s3_bucket_name, "Key": self.s3_key_name},
              "Metadata": {
                  AV_SIGNATURE_METADATA: scan_signature,
                  AV_STATUS_METADATA: scan_result,
                  AV_TIMESTAMP_METADATA: timestamp,
              },
              "MetadataDirective": "REPLACE",
          }
          s3_stubber_resource.add_response(
              "copy_object", copy_object_response, copy_object_expected_params
          )
  
          with s3_stubber_resource:
              set_av_metadata(s3_obj, scan_result, scan_signature, timestamp)
  
      def test_set_av_tags(self):
          scan_result = "CLEAN"
          scan_signature = AV_SIGNATURE_OK
          timestamp = get_timestamp()
          tag_set = {
              "TagSet": [
                  {"Key": AV_SIGNATURE_METADATA, "Value": scan_signature},
                  {"Key": AV_STATUS_METADATA, "Value": scan_result},
                  {"Key": AV_TIMESTAMP_METADATA, "Value": timestamp},
              ]
          }
  
          s3_stubber = Stubber(self.s3_client)
          get_object_tagging_response = tag_set
          get_object_tagging_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
          }
          s3_stubber.add_response(
              "get_object_tagging",
              get_object_tagging_response,
              get_object_tagging_expected_params,
          )
          put_object_tagging_response = {}
          put_object_tagging_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
              "Tagging": tag_set,
          }
          s3_stubber.add_response(
              "put_object_tagging",
              put_object_tagging_response,
              put_object_tagging_expected_params,
          )
  
          with s3_stubber:
              s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
              set_av_tags(self.s3_client, s3_obj, scan_result, scan_signature, timestamp)
  
      def test_sns_scan_results(self):
          sns_stubber = Stubber(self.sns_client)
          s3_stubber_resource = Stubber(self.s3.meta.client)
  
          status_sns_arn = "some_status_arn"
          version_id = "version-id"
          scan_result = "CLEAN"
          scan_signature = AV_SIGNATURE_OK
          timestamp = get_timestamp()
          message = {
              "bucket": self.s3_bucket_name,
              "key": self.s3_key_name,
              "version": version_id,
              AV_SIGNATURE_METADATA: scan_signature,
              AV_STATUS_METADATA: scan_result,
              AV_TIMESTAMP_METADATA: timestamp,
          }
          publish_response = {"MessageId": "message_id"}
          publish_expected_status_params = {
              "TargetArn": status_sns_arn,
              "Message": json.dumps({"default": json.dumps(message)}),
              "MessageAttributes": {
                  "av-status": {"DataType": "String", "StringValue": scan_result},
                  "av-signature": {"DataType": "String", "StringValue": scan_signature},
              },
              "MessageStructure": "json",
          }
          sns_stubber.add_response(
              "publish", publish_response, publish_expected_status_params
          )
  
          head_object_response = {"VersionId": version_id}
          head_object_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
          }
          s3_stubber_resource.add_response(
              "head_object", head_object_response, head_object_expected_params
          )
          with sns_stubber, s3_stubber_resource:
              s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
              sns_scan_results(
                  self.sns_client,
                  s3_obj,
                  status_sns_arn,
                  scan_result,
                  scan_signature,
                  timestamp,
              )
  
      def test_delete_s3_object(self):
          s3_stubber = Stubber(self.s3.meta.client)
          delete_object_response = {}
          delete_object_expected_params = {
              "Bucket": self.s3_bucket_name,
              "Key": self.s3_key_name,
          }
          s3_stubber.add_response(
              "delete_object", delete_object_response, delete_object_expected_params
          )
  
          with s3_stubber:
              s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
              delete_s3_object(s3_obj)
  
      def test_delete_s3_object_exception(self):
          s3_stubber = Stubber(self.s3.meta.client)
  
          with self.assertRaises(Exception) as cm:
              with s3_stubber:
                  s3_obj = self.s3.Object(self.s3_bucket_name, self.s3_key_name)
                  delete_s3_object(s3_obj)
              self.assertEquals(
                  cm.exception.message,
                  "Failed to delete infected file: {}.{}".format(
                      self.s3_bucket_name, self.s3_key_name
                  ),
              )
  