diff --git a/ileapp.py b/ileapp.py index 8caf43795..4e4442ac2 100755 --- a/ileapp.py +++ b/ileapp.py @@ -10,9 +10,14 @@ import scripts.report as report import traceback import sys +import multiprocessing +import signal +import time as time_module + import scripts.plugin_loader as plugin_loader import leapp_functions.app.history as history +import scripts.lavafuncs as lavafuncs from shutil import copy2 from getpass import getpass @@ -25,7 +30,7 @@ from scripts.context import Context from scripts.ios_keychain import report_supplied_keychain from scripts.lavafuncs import lava_json_name - +from scripts.mp_plugin_runner import run_one_plugin def validate_args(args): if args.artifact_paths or args.create_profile_casedata: @@ -203,13 +208,32 @@ def main(): help=("Path to a keychain file captured from the device. Some apps keep " "their database key in the keychain, which is collected separately " "from the file system extraction.")) + parser.add_argument('--mp_per_plugin', required=False, action="store_true", default=False, + help=("EXPERIMENTAL: Run each plugin in its own subprocess (spawn). " + "This enables skipping a long-running plugin (planned) without stopping the whole run.")) # Check if no arguments were provided if len(sys.argv) == 1: parser.print_help(sys.stderr) sys.exit() - args = parser.parse_args() + # Filter out multiprocessing spawn arguments that argparse doesn't recognize + # When using multiprocessing with 'spawn', Python may inject internal flags that + # our CLI parser doesn't know about (and would otherwise treat as fatal errors). + filtered_argv = [] + for arg in sys.argv[1:]: + # Skip multiprocessing / interpreter-injected arguments + if ( + arg.startswith('--multiprocessing-') + or arg.startswith('tracker_fd=') + or arg.startswith('pipe_handle=') + or arg in ['-B', '-S', '-I'] # Python optimization / startup flags + ): + continue + filtered_argv.append(arg) + + # Use filtered arguments for parsing so frozen mp children don't break argparse + args = parser.parse_args(filtered_argv) available_plugins = [] loader_paths = [plugin_loader.PLUGINPATH] @@ -352,6 +376,9 @@ def main(): out_params = OutputParameters(output_path, custom_output_folder) Context.set_output_params(out_params) + if args.mp_per_plugin: + logfunc("EXPERIMENTAL MODE ENABLED: --mp_per_plugin (per-plugin subprocess execution)") + initialize_lava(input_path, out_params.output_folder_base, extracttype) # Record history if enabled @@ -359,13 +386,14 @@ def main(): history.record_output_path(output_path) crunch_artifacts(selected_plugins, extracttype, input_path, out_params, wrap_text, loader, casedata, time_offset, - profile_filename, itunes_backup_password) + profile_filename, itunes_backup_password, decryption_keys=None, mp_per_plugin=args.mp_per_plugin) lava_finalize_output(out_params.output_folder_base) def crunch_artifacts( plugins: typing.Sequence[plugin_loader.PluginSpec], extracttype, input_path, out_params, wrap_text, - loader: plugin_loader.PluginLoader, casedata, time_offset, profile_filename, itunes_backup_password=None, decryption_keys=None): + loader: plugin_loader.PluginLoader, casedata, time_offset, profile_filename, + itunes_backup_password=None, decryption_keys=None, mp_per_plugin: bool = False): start = process_time() start_wall = perf_counter() @@ -442,6 +470,168 @@ def crunch_artifacts( log.write(f'Extraction/Path selected: {input_path}

') log.write(f'Timezone selected: {time_offset}

') + ctx_mp = multiprocessing.get_context("spawn") if mp_per_plugin else None + installed_os_version = '' + current_proc = None + last_interrupt_ts = 0.0 + # Tracks subprocess state for Ctrl+C handling; kept lightweight (no debug logging). + seeker_all_files = [] + try: + # For iTunes backups, seeker._all_files is a dict keyed by "virtual paths" in backup + # (eg. private/var/mobile/Library/...). + if hasattr(seeker, "_all_files") and isinstance(seeker._all_files, dict): + seeker_all_files = list(seeker._all_files.keys()) + elif hasattr(seeker, "_all_files") and isinstance(seeker._all_files, list): + seeker_all_files = list(seeker._all_files) + except Exception: + seeker_all_files = [] + + def _terminate_current_plugin_proc(reason: str): + nonlocal current_proc + if current_proc is not None and current_proc.is_alive(): + logfunc(f"Skip requested ({reason}). Terminating current plugin subprocess (pid={current_proc.pid}) ...") + try: + current_proc.terminate() + except Exception: + pass + + def _interrupt_handler(signum, frame): + """ + Ctrl+C / Ctrl+Break handling for mp mode: + - first press: terminate current plugin process and continue + - second press within 2 seconds: abort run + """ + nonlocal last_interrupt_ts + # IMPORTANT: use time module explicitly; `time` can be shadowed by datetime.time via star-imports. + now = time_module.time() + if now - last_interrupt_ts < 2.0: + logfunc("Second interrupt received. Aborting run.") + raise KeyboardInterrupt + last_interrupt_ts = now + _terminate_current_plugin_proc("SIGINT/SIGBREAK") + + def _skip_handler(signum, frame): + """ + SIGUSR1/SIGUSR2 handling for mp mode: + - Always skip current plugin (no abort risk) + - Can be sent externally via: kill -USR1 or kill -USR2 + """ + signal_name = "SIGUSR1" if signum == signal.SIGUSR1 else "SIGUSR2" + _terminate_current_plugin_proc(signal_name) + + if mp_per_plugin: + # Register interrupt handler (cross-platform): + # - SIGINT is Ctrl+C everywhere + # - SIGBREAK is Ctrl+Break on Windows (if present) + signal.signal(signal.SIGINT, _interrupt_handler) + if hasattr(signal, "SIGBREAK"): + signal.signal(signal.SIGBREAK, _interrupt_handler) + # Register skip handlers (Unix-like systems only): + # - SIGUSR1: dedicated skip signal (no abort risk) + # - SIGUSR2: dedicated skip signal (no abort risk) + if hasattr(signal, "SIGUSR1"): + signal.signal(signal.SIGUSR1, _skip_handler) + if hasattr(signal, "SIGUSR2"): + signal.signal(signal.SIGUSR2, _skip_handler) + + def _run_plugin_subprocess(plugin_key: str, files_found: list, category_folder: str, file_infos_subset: dict | None = None): + """Run a plugin in a spawned subprocess and merge returned deltas into parent globals.""" + nonlocal installed_os_version + nonlocal current_proc + q = ctx_mp.SimpleQueue() + q_reader = getattr(q, "_reader", None) + payload = { + "plugin_key": plugin_key, + "files_found": files_found, + "category_folder": category_folder, + "wrap_text": wrap_text, + "time_offset": time_offset, + "output_folder_base": out_params.output_folder_base, + "input_path": input_path, + "extracttype": extracttype, + "file_infos_subset": file_infos_subset or {}, + "installed_os_version": installed_os_version, + "seeker_all_files": seeker_all_files, + } + # Parent doesn't hold a write lock while the child writes to the same DB file + # (lavafuncs uses plain sqlite3.connect with no WAL/busy_timeout, so a lingering + # parent connection risks "database is locked"). Reopened right after the child + # finishes, since the loop needs it again for the next plugin's bookkeeping. + lava_close_db() + proc = ctx_mp.Process(target=run_one_plugin, args=(payload, q)) + current_proc = proc + proc.start() + + # Join in a loop so we can react to Ctrl+C and treat it as "skip current plugin". + # Also proactively drain the queue while the child is alive to avoid the child blocking on put(). + result = None + while proc.is_alive(): + # Try to drain once per tick if data is available, but never block. + try: + if result is None and q_reader is not None and hasattr(q_reader, "poll") and q_reader.poll(0): + result = q.get() + except Exception: + pass + proc.join(timeout=0.25) + + # Child is done writing; safe to reopen the parent's connection for the + # rest of this iteration's bookkeeping and the next plugin's. + lava_open_existing(out_params.output_folder_base) + + # If we killed the process due to interrupt/skip, treat it as a skip and keep going. + # IMPORTANT: do this BEFORE attempting any queue reads, because queue state can be corrupted + # if the child was terminated mid-put. + if proc.exitcode is not None and proc.exitcode < 0: + logfunc(f"Plugin {plugin_key} was interrupted (exitcode={proc.exitcode}). Skipping.") + current_proc = None + return {"ok": True, "plugin_key": plugin_key, "skipped": True} + + if result is None: + result = None + try: + # multiprocessing.SimpleQueue.get() has no timeout; use underlying reader.poll() to avoid blocking. + has_data = None + try: + if q_reader is not None and hasattr(q_reader, "poll"): + has_data = bool(q_reader.poll(0)) + except Exception: + has_data = None + if result is None and has_data: + result = q.get() + except Exception: + result = None + if not result or not result.get("ok"): + err = (result or {}).get("error") if result else "Child process failed without result" + tb = (result or {}).get("traceback") if result else "" + raise RuntimeError(f"Subprocess plugin run failed for {plugin_key}: {err}\n{tb}") + + icons_delta = result.get("icons_delta") or {} + for cat, icon_map in icons_delta.items(): + icons.setdefault(cat, {}).update(icon_map) + + lava_meta_delta = result.get("lava_meta_delta") + if lava_meta_delta: + # IMPORTANT: use module global, not the `lava_data` name imported into this module. + lavafuncs.lava_merge_meta_delta(lavafuncs.lava_data, lava_meta_delta) + + for item in (result.get("lava_only_delta") or []): + try: + lava_only_info(item["category"], item["artifact_name"], item["table_name"], item["records"]) + except Exception: + pass + + # Keep installed OS version in parent so we can pass it to future subprocesses + discovered_os_version = result.get("installed_os_version") or '' + if discovered_os_version: + installed_os_version = discovered_os_version + try: + iOS.set_version(discovered_os_version) + Context.set_installed_os_version(discovered_os_version) + except Exception: + pass + + return result + # Special processing for iTunesBackup Info.plist as it is a seperate entity, not part of the Manifest.db. Seeker won't find it if extracttype == 'itunes': info_plist_path = os.path.join(input_path, 'Info.plist') @@ -455,7 +645,10 @@ def crunch_artifacts( except (FileExistsError, FileNotFoundError) as ex: logfunc('Error creating report directory at path {}'.format(report_folder)) logfunc('Error was {}'.format(str(ex))) - loader["itunes_backup_info"].method([info_plist_path], report_folder, seeker, wrap_text, time_offset) + if not mp_per_plugin: + loader["itunes_backup_info"].method([info_plist_path], report_folder, seeker, wrap_text, time_offset) + else: + _run_plugin_subprocess("itunes_backup_info", [info_plist_path], report_folder) report_folder = os.path.join(out_params.output_folder_base, '_HTML', 'Installed Apps') if not os.path.exists(report_folder): try: @@ -463,7 +656,10 @@ def crunch_artifacts( except (FileExistsError, FileNotFoundError) as ex: logfunc('Error creating report directory at path {}'.format(report_folder)) logfunc('Error was {}'.format(str(ex))) - loader["itunes_backup_installed_applications"].method([info_plist_path], report_folder, seeker, wrap_text, time_offset) + if not mp_per_plugin: + loader["itunes_backup_installed_applications"].method([info_plist_path], report_folder, seeker, wrap_text, time_offset) + else: + _run_plugin_subprocess("itunes_backup_installed_applications", [info_plist_path], report_folder) #del search_list['last_build'] # removing last_build as this takes its place print([info_plist_path]) # Future: remove special consideration for itunes? Merge into main search else: @@ -538,11 +734,26 @@ def crunch_artifacts( logfunc('Error was {}'.format(str(ex))) continue # cannot do work try: - plugin.method(files_found, category_folder, seeker, wrap_text, time_offset) + if not mp_per_plugin: + plugin.method(files_found, category_folder, seeker, wrap_text, time_offset) + else: + file_infos_subset = {} + try: + for pth in files_found: + fi = getattr(seeker, "file_infos", {}).get(pth) + if fi: + file_infos_subset[pth] = (fi.source_path, fi.creation_date, fi.modification_date) + except Exception: + file_infos_subset = {} + _run_plugin_subprocess(plugin.name, files_found, category_folder, file_infos_subset) + if plugin.name == 'logarchive': lava_db_path = os.path.join(out_params.output_folder_base, '_lava_artifacts.db') if does_table_exist_in_db(lava_db_path, 'logarchive'): - loader["logarchive_artifacts"].method([lava_db_path], category_folder, seeker, wrap_text, time_offset) + if not mp_per_plugin: + loader["logarchive_artifacts"].method([lava_db_path], category_folder, seeker, wrap_text, time_offset) + else: + _run_plugin_subprocess("logarchive_artifacts", [lava_db_path], category_folder, {}) if does_table_exist_in_db(lava_db_path, 'logarchive_artifacts'): unifed_logs_artifacts = [] unifed_logs_artifacts = [plugin.name for plugin in loader.plugins @@ -550,7 +761,10 @@ def crunch_artifacts( and plugin.name != 'logarchive' and plugin.name != 'logarchive_artifacts'] for unifed_log_artifact in unifed_logs_artifacts: - loader[unifed_log_artifact].method([lava_db_path], category_folder, seeker, wrap_text, time_offset) + if not mp_per_plugin: + loader[unifed_log_artifact].method([lava_db_path], category_folder, seeker, wrap_text, time_offset) + else: + _run_plugin_subprocess(unifed_log_artifact, [lava_db_path], category_folder, {}) except Exception as ex: # pylint: disable=broad-exception-caught logfunc('Reading {} artifact had errors!'.format(plugin.name)) logfunc('Error was {}'.format(str(ex))) @@ -600,4 +814,6 @@ def crunch_artifacts( return True if __name__ == '__main__': + # Support for multiprocessing with PyInstaller frozen binaries + multiprocessing.freeze_support() main() diff --git a/scripts/ilapfuncs.py b/scripts/ilapfuncs.py index 849768480..5606992eb 100644 --- a/scripts/ilapfuncs.py +++ b/scripts/ilapfuncs.py @@ -2,6 +2,7 @@ import codecs # pylint: disable=unused-import # re-exported import contextlib import csv +import dataclasses import hashlib import inspect import io @@ -15,6 +16,7 @@ import shutil import sqlite3 import sys +import typing import xml from datetime import datetime, timezone, timedelta @@ -111,6 +113,125 @@ def __init__(self, output_folder, custom_folder_name=None): os.makedirs(self.media_folder, exist_ok=True) os.makedirs(self.html_media_folder, exist_ok=True) + +@dataclasses.dataclass(frozen=True) +class OutputParametersExisting: + """ + Lightweight, picklable output parameters for subprocess use. + + Unlike OutputParameters.__init__, this does NOT create a new timestamped report folder. + It is meant to point at an already-initialized output folder tree created by the parent. + """ + output_folder_base: str + data_folder: str + media_folder: str + html_media_folder: str + + screen_output_file_path: str + screen_output_file_path_devinfo: str + screen_output_file_path_lava_only: str + + +def output_params_from_existing_output_folder_base( + output_folder_base: str, + *, + ensure_dirs: bool = False) -> OutputParametersExisting: + """ + Build OutputParametersExisting from an existing output folder base. + + If ensure_dirs=True, ensures the expected log/data/media directories exist (exist_ok=True). + """ + data_folder = os.path.join(output_folder_base, 'data') + media_folder = os.path.join(output_folder_base, 'media') + html_media_folder = os.path.join(output_folder_base, '_HTML', 'media') + + screen_output_file_path = os.path.join( + output_folder_base, '_HTML', '_Script_Logs', 'Screen_Output.html') + screen_output_file_path_devinfo = os.path.join( + output_folder_base, '_HTML', '_Script_Logs', 'DeviceInfo.html') + screen_output_file_path_lava_only = os.path.join( + output_folder_base, '_HTML', '_Script_Logs', 'Lava_only_artifacts_log.html') + + if ensure_dirs: + os.makedirs(os.path.join(output_folder_base, '_HTML', '_Script_Logs'), exist_ok=True) + os.makedirs(data_folder, exist_ok=True) + os.makedirs(media_folder, exist_ok=True) + os.makedirs(html_media_folder, exist_ok=True) + + # logfunc() writes to these static paths, so we must set them in subprocesses too. + OutputParameters.screen_output_file_path = screen_output_file_path + OutputParameters.screen_output_file_path_devinfo = screen_output_file_path_devinfo + OutputParameters.screen_output_file_path_lava_only = screen_output_file_path_lava_only + + return OutputParametersExisting( + output_folder_base=output_folder_base, + data_folder=data_folder, + media_folder=media_folder, + html_media_folder=html_media_folder, + screen_output_file_path=screen_output_file_path, + screen_output_file_path_devinfo=screen_output_file_path_devinfo, + screen_output_file_path_lava_only=screen_output_file_path_lava_only, + ) + + +@dataclasses.dataclass(frozen=True) +class FileInfoSnapshot: + """ + Picklable snapshot of scripts.search_files.FileInfo. + """ + source_path: str + creation_date: float + modification_date: float + + +class SeekerProxy: + """ + Minimal seeker stand-in for subprocesses. + + Currently only provides .file_infos, which is used by media helpers to map extracted + paths back to original source paths and timestamps. + """ + def __init__(self, file_infos: dict[str, FileInfoSnapshot], all_files: typing.Optional[list[str]] = None): + self.file_infos = file_infos + # Optional lightweight search index: list of "virtual paths" (eg. iTunes full paths like + # private/var/mobile/Library/...). + self._all_files = all_files or [] + self.searched = {} + + def search(self, filepattern, return_on_first_hit: bool = False, force: bool = False): + """ + Lightweight seeker.search() for subprocess mode. + + This mimics FileSeekerItunes.search() behavior for matching (fnmatch against virtual paths), + but returns *already-extracted* paths only. It does NOT perform extraction/copying. + It exists so plugins that call seeker.search() for auxiliary files don't crash in mp mode. + """ + try: + if filepattern in self.searched and not force: + pathlist = self.searched[filepattern] + return pathlist[0] if return_on_first_hit and pathlist else pathlist + + import fnmatch as _fnmatch # local import to keep module import side-effects minimal + + matches = _fnmatch.filter(self._all_files, filepattern) if self._all_files else [] + + # Map virtual paths to extracted paths via known file_infos (only those copied/extracted so far) + found = [] + if matches: + for extracted_path, finfo in (self.file_infos or {}).items(): + try: + if finfo and finfo.source_path in matches: + found.append(extracted_path) + if return_on_first_hit: + break + except Exception: + continue + + self.searched[filepattern] = found + return found[0] if return_on_first_hit and found else found + except Exception: + return "" if return_on_first_hit else [] + class GuiWindow: '''This only exists to hold window handle if script is run from GUI''' window_handle = None # static variable diff --git a/scripts/lavafuncs.py b/scripts/lavafuncs.py index b229978b7..9db8d0849 100644 --- a/scripts/lavafuncs.py +++ b/scripts/lavafuncs.py @@ -34,6 +34,7 @@ from collections import OrderedDict import re import datetime +import typing from scripts.version_info import leapp_name, leapp_version from scripts.context import Context @@ -44,6 +45,9 @@ lava_db_name = '_lava_artifacts.db' lava_json_name = '_lava_data.lava' +def _get_lava_db_path(output_path: str) -> str: + return os.path.join(output_path, lava_db_name) + def sanitize_sql_name(name): """ @@ -137,7 +141,7 @@ def initialize_lava(input_path, output_path, input_type): } } - db_path = os.path.join(output_path, lava_db_name) + db_path = _get_lava_db_path(output_path) lava_db = sqlite3.connect(db_path) cursor = lava_db.cursor() @@ -155,7 +159,7 @@ def initialize_lava(input_path, output_path, input_type): file_path_id INTEGER NOT NULL, FOREIGN KEY (artifact_search_pattern_id) REFERENCES _artifact_search_patterns(id), FOREIGN KEY (file_path_id) REFERENCES _file_path_list(id))''') - cursor.execute('''CREATE TABLE _lava_media_items ( + cursor.execute('''CREATE TABLE IF NOT EXISTS _lava_media_items ( id TEXT PRIMARY KEY, source_path TEXT, extraction_path TEXT, @@ -164,13 +168,15 @@ def initialize_lava(input_path, output_path, input_type): created_at INTEGER, updated_at INTEGER, is_embedded INTEGER)''') - cursor.execute('''CREATE TABLE _lava_media_references ( + cursor.execute('''CREATE TABLE IF NOT EXISTS _lava_media_references ( id TEXT PRIMARY KEY, media_item_id TEXT, module_name TEXT, artifact_name TEXT, name TEXT, FOREIGN KEY (media_item_id) REFERENCES _lava_media_items(id))''') + # Make view creation idempotent (eg, safe across subprocesses) + cursor.execute('''DROP VIEW IF EXISTS _lava_media_info''') cursor.execute('''CREATE VIEW _lava_media_info AS SELECT lmr.id as 'media_ref_id', @@ -187,6 +193,185 @@ def initialize_lava(input_path, output_path, input_type): lmi.is_embedded FROM _lava_media_references as lmr LEFT JOIN _lava_media_items as lmi ON lmr.media_item_id = lmi.id''') + lava_db.commit() + + +def lava_open_existing(output_path: str) -> None: + """ + Open an existing `_lava_artifacts.db` inside a child process. + + The parent process is expected to have created the database file already. + """ + global lava_db + db_path = _get_lava_db_path(output_path) + lava_db = sqlite3.connect(db_path) + + +def lava_close_db() -> None: + """Close the global LAVA sqlite connection if open.""" + global lava_db + if lava_db is not None: + try: + lava_db.close() + finally: + lava_db = None + + +def _normalize_data_views_for_lava(data_views: typing.Optional[dict], data_headers) -> typing.Optional[dict]: + """ + Normalize & sanitize data views for storing in LAVA metadata. + + - Upgrades legacy 'chat' view to 'conversation' + - Sanitizes values that reference column names so they match table schema sanitization + """ + if not data_views: + return None + + # Deep copy (avoid mutating caller/module globals) + try: + normalized = json.loads(json.dumps(data_views)) + except Exception: + # If non-JSONable, fall back to shallow copy + normalized = dict(data_views) + + # Backward compatibility for chat view. Remove 'chat' once modules are updated. + if "chat" in normalized: + view_params = normalized.pop("chat") + normalized["conversation"] = view_params + + view_params = normalized.get("conversation") + if not view_params: + return normalized + + # Get original column names for dynamic sanitization check + column_names = [item[0] if isinstance(item, tuple) else item for item in data_headers] + + # Conversion map for backward compatibility. Remove once modules are updated. + convert_map = { + "threadDiscriminatorColumn": "conversationDiscriminatorColumn", + "threadLabelColumn": "conversationLabelColumn", + } + + sanitized_params = {} + for key, value in view_params.items(): + final_key = convert_map.get(key, key) + + # Sanitize value if it's a column name, otherwise pass through + if value in column_names: + sanitized_params[final_key] = sanitize_sql_name(value) + else: + sanitized_params[final_key] = value + + normalized["conversation"] = sanitized_params + return normalized + + +def lava_build_artifact_meta_delta( + *, + category: str, + module_name: str, + module_filename: str, + artifact_name: str, + artifact_info: dict, + table_name: str, + column_map: dict, + object_columns: typing.Optional[dict] = None, + record_count: typing.Optional[int] = None, + artifact_icon: typing.Optional[str] = None, + source_path: typing.Optional[str] = None, + data_headers=None, + data_views: typing.Optional[dict] = None, +) -> dict: + """ + Build a pure-data delta for the LAVA metadata structure (picklable). + + Intended to be produced in a subprocess and merged into the parent's in-memory + lava_data, so the parent can write a correct `_lava_data.lava` at the end. + """ + artifact_meta = { + "artifact_key": table_name, + "tablename": table_name, + "name": artifact_name, + "description": artifact_info.get('description', ''), + "author": artifact_info.get('author', ''), + "created_date": artifact_info.get('creation_date', ''), + "last_updated_date": artifact_info.get('last_update_date', ''), + "notes": artifact_info.get('notes', ''), + "category": category + } + + artifact = { + "name": artifact_name, + "tablename": table_name, + "module": module_name, + "column_map": column_map + } + + if artifact_icon: + artifact['artifact_icon'] = artifact_icon + + if record_count is not None: + artifact["record_count"] = record_count + + if source_path: + artifact['source_path'] = source_path + + if object_columns: + artifact["object_columns"] = [{"name": name, "type": type_} for name, type_ in object_columns.items()] + + if data_views: + normalized = _normalize_data_views_for_lava(data_views, data_headers) if data_headers else data_views + if normalized: + artifact['data_views'] = normalized + + return { + "meta_modules": [{ + "module_name": module_name, + "module_filename": module_filename, + "artifacts": [artifact_meta] + }], + "artifacts": { + category: [artifact] + } + } + + +def lava_merge_meta_delta(lava_data_obj: dict, delta: dict) -> None: + """ + Merge a delta created by lava_build_artifact_meta_delta() into lava_data_obj. + """ + if not delta: + return + + lava_data_obj.setdefault("artifacts", OrderedDict()) + lava_data_obj.setdefault("meta", {}).setdefault("modules", []) + + # Merge artifacts by category + for category, artifacts in (delta.get("artifacts") or {}).items(): + if category not in lava_data_obj["artifacts"]: + lava_data_obj["artifacts"][category] = [] + existing = {a.get("tablename") for a in lava_data_obj["artifacts"][category]} + for artifact in artifacts: + if artifact.get("tablename") not in existing: + lava_data_obj["artifacts"][category].append(artifact) + + # Merge meta modules + artifact meta + for mod in delta.get("meta_modules") or []: + module_name = mod.get("module_name") + module_filename = mod.get("module_filename") + module_info = next((m for m in lava_data_obj["meta"]["modules"] if m.get("module_name") == module_name), None) + if not module_info: + module_info = { + "module_name": module_name, + "module_filename": module_filename, + "artifacts": [] + } + lava_data_obj["meta"]["modules"].append(module_info) + + existing_keys = {a.get("artifact_key") for a in (module_info.get("artifacts") or [])} + for artifact_meta in mod.get("artifacts") or []: + if artifact_meta.get("artifact_key") not in existing_keys: + module_info["artifacts"].append(artifact_meta) def lava_process_artifact( @@ -267,40 +452,9 @@ def lava_process_artifact( artifact["object_columns"] = [{"name": name, "type": type_} for name, type_ in object_columns.items()] if data_views: - view_params = None - - # Backward compatibility for chat view. Remove 'chat' once modules are updated. - if "chat" in data_views: - view_params = data_views.pop("chat") - data_views["conversation"] = view_params # Upgrade to conversation - elif "conversation" in data_views: - view_params = data_views.get("conversation") - - if view_params: - sanitized_params = {} - - # Get original column names for dynamic sanitization check - column_names = [item[0] if isinstance(item, tuple) else item for item in data] - - # Conversion map for backward compatibility. Remove once modules are updated. - convert_map = { - "threadDiscriminatorColumn": "conversationDiscriminatorColumn", - "threadLabelColumn": "conversationLabelColumn" - } - - for key, value in view_params.items(): - # Remap old keys to new keys - final_key = convert_map.get(key, key) - - # Sanitize value if it's a column name, otherwise pass through - if value in column_names: - sanitized_params[final_key] = sanitize_sql_name(value) - else: - sanitized_params[final_key] = value - - data_views["conversation"] = sanitized_params - - artifact['data_views'] = data_views + normalized = _normalize_data_views_for_lava(data_views, data) + if normalized: + artifact['data_views'] = normalized lava_data["artifacts"][category].append(artifact) @@ -698,5 +852,6 @@ def lava_finalize_output(output_path): with open(os.path.join(output_path, lava_json_name), 'w', encoding='utf-8') as f: json.dump(lava_data, f, indent=4) - # Close the SQLite database - lava_db.close() + # Close the SQLite database (may be None in per-plugin multiprocessing mode) + if lava_db is not None: + lava_db.close() diff --git a/scripts/mp_plugin_runner.py b/scripts/mp_plugin_runner.py new file mode 100644 index 000000000..5a06e0b65 --- /dev/null +++ b/scripts/mp_plugin_runner.py @@ -0,0 +1,217 @@ +""" +Multiprocessing helpers for iLEAPP. + +This module is designed to be importable in a spawned subprocess. +It provides a per-plugin runner that rehydrates minimal Context state and +returns small, picklable deltas back to the parent process (icons + LAVA metadata). +""" + +from __future__ import annotations + +import os +import traceback +import typing +import warnings +from collections import OrderedDict + +import scripts.plugin_loader as plugin_loader +from scripts.context import Context +from scripts.ilapfuncs import ( + FileInfoSnapshot, + SeekerProxy, + check_output_types, + iOS, + output_params_from_existing_output_folder_base, + logfunc, +) +import scripts.lavafuncs as lavafuncs + + +def _build_file_infos_snapshot(file_infos_subset: dict[str, tuple[str, float, float]]) -> dict[str, FileInfoSnapshot]: + """ + Convert a picklable subset dict from parent into FileInfoSnapshot objects. + + Input format: { extracted_path: (source_path, ctime, mtime) }. + """ + snapshot: dict[str, FileInfoSnapshot] = {} + for extracted_path, info_tuple in (file_infos_subset or {}).items(): + try: + source_path, ctime, mtime = info_tuple + snapshot[extracted_path] = FileInfoSnapshot( + source_path=str(source_path), + creation_date=float(ctime), + modification_date=float(mtime), + ) + except Exception: + # Best-effort: ignore malformed entries + continue + return snapshot + + +def _init_lava_metadata_only(input_path: str, output_path: str, input_type: str) -> None: + """ + Initialize lavafuncs.lava_data without creating/initializing tables. + + The parent process should have already created the DB file & base schema. + """ + lavafuncs.lava_data = { + "param_input": input_path, + "param_output": output_path, + "param_type": input_type, + "processing_status": "In Progress", + "lava_db_name": lavafuncs.lava_db_name, + "modules": [], + "artifacts": OrderedDict(), + "meta": {"modules": []}, + } + + +def run_one_plugin(payload: dict, result_queue) -> None: + """ + Run a single plugin in a subprocess. + + Expected payload keys: + - plugin_key: str + - files_found: list[str] + - category_folder: str + - wrap_text: bool + - time_offset: str + - output_folder_base: str + - input_path: str + - extracttype: str + - file_infos_subset: dict[str, tuple[str, float, float]] + - installed_os_version: str (optional) + - seeker_all_files: list[str] (optional) + """ + # Suppress common deprecation warnings that appear in every subprocess, we need to do this because the plugin loader is using pkg_resources. + warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + warnings.filterwarnings("ignore", message=".*pkg_resources is deprecated.*") + + try: + plugin_key: str = payload["plugin_key"] + files_found: list[str] = payload.get("files_found") or [] + category_folder: str = payload["category_folder"] + wrap_text: bool = bool(payload.get("wrap_text", True)) + time_offset: str = payload.get("time_offset") or "UTC" + output_folder_base: str = payload["output_folder_base"] + input_path: str = payload.get("input_path") or "" + extracttype: str = payload.get("extracttype") or "" + file_infos_subset: dict[str, tuple[str, float, float]] = payload.get("file_infos_subset") or {} + installed_os_version: str = payload.get("installed_os_version") or "" + seeker_all_files: list[str] = payload.get("seeker_all_files") or [] + + # Rehydrate output params + log paths for logfunc() + out_params_existing = output_params_from_existing_output_folder_base( + output_folder_base, ensure_dirs=True + ) + Context.set_output_params(out_params_existing) + + # Propagate installed iOS version into the subprocess (many plugins depend on it) + if installed_os_version: + try: + iOS.set_version(installed_os_version) + Context.set_installed_os_version(installed_os_version) + except Exception: + pass + + # Minimal seeker proxy (media helpers depend on seeker.file_infos) + seeker_proxy = SeekerProxy(_build_file_infos_snapshot(file_infos_subset), seeker_all_files) + + # Load plugin inside child to avoid pickling callables + loader = plugin_loader.PluginLoader() + plugin_spec = loader[plugin_key] + + artifact_info = plugin_spec.artifact_info or {} + output_types = artifact_info.get("output_types", ["html", "tsv", "timeline", "lava", "kml"]) + + # Setup LAVA connection only if needed by this plugin + wants_lava = check_output_types("lava", output_types) or check_output_types("lava_only", output_types) + if wants_lava: + _init_lava_metadata_only(input_path, output_folder_base, extracttype) + lavafuncs.lava_open_existing(output_folder_base) + + # Execute plugin + data_headers, data_list, source_path = plugin_spec.method( + files_found, category_folder, seeker_proxy, wrap_text, time_offset + ) + + # Capture installed OS version if this plugin discovered it + discovered_os_version = "" + try: + discovered_os_version = Context.get_installed_os_version() or "" + except Exception: + discovered_os_version = "" + + # Build deltas to send back + category = artifact_info.get("category", "") + artifact_name = artifact_info.get("name", plugin_key) + artifact_icon = artifact_info.get("artifact_icon", "") + + icons_delta: dict[str, dict[str, str]] = {} + if data_list and category: + icons_delta = {category: {artifact_name: artifact_icon}} + + lava_meta_delta: dict | None = None + lava_only_delta: list[dict] = [] + + if wants_lava and data_headers and data_list: + # Derive table name + maps consistently with lava_create_sqlite_table() + func_name = plugin_spec.name + table_name, column_map, object_columns = lavafuncs.lava_create_sqlite_table(func_name, data_headers) + + module_filename = "" + try: + module_filename = os.path.basename(Context.get_module_file_path()) + except Exception: + module_filename = f"{plugin_spec.module_name}.py" + + lava_meta_delta = lavafuncs.lava_build_artifact_meta_delta( + category=category, + module_name=plugin_spec.module_name, + module_filename=module_filename, + artifact_name=artifact_name, + artifact_info=artifact_info, + table_name=table_name, + column_map=column_map, + object_columns=object_columns, + record_count=len(data_list) if data_list else 0, + artifact_icon=artifact_icon, + source_path=source_path, + data_headers=data_headers, + data_views=artifact_info.get("data_views"), + ) + + # If lava_only, record for the parent's lava-only log + if check_output_types("lava_only", output_types): + lava_only_delta.append({ + "category": category, + "artifact_name": artifact_name, + "table_name": table_name, + "records": len(data_list) if data_list else 0, + }) + + result_queue.put({ + "ok": True, + "plugin_key": plugin_key, + "record_count": len(data_list) if data_list else 0, + "icons_delta": icons_delta, + "lava_meta_delta": lava_meta_delta, + "lava_only_delta": lava_only_delta, + "installed_os_version": discovered_os_version, + }) + + except Exception as ex: + result_queue.put({ + "ok": False, + "plugin_key": payload.get("plugin_key"), + "error": str(ex), + "traceback": traceback.format_exc(), + }) + finally: + # Ensure DB connection is closed in the child + try: + lavafuncs.lava_close_db() + except Exception: + pass + +