Improvements (log2timeline#764)

* Save published messages

* Accept list-based params

* Fix this annoying pytype thing

* workflow_uuid as optional parameter

* Revert "Save published messages"

This reverts commit 704f251.

* Fix linter

dftimewolf/cli/dftimewolf_recipes.py

@@ -96,7 +96,10 @@ class DFTimewolfTool(object):
  _DEFAULT_DATA_FILES_PATH = os.path.join(
  os.sep, 'usr', 'local', 'share', 'dftimewolf')
 
- def __init__(self, cdm: Optional[CursesDisplayManager] = None) -> None:
+ def __init__(
+ self,
+ cdm: Optional[CursesDisplayManager] = None,
+ workflow_uuid: Optional[str] = None) -> None:
  """Initializes a DFTimewolf tool."""
  super(DFTimewolfTool, self).__init__()
  self._command_line_options: Optional[argparse.Namespace]
@@ -107,7 +110,9 @@ def __init__(self, cdm: Optional[CursesDisplayManager] = None) -> None:
  self.dry_run = False
  self.cdm = cdm
  self._state: "dftw_state.DFTimewolfState" # for pytype
- self.uuid = str(uuid.uuid4())
+ if not workflow_uuid:
+ workflow_uuid = str(uuid.uuid4())
+ self.uuid = workflow_uuid
  self.telemetry = telemetry.GetTelemetry(uuid=self.uuid)
 
  self._DetermineDataFilesPath()

dftimewolf/lib/collectors/gce_disk_copy.py

@@ -62,8 +62,8 @@ def SetUp(self,
  destination_project_name: Optional[str],
  source_project_name: str,
  destination_zone: str,
- remote_instance_names: Optional[str],
- disk_names: Optional[str],
+ remote_instance_names: Optional[str | List[str]],
+ disk_names: Optional[str | List[str]],
  all_disks: bool,
  stop_instances: bool) -> None:
  """Sets up a GCEDiskCopyCollector.
@@ -110,9 +110,16 @@ def SetUp(self,
  else:
  self.destination_project = self.source_project
 
- self.remote_instance_names = (
- remote_instance_names.split(',') if remote_instance_names else [])
- self.disk_names = disk_names.split(',') if disk_names else []
+ if isinstance(remote_instance_names, str):
+ self.remote_instance_names = (
+ remote_instance_names.split(',') if remote_instance_names else [])
+ else:
+ self.remote_instance_names = remote_instance_names or []
+
+ if isinstance(disk_names, str):
+ self.disk_names = disk_names.split(',') if disk_names else []
+ else:
+ self.disk_names = disk_names or []
  self.all_disks = all_disks
  self.stop_instances = stop_instances
 

dftimewolf/lib/collectors/gsheets.py

@@ -194,7 +194,7 @@ def _ValidateSpreadSheetId(self, spreadsheet: str) -> str:
 
  def _ExtractEntriesFromSheet(self,
  spreadsheet_id: str,
- sheet_title: str) -> pd.DataFrame:
+ sheet_title: str) -> pd.DataFrame | None:
  """Extract entries from the sheet inside the spreadsheet.
 
  Args: