Skip to content

Improve logging successful and failed acquire runs #177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
218 changes: 128 additions & 90 deletions acquire/acquire.py
Original file line number Diff line number Diff line change
Expand Up @@ -1658,7 +1658,7 @@
return modules_selected


def acquire_target(target: Target, args: argparse.Namespace, output_ts: Optional[str] = None) -> list[str]:
def acquire_target(target: Target, args: argparse.Namespace, output_ts: Optional[str] = None) -> list[str | Path]:
acquire_gui = GUI()
files = []
output_ts = output_ts or get_utc_now_str()
Expand Down Expand Up @@ -1866,16 +1866,18 @@
return files


def upload_files(paths: list[Path], upload_plugin: UploaderPlugin, no_proxy: bool = False) -> None:
def upload_files(paths: list[str | Path], upload_plugin: UploaderPlugin, no_proxy: bool = False) -> None:
proxies = None if no_proxy else urllib.request.getproxies()
log.debug("Proxies: %s (no_proxy = %s)", proxies, no_proxy)

log.info('Uploading files: "%s"', " ".join(map(str, paths)))

Check warning on line 1873 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1873

Added line #L1873 was not covered by tests
try:
upload_files_using_uploader(upload_plugin, paths, proxies)
except Exception:
log.error("Upload %s FAILED. See log file for details.", paths)
GUI().message("Upload failed.")
log.exception("")
log.error('Upload FAILED for files: "%s". See log file for details.', " ".join(map(str, paths)))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is the GUI message removed here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because it is already displayed in main() I didn't want to add extra GUI dependencies in these functions.

raise

Check warning on line 1878 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1877-L1878

Added lines #L1877 - L1878 were not covered by tests
else:
log.info("Upload succeeded.")

Check warning on line 1880 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L1880

Added line #L1880 was not covered by tests


class WindowsProfile:
Expand Down Expand Up @@ -2036,29 +2038,30 @@
}


def exit_success(default_args: list[str]):
log.info("Acquire finished successful")
log.info("Arguments: %s", " ".join(sys.argv[1:]))
log.info("Default Arguments: %s", " ".join(default_args))
log.info("Exiting with status code 0 (SUCCESS)")
sys.exit(0)

Check warning on line 2046 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2042-L2046

Added lines #L2042 - L2046 were not covered by tests


def exit_failure(default_args: list[str]):
log.error("Acquire FAILED")
log.error("Arguments: %s", " ".join(sys.argv[1:]))
log.error("Default Arguments: %s", " ".join(default_args))
log.error("Exiting with status code 1 (FAILURE)")
sys.exit(1)

Check warning on line 2054 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2050-L2054

Added lines #L2050 - L2054 were not covered by tests


def main() -> None:
parser = create_argument_parser(PROFILES, VOLATILE, MODULES)
args, rest = parse_acquire_args(parser, config=CONFIG)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does the GUI still work?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes it still works :)

# start GUI if requested through CLI / config
flavour = None
if args.gui == "always" or (
args.gui == "depends" and os.environ.get("PYS_KEYSOURCE") == "prompt" and len(sys.argv) == 1
):
flavour = platform.system()

acquire_gui = GUI(flavour=flavour, upload_available=args.auto_upload)
args.output, args.auto_upload, cancel = acquire_gui.wait_for_start(args)

# Since output has a default value, set it to None when output_file is defined
if args.output_file:
args.output = None

if cancel:
parser.exit(0)

# From here onwards, the GUI will be locked and cannot be closed because we're acquiring

try:
check_and_set_log_args(args)
except ValueError as err:
Expand All @@ -2075,65 +2078,118 @@

setup_logging(log, log_file, args.verbose, delay=args.log_delay)

log.info(ACQUIRE_BANNER)
log.info("User: %s | Admin: %s", get_user_name(), is_user_admin())
log.info("Arguments: %s", " ".join(sys.argv[1:]))
log.info("Default Arguments: %s", " ".join(args.config.get("arguments")))
log.info("")
acquire_successful = True
files_to_upload = [log_file]
acquire_gui = None
try:
log.info(ACQUIRE_BANNER)
log.info("User: %s | Admin: %s", get_user_name(), is_user_admin())
log.info("Arguments: %s", " ".join(sys.argv[1:]))
log.info("Default Arguments: %s", " ".join(args.config.get("arguments")))
log.info("")

Check warning on line 2089 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2081-L2089

Added lines #L2081 - L2089 were not covered by tests

# start GUI if requested through CLI / config
flavour = None
if args.gui == "always" or (

Check warning on line 2093 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2092-L2093

Added lines #L2092 - L2093 were not covered by tests
args.gui == "depends" and os.environ.get("PYS_KEYSOURCE") == "prompt" and len(sys.argv) == 1
):
flavour = platform.system()
acquire_gui = GUI(flavour=flavour, upload_available=args.auto_upload)

Check warning on line 2097 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2096-L2097

Added lines #L2096 - L2097 were not covered by tests

plugins_to_load = [("cloud", MinIO)]
upload_plugins = UploaderRegistry("acquire.plugins", plugins_to_load)
args.output, args.auto_upload, cancel = acquire_gui.wait_for_start(args)
if cancel:
log.info("Acquire cancelled")
exit_success(args.config.get("arguments"))

Check warning on line 2102 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2099-L2102

Added lines #L2099 - L2102 were not covered by tests
# From here onwards, the GUI will be locked and cannot be closed because we're acquiring

plugins_to_load = [("cloud", MinIO)]
upload_plugins = UploaderRegistry("acquire.plugins", plugins_to_load)

Check warning on line 2106 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2105-L2106

Added lines #L2105 - L2106 were not covered by tests

try:
check_and_set_acquire_args(args, upload_plugins)
except ValueError as err:
log.exception(err)
parser.exit(1)

if args.upload:
if args.upload:
try:
upload_files(args.upload, args.upload_plugin, args.no_proxy)
except Exception as err:
acquire_gui.message("Failed to upload files")
log.exception(err)
exit_failure(args.config.get("arguments"))
exit_success(args.config.get("arguments"))

Check warning on line 2117 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2110-L2117

Added lines #L2110 - L2117 were not covered by tests

target_paths = []
for target_path in args.targets:
target_path = args_to_uri([target_path], args.loader, rest)[0] if args.loader else target_path
if target_path == "local":
target_query = {}
if args.force_fallback:
target_query.update({"force-directory-fs": 1})

Check warning on line 2125 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2119-L2125

Added lines #L2119 - L2125 were not covered by tests

if args.fallback:
target_query.update({"fallback-to-directory-fs": 1})

Check warning on line 2128 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2127-L2128

Added lines #L2127 - L2128 were not covered by tests

target_query = urllib.parse.urlencode(target_query)
target_path = f"{target_path}?{target_query}"
target_paths.append(target_path)

Check warning on line 2132 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2130-L2132

Added lines #L2130 - L2132 were not covered by tests

try:
upload_files(args.upload, args.upload_plugin, args.no_proxy)
target_name = "Unknown" # just in case open_all already fails
for target in Target.open_all(target_paths):
target_name = "Unknown" # overwrite previous target name
target_name = target.name
log.info("Loading target %s", target_name)
log.info(target)
if target.os == "esxi" and target.name == "local":

Check warning on line 2141 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2135-L2141

Added lines #L2135 - L2141 were not covered by tests
# Loader found that we are running on an esxi host
# Perform operations to "enhance" memory
with esxi_memory_context_manager():
files_to_upload = acquire_children_and_targets(target, args)

Check warning on line 2145 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2144-L2145

Added lines #L2144 - L2145 were not covered by tests
else:
files_to_upload = acquire_children_and_targets(target, args)

Check warning on line 2147 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2147

Added line #L2147 was not covered by tests
except Exception:
log.exception("Failed to upload files")
return

target_paths = []
for target_path in args.targets:
target_path = args_to_uri([target_path], args.loader, rest)[0] if args.loader else target_path
if target_path == "local":
target_query = {}
if args.force_fallback:
target_query.update({"force-directory-fs": 1})
log.error("Failed to acquire target: %s", target_name)
if not is_user_admin():
log.error("Try re-running as administrator/root")
acquire_gui.message("This application must be run as administrator.")
raise

Check warning on line 2153 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2149-L2153

Added lines #L2149 - L2153 were not covered by tests

if args.fallback:
target_query.update({"fallback-to-directory-fs": 1})
files_to_upload = sort_files(files_to_upload)

Check warning on line 2155 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2155

Added line #L2155 was not covered by tests

target_query = urllib.parse.urlencode(target_query)
target_path = f"{target_path}?{target_query}"
target_paths.append(target_path)
except Exception as err:
log.error("Acquiring artifacts FAILED")
log.exception(err)
acquire_successful = False

Check warning on line 2160 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2157-L2160

Added lines #L2157 - L2160 were not covered by tests
else:
log.info("Acquiring artifacts succeeded")

Check warning on line 2162 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2162

Added line #L2162 was not covered by tests

try:
target_name = "Unknown" # just in case open_all already fails
for target in Target.open_all(target_paths):
target_name = "Unknown" # overwrite previous target name
target_name = target.name
log.info("Loading target %s", target_name)
log.info(target)
if target.os == "esxi" and target.name == "local":
# Loader found that we are running on an esxi host
# Perform operations to "enhance" memory
with esxi_memory_context_manager():
acquire_children_and_targets(target, args)
else:
acquire_children_and_targets(target, args)
except Exception:
if not is_user_admin():
log.error("Failed to load target: %s, try re-running as administrator/root", target_name)
acquire_gui.message("This application must be run as administrator.")
# The auto-upload of files is done at the very very end to make sure any
# logged exceptions are written to the log file before uploading.
# This means that any failures from this point on will not be part of the
# uploaded log files, they will be written to the logfile on disk though.
if args.auto_upload and args.upload_plugin and files_to_upload:
try:
log_file_handler = get_file_handler(log)
if log_file_handler:
log_file_handler.close()

Check warning on line 2173 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2169-L2173

Added lines #L2169 - L2173 were not covered by tests

upload_files(files_to_upload, args.upload_plugin)
except Exception:
if acquire_gui:
acquire_gui.message("Failed to upload files")
raise

Check warning on line 2179 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2175-L2179

Added lines #L2175 - L2179 were not covered by tests

if acquire_gui:
acquire_gui.finish()

Check warning on line 2182 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2181-L2182

Added lines #L2181 - L2182 were not covered by tests
acquire_gui.wait_for_quit()
parser.exit(1)
log.exception("Failed to load target: %s", target_name)
raise

except Exception as err:
acquire_successful = False
log.exception(err)

Check warning on line 2187 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2185-L2187

Added lines #L2185 - L2187 were not covered by tests

if acquire_successful:
exit_success(args.config.get("arguments"))

Check warning on line 2190 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2189-L2190

Added lines #L2189 - L2190 were not covered by tests
else:
exit_failure(args.config.get("arguments"))

Check warning on line 2192 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2192

Added line #L2192 was not covered by tests


def load_child(target: Target, child_path: Path) -> None:
Expand All @@ -2149,7 +2205,7 @@
return child


def acquire_children_and_targets(target: Target, args: argparse.Namespace) -> None:
def acquire_children_and_targets(target: Target, args: argparse.Namespace) -> list[str | Path]:
if args.child:
target = load_child(target, args.child)

Expand All @@ -2172,7 +2228,7 @@
files.extend(acquire_target(target, args, args.start_time))

except Exception:
log.exception("Failed to acquire target")
log.error("Failed to acquire main target")

Check warning on line 2231 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2231

Added line #L2231 was not covered by tests
acquire_gui.message("Failed to acquire target")
acquire_gui.wait_for_quit()
raise
Expand All @@ -2192,29 +2248,11 @@
child_files = acquire_target(child_target, args)
files.extend(child_files)
except Exception:
log.exception("Failed to acquire child target")
log.exception("Failed to acquire child target %s", child_target.name)

Check warning on line 2251 in acquire/acquire.py

View check run for this annotation

Codecov / codecov/patch

acquire/acquire.py#L2251

Added line #L2251 was not covered by tests
acquire_gui.message("Failed to acquire child target")
continue

files = sort_files(files)

if args.auto_upload:
log_file_handler = get_file_handler(log)
if log_file_handler:
log_file_handler.close()

log.info("")
try:
upload_files(files, args.upload_plugin)
acquire_gui.finish()
acquire_gui.wait_for_quit()
except Exception:
log.exception("Failed to upload files")
acquire_gui.message("Failed to upload files")
acquire_gui.wait_for_quit()
else:
acquire_gui.finish()
acquire_gui.wait_for_quit()
return files


def sort_files(files: list[Union[str, Path]]) -> list[Path]:
Expand Down
4 changes: 3 additions & 1 deletion acquire/uploaders/plugin.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Optional
Expand Down Expand Up @@ -31,7 +33,7 @@ def finish(self, client: Any) -> None:


def upload_files_using_uploader(
uploader: UploaderPlugin, paths: list[Path], proxies: Optional[dict[str, str]] = None
uploader: UploaderPlugin, paths: list[str | Path], proxies: Optional[dict[str, str]] = None
) -> None:
"""Uploads the files in ``paths`` to a destination.

Expand Down
7 changes: 5 additions & 2 deletions acquire/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -278,8 +278,12 @@
ValueError: When an invalid combination of arguments is found.
"""
upload_plugin = None
setattr(args, "upload_plugin", upload_plugin)

# check & set upload related configuration
if args.upload and args.auto_upload:
raise ValueError("only one of --upload or --auto-upload can be specified")

Check warning on line 285 in acquire/utils.py

View check run for this annotation

Codecov / codecov/patch

acquire/utils.py#L285

Added line #L285 was not covered by tests

if args.upload or args.auto_upload:
upload_mode = args.config.get("upload", {}).get("mode")
if not upload_mode:
Expand All @@ -291,8 +295,7 @@

# If initialization of the plugin fails, a ValueError is raised
upload_plugin = upload_plugin_cls(**args.config)

setattr(args, "upload_plugin", upload_plugin)
setattr(args, "upload_plugin", upload_plugin)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this set two times?

Copy link
Contributor Author

@pyrco pyrco Jul 4, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It needs to be set (and set to a default value), as the auto upload code checks on this parameter. The auto upload code is run in a separate try ... except from the check around setting parameters and acquiring, so it can't be sure is has been set (or set to a proper upload plugin) as that may be the thing that failed.


if not args.upload:
# check output related configuration
Expand Down
Loading