diff --git a/core/engine.py b/core/engine.py index deb1aa7..3910a65 100644 --- a/core/engine.py +++ b/core/engine.py @@ -254,7 +254,7 @@ def create_resource_inventory( raw_data_path: str, ) -> dict[str, Any]: # Copy assets and datasets folders data - copy_assets(report_path) + copy_assets(report_path, cloud_service_provider) try: diff --git a/core/utils.py b/core/utils.py index 399882f..6d8ea38 100644 --- a/core/utils.py +++ b/core/utils.py @@ -5,15 +5,44 @@ logger = logging.getLogger("core.engine.utils") +# Icon folder per cloud service provider id +PROVIDER_ICON_DIRS = {1: "azure", 2: "aws"} -def copy_assets(report_path: str) -> None: - assets_folders = ["css", "icons", "img"] +# Icon folders every report needs, whichever provider was assessed +SHARED_ICON_DIRS = ("severity", "misc") + + +def _png_only(src: str, names: list[str]) -> list[str]: + # Keep directories so copytree still walks the whole tree, drop every + # file that isn't a PNG. The renderers never load an icon SVG. + return [ + name + for name in names + if not os.path.isdir(os.path.join(src, name)) + and not name.lower().endswith(".png") + ] + + +def _icon_dirs(cloud_service_provider: int) -> list[str]: + provider_dir = PROVIDER_ICON_DIRS.get(cloud_service_provider) + + if provider_dir is None: + logger.warning( + "Unknown cloud service provider %s, copying every provider icon set", + cloud_service_provider, + ) + return [*PROVIDER_ICON_DIRS.values(), *SHARED_ICON_DIRS] + + return [provider_dir, *SHARED_ICON_DIRS] + + +def copy_assets(report_path: str, cloud_service_provider: int) -> None: assets_path = os.path.join(report_path, "assets") # Create the 'assets' directory if it doesn't exist os.makedirs(assets_path, exist_ok=True) - for folder in assets_folders: + for folder in ("css", "img"): src_path = os.path.join("assets", folder) dest_path = os.path.join(assets_path, folder) @@ -21,6 +50,15 @@ def copy_assets(report_path: str) -> None: if not os.path.exists(dest_path): shutil.copytree(src_path, dest_path, dirs_exist_ok=True) + # Only the assessed provider's icons travel with the report + for folder in _icon_dirs(cloud_service_provider): + src_path = os.path.join("assets", "icons", folder) + dest_path = os.path.join(assets_path, "icons", folder) + + # Only copy if the destination doesn't already exist + if not os.path.exists(dest_path): + shutil.copytree(src_path, dest_path, ignore=_png_only, dirs_exist_ok=True) + # Copy datasets/data.db to data/assessment.db db_src_path = "datasets/data.db" db_dest_dir = os.path.join(report_path, "data") diff --git a/tests/test_copy_assets.py b/tests/test_copy_assets.py new file mode 100644 index 0000000..07cc1bc --- /dev/null +++ b/tests/test_copy_assets.py @@ -0,0 +1,124 @@ +import os +import shutil +import tempfile +import unittest +from contextlib import contextmanager +from pathlib import Path +from unittest import mock + +from core.utils import copy_assets + +SOURCE_ICONS = Path("assets/icons") + + +@contextmanager +def staged(cloud_service_provider): + # datasets/data.db is downloaded at runtime and is absent in CI, so the + # assessment.db copy is stubbed out -- it isn't what these tests cover. + # Replace the module binding inside core.utils rather than the attribute + # on shutil itself: copytree calls copyfile internally, so patching it + # globally would break the asset copy this module is testing. + fake_shutil = mock.MagicMock(wraps=shutil) + fake_shutil.copyfile = mock.MagicMock() + + with tempfile.TemporaryDirectory() as report_dir: + with mock.patch("core.utils.shutil", fake_shutil): + copy_assets(report_dir, cloud_service_provider) + + yield Path(report_dir), fake_shutil.copyfile + + +def icon_dirs(report_path): + return {p.name for p in (report_path / "assets" / "icons").iterdir() if p.is_dir()} + + +def relative_pngs(root): + return {p.relative_to(root) for p in root.rglob("*.png")} + + +class CopyAssetsProviderScopeTests(unittest.TestCase): + def test_azure_assessment_copies_only_azure_icons(self): + with staged(1) as (report_path, _): + self.assertEqual(icon_dirs(report_path), {"azure", "severity", "misc"}) + + def test_aws_assessment_copies_only_aws_icons(self): + with staged(2) as (report_path, _): + self.assertEqual(icon_dirs(report_path), {"aws", "severity", "misc"}) + + def test_unknown_provider_falls_back_to_every_icon_set(self): + with staged(99) as (report_path, _): + self.assertEqual( + icon_dirs(report_path), {"azure", "aws", "severity", "misc"} + ) + + def test_provider_icon_set_is_copied_in_full(self): + # The filter must drop non-PNG files only, never an icon. + for cloud_service_provider, provider in ((1, "azure"), (2, "aws")): + with self.subTest(provider=provider): + with staged(cloud_service_provider) as (report_path, _): + self.assertEqual( + relative_pngs(report_path / "assets" / "icons" / provider), + relative_pngs(SOURCE_ICONS / provider), + ) + + +class CopyAssetsFileTypeTests(unittest.TestCase): + def test_no_svg_is_copied(self): + for cloud_service_provider in (1, 2): + with self.subTest(cloud_service_provider=cloud_service_provider): + with staged(cloud_service_provider) as (report_path, _): + self.assertEqual( + list((report_path / "assets" / "icons").rglob("*.svg")), [] + ) + + def test_only_png_files_are_copied(self): + with staged(1) as (report_path, _): + non_png = [ + p + for p in (report_path / "assets" / "icons").rglob("*") + if p.is_file() and p.suffix.lower() != ".png" + ] + + self.assertEqual(non_png, []) + + def test_category_subfolders_are_preserved(self): + with staged(2) as (report_path, _): + copied = report_path / "assets" / "icons" / "aws" + + self.assertEqual( + {p.name for p in copied.iterdir() if p.is_dir()}, + {p.name for p in (SOURCE_ICONS / "aws").iterdir() if p.is_dir()}, + ) + + +class CopyAssetsUnrelatedAssetTests(unittest.TestCase): + def test_css_and_img_are_still_copied(self): + with staged(1) as (report_path, _): + self.assertTrue((report_path / "assets" / "css").is_dir()) + self.assertTrue((report_path / "assets" / "img").is_dir()) + + def test_assessment_db_is_still_copied(self): + with staged(1) as (report_path, copyfile): + self.assertTrue((report_path / "data").is_dir()) + copyfile.assert_called_once_with( + "datasets/data.db", + os.path.join(str(report_path), "data", "assessment.db"), + ) + + def test_shared_icons_are_copied_for_every_provider(self): + # no_image.png is the fallback icon and the severity icons are used by + # the PDF renderer, so both must travel with any report. + for cloud_service_provider in (1, 2): + with self.subTest(cloud_service_provider=cloud_service_provider): + with staged(cloud_service_provider) as (report_path, _): + icons = report_path / "assets" / "icons" + + self.assertTrue((icons / "misc" / "no_image.png").is_file()) + for severity in ("high", "medium", "low"): + self.assertTrue( + (icons / "severity" / f"{severity}.png").is_file() + ) + + +if __name__ == "__main__": + unittest.main()