Merge pull request #23 from mbrg/split_dump_command

Split dump command

src/powerpwn/cli/arguments.py

@@ -13,16 +13,22 @@ def module_gui(sub_parser: argparse.ArgumentParser):
 
 def module_dump(sub_parser: argparse.ArgumentParser):
  dump_parser = sub_parser.add_parser(
- "dump",
- description="Recon for available data connections and dump their content",
- help="Recon for available data connections and dump their content.",
+ "dump", description="Dump content for all available connection from recon", help="Dump content for all available connection from recon"
  )
  dump_parser.add_argument("-c", "--clear-cache", action="store_true", help="Clear local disk cache")
  dump_parser.add_argument("--cache-path", default=CACHE_PATH, help="Path to store collected resources and data.")
  dump_parser.add_argument("-t", "--tenant", required=False, type=str, help="Tenant id to connect.")
  dump_parser.add_argument("-g", "--gui", action="store_true", help="Run local server for gui.")
 
 
+def module_recon(sub_parser: argparse.ArgumentParser):
+ dump_parser = sub_parser.add_parser("recon", description="Recon for available data connections", help="Recon for available data connections.")
+ dump_parser.add_argument("-c", "--clear-cache", action="store_true", help="Clear local disk cache")
+ dump_parser.add_argument("--cache-path", default=CACHE_PATH, help="Path to store collected resources and data.")
+ dump_parser.add_argument("-t", "--tenant", required=False, type=str, help="Tenant id to connect.")
+ dump_parser.add_argument("-g", "--gui", action="store_true", help="Run local server for gui.")
+
+
 def module_nocodemalware(command_subparsers: argparse.ArgumentParser):
  nocodemalware_parser = command_subparsers.add_parser(
  "nocodemalware",
@@ -126,6 +132,7 @@ def parse_arguments():
  command_subparsers = parser.add_subparsers(help="command", dest="command")
 
  module_dump(command_subparsers)
+ module_recon(command_subparsers)
  module_gui(command_subparsers)
  module_backdoor(command_subparsers)
  module_nocodemalware(command_subparsers)

src/powerpwn/cli/runners.py

@@ -28,17 +28,7 @@ def __init_command_token(args, scope: str) -> str:
  return acquire_token(scope=scope, tenant=args.tenant)
 
 
-def run_dump_command(args):
- _run_collect_resources_command(args)
- _run_collect_data_command(args)
- logger.info(f"Dump is completed in {args.cache_path}")
-
- if args.gui:
- logger.info("Going to run local server for gui")
- run_gui_command(args)
-
-
-def _run_collect_resources_command(args):
+def run_recon_command(args):
  # cache
  if args.clear_cache:
  try:
@@ -52,14 +42,29 @@ def _run_collect_resources_command(args):
  entities_fetcher = ResourcesCollector(token=token, cache_path=args.cache_path)
  entities_fetcher.collect_and_cache()
 
+ logger.info(f"Recon is completed in {args.cache_path}/resources")
+
+ if args.gui:
+ logger.info("Going to run local server for gui")
+ run_gui_command(args)
+
 
 def run_gui_command(args):
  Gui().run(cache_path=args.cache_path)
 
 
-def _run_collect_data_command(args):
+def run_dump_command(args):
  token = __init_command_token(args, API_HUB_SCOPE)
- DataCollector(token=token, cache_path=args.cache_path).collect()
+ is_data_collected = DataCollector(token=token, cache_path=args.cache_path).collect()
+ if not is_data_collected:
+ logger.info("No data dumped. Please run recon first.")
+ return None
+
+ logger.info(f"Dump is completed in {args.cache_path}/data")
+
+ if args.gui:
+ logger.info("Going to run local server for gui")
+ run_gui_command(args)
 
 
 def run_backdoor_flow_command(args):

src/powerpwn/main.py

@@ -4,7 +4,14 @@
 
 from powerpwn.cli.arguments import parse_arguments
 from powerpwn.cli.const import LOGGER_NAME
-from powerpwn.cli.runners import run_backdoor_flow_command, run_dump_command, run_gui_command, run_nocodemalware_command, run_phishing_command
+from powerpwn.cli.runners import (
+ run_backdoor_flow_command,
+ run_dump_command,
+ run_gui_command,
+ run_nocodemalware_command,
+ run_phishing_command,
+ run_recon_command,
+)
 
 logger = logging.getLogger(LOGGER_NAME)
 
@@ -22,6 +29,8 @@ def main():
 
  if command == "dump":
  run_dump_command(args)
+ elif command == "recon":
+ run_recon_command(args)
  elif command == "gui":
  run_gui_command(args)
  elif command == "backdoor":

src/powerpwn/powerdump/collect/data_collectors/data_collector.py

@@ -17,7 +17,11 @@ def __init__(self, cache_path: str, token: str) -> None:
  self.__session = init_session(token=token)
  self.__data_collectors = [ConnectionsDataCollector]
 
- def collect(self) -> None:
+ def collect(self) -> bool:
+ environment_ids = get_environment_ids(self.__cache_path)
+ if len(environment_ids) == 0:
+ return False
+
  for env_id in get_environment_ids(self.__cache_path):
  env_dumps_root_dir = env_collected_data_path(env_id, self.__cache_path)
  if os.path.isdir(env_dumps_root_dir):
@@ -26,3 +30,4 @@ def collect(self) -> None:
  for data_collector in self.__data_collectors:
  data_collector_instance = data_collector(self.__cache_path)
  data_collector_instance.collect(self.__session, env_id, env_dumps_root_dir)
+ return True

src/powerpwn/powerdump/utils/model_loaders.py

@@ -102,6 +102,8 @@ def load_connectors(cache_path: str, env_id: Optional[str] = None) -> Generator[
 
 
 def get_environment_ids(cache_path: str) -> List[str]:
+ if not os.path.exists(cache_path):
+ return []
  return os.listdir(entities_path(cache_path))