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))