From cfb57843c34a9ba9378c0a0de91adc026ae0142d Mon Sep 17 00:00:00 2001
From: 4shen0ne <4shen.01@gmail.com>
Date: Sun, 4 Aug 2024 15:18:51 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A7=20improve=20output=20in=20asynchro?=
 =?UTF-8?q?nous=20scanning?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

change default value of `with_output` to False
---
 nmass/masscan.py |  5 +++--
 nmass/nmap.py    |  4 ++--
 nmass/scanner.py | 49 +++++++++++++++++++++---------------------------
 3 files changed, 26 insertions(+), 32 deletions(-)

diff --git a/nmass/masscan.py b/nmass/masscan.py
index 474eae9..0d644a4 100644
--- a/nmass/masscan.py
+++ b/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)
diff --git a/nmass/nmap.py b/nmass/nmap.py
index 059094b..30f7d1d 100644
--- a/nmass/nmap.py
+++ b/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)
diff --git a/nmass/scanner.py b/nmass/scanner.py
index 67cc335..93c2325 100644
--- a/nmass/scanner.py
+++ b/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())