Skip to content

Commit

Permalink
Make acquire more robust against missing collection profiles
Browse files Browse the repository at this point in the history
  • Loading branch information
pyrco committed Nov 4, 2024
1 parent 910e364 commit 1266399
Showing 1 changed file with 22 additions and 20 deletions.
42 changes: 22 additions & 20 deletions acquire/acquire.py
Original file line number Diff line number Diff line change
Expand Up @@ -1696,17 +1696,20 @@ def print_acquire_warning(target: Target) -> None:
log.warning("========================================== WARNING ==========================================")


def _add_modules_for_profile(choice: str, operating_system: str, profile: dict, msg: str) -> Optional[dict]:
modules_selected = dict()

if choice and choice != "none":
profile_dict = profile[choice]
if operating_system not in profile_dict:
log.error(msg, operating_system, choice)
return None
def _get_modules_for_profile(
profile_name: str,
operating_system: str,
profiles: dict[str, dict[str, list[type[Module]]]],
err_msg: str,
) -> dict[str, type[Module]]:
modules_selected = {}

Check warning on line 1705 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1705

Added line #L1705 was not covered by tests

for mod in profile_dict[operating_system]:
modules_selected[mod.__modname__] = mod
if profile_name != "none":
if (profile := profiles.get(profile_name, {}).get(operating_system)) is not None:
for mod in profile:
modules_selected[mod.__modname__] = mod

Check warning on line 1710 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1707-L1710

Added lines #L1707 - L1710 were not covered by tests
else:
log.error(err_msg, operating_system, profile_name)

Check warning on line 1712 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1712

Added line #L1712 was not covered by tests

return modules_selected

Expand Down Expand Up @@ -1775,24 +1778,23 @@ def acquire_target(target: Target, args: argparse.Namespace, output_ts: Optional
profile = "default"
log.info("")

profile_modules = _add_modules_for_profile(
profile, target.os, PROFILES, "No collection set for OS %s with profile %s"
normal_modules = _get_modules_for_profile(

Check warning on line 1781 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1781

Added line #L1781 was not covered by tests
profile, target.os, PROFILES, "No collection set for OS '%s' with profile '%s'"
)
modules_selected.update(normal_modules)

Check warning on line 1784 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1784

Added line #L1784 was not covered by tests

if not (volatile_profile := args.volatile_profile):
volatile_profile = "none"

volatile_modules = _add_modules_for_profile(
volatile_profile, target.os, VOLATILE, "No collection set for OS %s with volatile profile %s"
volatile_modules = _get_modules_for_profile(

Check warning on line 1789 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1789

Added line #L1789 was not covered by tests
volatile_profile, target.os, VOLATILE, "No collection set for OS '%s' with volatile profile '%s'"
)

if (profile_modules or volatile_modules) is None:
return files

modules_selected.update(profile_modules)
modules_selected.update(volatile_modules)

log.info("Modules selected: %s", ", ".join(sorted(modules_selected)))
if not modules_selected:
log.warn("NO modules selected!")

Check warning on line 1795 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1794-L1795

Added lines #L1794 - L1795 were not covered by tests
else:
log.info("Modules selected: %s", ", ".join(sorted(modules_selected)))

Check warning on line 1797 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1797

Added line #L1797 was not covered by tests

local_only_modules = {name: module for name, module in modules_selected.items() if hasattr(module, "__local__")}
if target.path.name != "local" and local_only_modules:
Expand Down

0 comments on commit 1266399

Please sign in to comment.