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",