Improved Ctrl + c behavior

nxc/netexec.py

@@ -14,11 +14,13 @@
 from nxc.logger import nxc_logger
 from nxc.config import nxc_config, nxc_workspace, config_log
 from nxc.database import create_db_engine
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import ThreadPoolExecutor
 import asyncio
 from nxc.helpers import powershell
+import signal
 import shutil
 import os
+import contextlib
 from os.path import exists
 from os.path import join as path_join
 from sys import exit
@@ -40,34 +42,51 @@
     resource.setrlimit(resource.RLIMIT_NOFILE, file_limit)
 
 
-async def start_run(protocol_obj, args, db, targets):  # noqa: RUF029
-    futures = []
+async def start_run(protocol_obj, args, db, targets):
     nxc_logger.debug("Creating ThreadPoolExecutor")
+    loop = asyncio.get_running_loop()
+
     if args.no_progress or len(targets) == 1:
         with ThreadPoolExecutor(max_workers=args.threads) as executor:
             nxc_logger.debug(f"Creating thread for {protocol_obj}")
-            futures = [executor.submit(protocol_obj, args, db, target) for target in targets]
+            futures = [loop.run_in_executor(executor, protocol_obj, args, db, target) for target in targets]
+            await asyncio.gather(*futures, return_exceptions=True)
     else:
         with Progress(console=nxc_console) as progress, ThreadPoolExecutor(max_workers=args.threads) as executor:
-            current = 0
             total = len(targets)
             tasks = progress.add_task(
                 f"[green]Running nxc against {total} {'target' if total == 1 else 'targets'}",
                 total=total,
             )
             nxc_logger.debug(f"Creating thread for {protocol_obj}")
-            futures = [executor.submit(protocol_obj, args, db, target) for target in targets]
-            for _ in as_completed(futures):
-                current += 1  # noqa: SIM113
+            futures = [loop.run_in_executor(executor, protocol_obj, args, db, target) for target in targets]
+            for current, future in enumerate(asyncio.as_completed(futures), 1):
+                with contextlib.suppress(Exception):
+                    await future
                 progress.update(tasks, completed=current)
-    for future in as_completed(futures):
+
+    for i, future in enumerate(futures):
         try:
             future.result()
         except Exception:
-            nxc_logger.exception(f"Exception for target {targets[futures.index(future)]}: {future.exception()}")
+            nxc_logger.exception(f"Exception for target {targets[i]}: {future.exception()}")
+
+
+def ctrl_c(sig, frame):
+    nxc_logger.debug("Got keyboard interrupt")
+    try:
+        if hasattr(nxc_console, "_live_stack") and nxc_console._live_stack:
+            for live in nxc_console._live_stack:
+                live.stop()
+        nxc_console.show_cursor(True)
+    except Exception:
+        pass
+    nxc_logger.highlight("[!] Interrupted, exiting.")
+    os._exit(0)
 
 
 def main():
+    signal.signal(signal.SIGINT, ctrl_c)
     first_run_setup(nxc_logger)
     args, version_info = gen_cli_args()
 
@@ -209,8 +228,6 @@ def main():
 
     try:
         asyncio.run(start_run(protocol_object, args, db, targets))
-    except KeyboardInterrupt:
-        nxc_logger.debug("Got keyboard interrupt")
     finally:
         db_engine.dispose()