🚧 improve output in asynchronous scanning

change default value of `with_output` to False

nmass/masscan.py

@@ -24,7 +24,7 @@ def __post_init__(self):
     def run(
         self,
         timeout: float | None = None,
-        with_output: bool = True,
+        with_output: bool = False,
     ) -> NmapRun | None:
         """Run masscan command
 
@@ -44,7 +44,8 @@ def run(
     async def arun(
         self,
         timeout: float | None = None,
-        with_output: bool = True,
+        # FIXME: 异步执行 masscan 时，没有输出进度和倒计时那一行
+        with_output: bool = False,
     ) -> NmapRun | None:
         try:
             return await self._arun_command(timeout, with_output)

nmass/nmap.py

@@ -24,7 +24,7 @@ def __post_init__(self):
     def run(
         self,
         timeout: float | None = None,
-        with_output: bool = True,
+        with_output: bool = False,
     ) -> NmapRun | None:
         """Run nmap command
 
@@ -43,7 +43,7 @@ def run(
     async def arun(
         self,
         timeout: float | None = None,
-        with_output: bool = True,
+        with_output: bool = False,
     ) -> NmapRun | None:
         try:
             return await self._arun_command(timeout, with_output)

nmass/scanner.py

@@ -2,6 +2,7 @@
 import logging
 import subprocess
 import tempfile
+import time
 from abc import abstractmethod
 from dataclasses import dataclass, field
 from typing import Self
@@ -17,18 +18,10 @@ class Scanner:
     _args: list[str] = field(default_factory=lambda: [], init=False)
 
     @abstractmethod
-    def run(
-        self,
-        timeout: float | None = None,
-        with_output: bool = True,
-    ) -> NmapRun | None:
+    def run(self, timeout: float | None, with_output: bool) -> NmapRun | None:
         raise NotImplementedError()
 
-    def _run_command(
-        self,
-        timeout: float | None = None,
-        with_output: bool = True,
-    ) -> NmapRun | None:
+    def _run_command(self, timeout: float | None, with_output: bool) -> NmapRun | None:
         with tempfile.NamedTemporaryFile() as xml_out:
             cmd = [self.bin_path, "-oX", xml_out.name, *self._args]
             try:
@@ -51,39 +44,39 @@ def _run_command(
             return None
 
     @abstractmethod
-    async def arun(
-        self,
-        timeout: float | None = None,
-        with_output: bool = True,
-    ) -> NmapRun | None:
+    async def arun(self, timeout: float | None, with_output: bool) -> NmapRun | None:
         raise NotImplementedError()
 
-    async def _arun_command(
-        self,
-        timeout: float | None = None,
-        with_output: bool = True,
-    ) -> NmapRun | None:
+    async def _arun_command(self, timeout: float | None, with_output: bool) -> NmapRun | None:
         async with atempfile.NamedTemporaryFile() as xml_out:
             proc = await asyncio.create_subprocess_exec(
                 self.bin_path,
                 *["-oX", xml_out.name, *self._args],
                 stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.STDOUT,
             )
+
+            if with_output:
+                start_time = time.time()
+                killed = False
+                async for line in proc.stdout:
+                    print(line.decode().rstrip())
+                    if killed:
+                        continue
+                    if timeout and time.time() - start_time > timeout:
+                        proc.kill()
+                        killed = True
+                if killed:
+                    raise asyncio.TimeoutError()
+
             try:
-                stdout, stderr = await asyncio.wait_for(
-                    proc.communicate(), timeout=timeout
-                )
-                # TODO: 运行中实时打印输出
-                if with_output:
-                    print(stdout.decode())
+                await asyncio.wait_for(proc.wait(), timeout=timeout)
             except asyncio.TimeoutError:
                 raise
             except Exception as why:
                 logging.exception(why)
             else:
                 if proc.returncode != 0:
-                    logging.error(stderr.decode())
                     raise subprocess.CalledProcessError(returncode=proc.returncode)
                 else:
                     return NmapRun.from_xml(await xml_out.read())