def test_since_filtered_multi_bucket(self, versions_page_size, dry_run, concurrency): av_api_client, s3_client = self._get_mock_clients( self.buckets_versions_responses, versions_page_size) with ThreadPoolExecutor( max_workers=concurrency) if concurrency else nullcontext( ) as executor: map_callable = map if executor is None else executor.map retval = virus_scan_bucket( s3_client, av_api_client, ( "spade", "martello", ), prefix="", since=datetime(2012, 11, 10, 9, 8, 7), dry_run=dry_run, map_callable=map_callable, ) assert s3_client.mock_calls == [ mock.call.get_paginator("list_object_versions"), mock.call.get_paginator().paginate(Bucket="spade", Prefix=""), mock.call.get_paginator("list_object_versions"), mock.call.get_paginator().paginate(Bucket="martello", Prefix=""), ] if dry_run: assert av_api_client.mock_calls == [] assert retval == Counter({"candidate": 6}) else: # taking string representations because call()s are not sortable and we want to disregard order assert sorted(str(c) for c in av_api_client.mock_calls) == sorted( str(c) for c in ( mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "oo_.BepoodlLml"), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "ooBmo_pe.ldoLl"), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "molo.oB_oLdelp"), mock.call.scan_and_tag_s3_object( "martello", "unmentionables.PNG", "lFHrwenroye_"), mock.call.scan_and_tag_s3_object( "martello", "sandy/mount.pdf", "nHwr_elFoyre"), mock.call.scan_and_tag_s3_object( "martello", "handy/mount.pdf", "Hn_olFerweyr"), )) assert retval == Counter({ "candidate": 6, "pass": 1, "fail": 2, "already_tagged": 2, "error": 1, })
def test_unfiltered_single_bucket(self, versions_page_size, dry_run, concurrency): av_api_client, s3_client = self._get_mock_clients( self.buckets_versions_responses, versions_page_size) with ThreadPoolExecutor( max_workers=concurrency) if concurrency else nullcontext( ) as executor: map_callable = map if executor is None else executor.map retval = virus_scan_bucket( s3_client, av_api_client, ("spade", ), prefix="", since=None, dry_run=dry_run, map_callable=map_callable, ) assert s3_client.mock_calls == [ mock.call.get_paginator("list_object_versions"), mock.call.get_paginator().paginate(Bucket="spade", Prefix=""), ] if dry_run: assert av_api_client.mock_calls == [] assert retval == Counter({"candidate": 7}) else: # taking string representations because call()s are not sortable and we want to disregard order assert sorted(str(c) for c in av_api_client.mock_calls) == sorted( str(c) for c in ( mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "oo_.BepoodlLml"), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "moB_eLplool.do"), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "ooBmo_pe.ldoLl"), mock.call.scan_and_tag_s3_object( "spade", "sandman/1234-deedaw.pdf", "epmlLoBodo_ol."), mock.call.scan_and_tag_s3_object( "spade", "sandman/1234-deedaw.pdf", "loleLoooBp_md."), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "molo.oB_oLdelp"), mock.call.scan_and_tag_s3_object( "spade", "dribbling/bib.jpeg", "ldmoo_.pBeolLo"), )) assert retval == Counter({ "candidate": 7, "pass": 4, "fail": 1, "already_tagged": 2, })
def test_prefix_filtered_multi_bucket(self, versions_page_size, dry_run, concurrency): av_api_client, s3_client = self._get_mock_clients( { bucket_name: tuple(v_r for v_r in versions_responses if v_r[0]["Key"].startswith("sand")) for bucket_name, versions_responses in self.buckets_versions_responses.items() }, versions_page_size, ) with ThreadPoolExecutor( max_workers=concurrency) if concurrency else nullcontext( ) as executor: map_callable = map if executor is None else executor.map retval = virus_scan_bucket( s3_client, av_api_client, ( "spade", "martello", ), prefix="sand", since=None, dry_run=dry_run, map_callable=map_callable, ) assert s3_client.mock_calls == [ mock.call.get_paginator("list_object_versions"), mock.call.get_paginator().paginate(Bucket="spade", Prefix="sand"), mock.call.get_paginator("list_object_versions"), mock.call.get_paginator().paginate(Bucket="martello", Prefix="sand"), ] if dry_run: assert av_api_client.mock_calls == [] assert retval == Counter({"candidate": 7}) else: # taking string representations because call()s are not sortable and we want to disregard order assert sorted(str(c) for c in av_api_client.mock_calls) == sorted( str(c) for c in ( mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "oo_.BepoodlLml"), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "moB_eLplool.do"), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "ooBmo_pe.ldoLl"), mock.call.scan_and_tag_s3_object( "spade", "sandman/1234-deedaw.pdf", "epmlLoBodo_ol."), mock.call.scan_and_tag_s3_object( "spade", "sandman/1234-deedaw.pdf", "loleLoooBp_md."), mock.call.scan_and_tag_s3_object( "spade", "sandman/4321-billy-winks.pdf", "molo.oB_oLdelp"), mock.call.scan_and_tag_s3_object( "martello", "sandy/mount.pdf", "nHwr_elFoyre"), )) assert retval == Counter({ "candidate": 7, "pass": 3, "fail": 1, "already_tagged": 3, })
"Configuration:\nTarget stage:\t%s\nTarget buckets:\t%s\nFilter prefix:\t%s\nModified since:\t%s\nDry run:\t%s", args.stage, args.buckets, args.prefix, args.since, args.dry_run, ) with ThreadPoolExecutor(max_workers=args.concurrency) if args.concurrency else nullcontext() as executor: map_callable = map if executor is None else executor.map try: counter = virus_scan_bucket( s3_client=boto3.client("s3", region_name="eu-west-1"), # actual region specified here doesn't matter antivirus_api_client=av_api_client, bucket_names=args.buckets.split(","), prefix=args.prefix, since=args.since, dry_run=args.dry_run, map_callable=map_callable, ) except Exception: if executor is not None: executor.shutdown(wait=False) raise logger.info( "Total files found:\t%s\n" "Total files passed:\t%s\n" "Total files failed:\t%s\n" "Total files already tagged:\t%s\n" "Total files errored:\t%s",