✨ asynchronous execution

nmass/errors.py

@@ -11,7 +11,7 @@ def __init__(self, message: str, nmap_arg: str) -> None:
         super().__init__(message, nmap_arg)
 
 
-class NmapExecutinoError(Exception):
+class NmapExecutionError(Exception):
     def __init__(self, message: str, retcode: int) -> None:
         super().__init__(message, retcode)
 
@@ -29,6 +29,6 @@ def __init__(self, message: str, masscan_arg: str) -> None:
         super().__init__(message, masscan_arg)
 
 
-class MasscanExecutinoError(Exception):
+class MasscanExecutionError(Exception):
     def __init__(self, message: str, retcode: int) -> None:
         super().__init__(message, retcode)

nmass/masscan.py

@@ -1,11 +1,11 @@
+import asyncio
 import logging
 import shutil
 import subprocess
-import tempfile
 from dataclasses import dataclass
 from typing import Self
 
-from nmass.errors import MasscanNotInstalledError
+from nmass.errors import MasscanExecutionError, MasscanNotInstalledError
 from nmass.models import NmapRun
 from nmass.scanner import Scanner
 from nmass.utils import as_root
@@ -32,26 +32,27 @@ def run(
         :param with_output: print masscan's output, defaults to True
         :return: NmapRun object or None
         """
-        with tempfile.NamedTemporaryFile(delete_on_close=True) as xml_out:
-            cmd = [self.bin_path, "-oX", xml_out.name, *self._args]
-            try:
-                subprocess.run(
-                    cmd,
-                    check=True,
-                    timeout=timeout,
-                    capture_output=not with_output,
-                )
-            except subprocess.TimeoutExpired:
-                logging.warn("masscan scanning timeout")
-            except subprocess.CalledProcessError as e:
-                logging.error(f"masscan's return code is {e.returncode}")
-                logging.error(e.stderr.decode())
-            except Exception as why:
-                logging.exception(why)
-            else:
-                return NmapRun.from_xml(xml_out.read())
+        try:
+            return self._run_command(timeout, with_output)
+        except subprocess.CalledProcessError as e:
+            raise MasscanExecutionError(retcode=e.returncode)
+        except subprocess.TimeoutExpired:
+            logging.warn("masscan scanning timeout")
+            raise
 
-            return None
+    @as_root
+    async def arun(
+        self,
+        timeout: float | None = None,
+        with_output: bool = True,
+    ) -> NmapRun | None:
+        try:
+            return await self._arun_command(timeout, with_output)
+        except subprocess.CalledProcessError as e:
+            raise MasscanExecutionError(retcode=e.returncode)
+        except asyncio.TimeoutError:
+            logging.warn("asynchronous masscan scanning timeout")
+            raise
 
     def with_rate(self, rate: int) -> Self:
         self._args.extend(("--rate", str(rate)))

nmass/nmap.py

@@ -1,12 +1,12 @@
+import asyncio
 import enum
 import logging
 import shutil
 import subprocess
-import tempfile
 from dataclasses import dataclass
 from typing import Literal, Self
 
-from nmass.errors import NmapArgumentError, NmapNotInstalledError
+from nmass.errors import NmapArgumentError, NmapExecutionError, NmapNotInstalledError
 from nmass.models import NmapRun
 from nmass.scanner import Scanner
 from nmass.utils import as_root
@@ -32,26 +32,26 @@ def run(
         :param with_output: print nmap's output, defaults to True
         :return: NmapRun object or None
         """
-        with tempfile.NamedTemporaryFile(delete_on_close=True) as xml_out:
-            cmd = [self.bin_path, "-oX", xml_out.name, *self._args]
-            try:
-                subprocess.run(
-                    cmd,
-                    check=True,
-                    timeout=timeout,
-                    capture_output=not with_output,
-                )
-            except subprocess.TimeoutExpired:
-                logging.warn("nmap scanning timeout")
-            except subprocess.CalledProcessError as e:
-                logging.error(f"nmap's return code is {e.returncode}")
-                logging.error(e.stderr.decode())
-            except Exception as why:
-                logging.exception(why)
-            else:
-                return NmapRun.from_xml(xml_out.read())
-
-            return None
+        try:
+            return self._run_command(timeout, with_output)
+        except subprocess.CalledProcessError as e:
+            raise NmapExecutionError(retcode=e.returncode)
+        except subprocess.TimeoutExpired:
+            logging.warn("nmap scanning timeout")
+            raise
+
+    async def arun(
+        self,
+        timeout: float | None = None,
+        with_output: bool = True,
+    ) -> NmapRun | None:
+        try:
+            return await self._arun_command(timeout, with_output)
+        except subprocess.CalledProcessError as e:
+            raise NmapExecutionError(retcode=e.returncode)
+        except asyncio.TimeoutError:
+            logging.warn("asynchronous nmap scanning timeout")
+            raise
 
     ### TARGET SPECIFICATION ###
 

nmass/scanner.py

@@ -1,8 +1,13 @@
+import asyncio
+import logging
+import subprocess
+import tempfile
 from abc import abstractmethod
 from dataclasses import dataclass, field
-import logging
 from typing import Self
 
+from aiofiles import tempfile as atempfile
+
 from nmass.models import Address, NmapRun
 
 
@@ -19,6 +24,72 @@ def run(
     ) -> NmapRun | None:
         raise NotImplementedError()
 
+    def _run_command(
+        self,
+        timeout: float | None = None,
+        with_output: bool = True,
+    ) -> NmapRun | None:
+        with tempfile.NamedTemporaryFile() as xml_out:
+            cmd = [self.bin_path, "-oX", xml_out.name, *self._args]
+            try:
+                subprocess.run(
+                    cmd,
+                    check=True,
+                    timeout=timeout,
+                    capture_output=not with_output,
+                )
+            except subprocess.TimeoutExpired:
+                raise
+            except subprocess.CalledProcessError as e:
+                logging.error(e.stderr.decode())
+                raise
+            except Exception as why:
+                logging.exception(why)
+            else:
+                return NmapRun.from_xml(xml_out.read())
+
+            return None
+
+    @abstractmethod
+    async def arun(
+        self,
+        timeout: float | None = None,
+        with_output: bool = True,
+    ) -> NmapRun | None:
+        raise NotImplementedError()
+
+    async def _arun_command(
+        self,
+        timeout: float | None = None,
+        with_output: bool = True,
+    ) -> 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,
+            )
+            try:
+                stdout, stderr = await asyncio.wait_for(
+                    proc.communicate(), timeout=timeout
+                )
+                # TODO: 运行中实时打印输出
+                if with_output:
+                    print(stdout.decode())
+            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())
+
+            return None
+
     def with_step(self, model: NmapRun) -> Self:
         # masscan 中同一个目标会有多个 host element
         targets = set()

pyproject.toml

@@ -8,6 +8,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.12"
 pydantic-xml = "^2.11.0"
+aiofiles = "^24.1.0"
 
 [build-system]
 requires = ["poetry-core"]