Attempt to adapt asynchronous I/O

pythonping/__init__.py

@@ -2,7 +2,7 @@
 from random import randint
 from . import network, executor, payload_provider
 from .utils import random_text
-
+from .aioping import aping
 
 # this needs to be available across all thread usages and will hold ints
 SEED_IDs = []
@@ -78,11 +78,11 @@ def ping(target,
             break
 
 
-    comm = executor.Communicator(target, provider, timeout, interval, socket_options=options, verbose=verbose, output=out,
-                                 seed_id=seed_id, source=source, repr_format=out_format)
+    with executor.Communicator(target, provider, timeout, interval, socket_options=options, verbose=verbose, output=out,
+                               seed_id=seed_id, source=source, repr_format=out_format) as comm:
 
-    comm.run(match_payloads=match)
+        comm.run(match_payloads=match)
 
-    SEED_IDs.remove(seed_id)
+        SEED_IDs.remove(seed_id)
 
-    return comm.responses
+        return comm.responses

pythonping/aioexecutor.py

@@ -0,0 +1,88 @@
+import asyncio
+
+from . import executor
+from . import aionetwork
+from . import icmp
+
+class AsyncCommunicator(executor.Communicator):
+    socket: aionetwork.AsyncSocket
+
+    @staticmethod
+    def _create_socket(
+        target,
+        protocol='icmp',
+        options=(),
+        buffer_size=2048,
+        source=None,
+    ):
+        return aionetwork.AsyncSocket(target, protocol, options, buffer_size, source)
+
+    async def send_ping(self, packet_id, sequence_number, payload):
+        """Sends one ICMP Echo Request on the socket
+
+        :param packet_id: The ID to use for the packet
+        :type packet_id: int
+        :param sequence_number: The sequence number to use for the packet
+        :type sequence_number: int
+        :param payload: The payload of the ICMP message
+        :type payload: Union[str, bytes]
+        :rtype: ICMP"""
+        i = icmp.ICMP(
+            icmp.Types.EchoRequest,
+            payload=payload,
+            identifier=packet_id, sequence_number=sequence_number)
+        await self.socket.send(i.packet)
+        return i
+
+    async def listen_for(self, packet_id, timeout, payload_pattern=None, source_request=None):
+        """Listens for a packet of a given id for a given timeout
+
+        :param packet_id: The ID of the packet to listen for, the same for request and response
+        :type packet_id: int
+        :param timeout: How long to listen for the specified packet, in seconds
+        :type timeout: float
+        :param payload_pattern: Payload reply pattern to match to request, if set to None, match by ID only
+        :type payload_pattern: Union[None, bytes]
+        :return: The response to the request with the specified packet_id
+        :rtype: Response"""
+        time_left = timeout
+        response = icmp.ICMP()
+        while time_left > 0:
+            # Keep listening until a packet arrives
+            raw_packet, source_socket, time_left = await self.socket.receive(time_left)
+            # If we actually received something
+            if raw_packet != b'':
+                response.unpack(raw_packet)
+
+                # Ensure we have not unpacked the packet we sent (RHEL will also listen to outgoing packets)
+                if response.id == packet_id and response.message_type != icmp.Types.EchoRequest.type_id:
+                    if payload_pattern is None:
+                        # To allow Windows-like behaviour (no payload inspection, but only match packet identifiers),
+                        # simply allow for it to be an always true in the legacy usage case
+                        payload_matched = True
+                    else:
+                        payload_matched = (payload_pattern == response.payload)
+
+                    if payload_matched:
+                        return executor.Response(executor.Message('', response, source_socket[0]), timeout - time_left, source_request, repr_format=self.repr_format)
+        return executor.Response(None, timeout, source_request, repr_format=self.repr_format)
+
+    async def run(self, match_payloads=False):
+        """Performs all the pings and stores the responses
+
+        :param match_payloads: optional to set to True to make sure requests and replies have equivalent payloads
+        :type match_payloads: bool"""
+        self.responses.clear()
+        identifier = self.seed_id
+        seq = 1
+        for payload in self.provider:
+            icmp_out = await self.send_ping(identifier, seq, payload)
+            if not match_payloads:
+                self.responses.append(await self.listen_for(identifier, self.timeout, None, icmp_out))
+            else:
+                self.responses.append(await self.listen_for(identifier, self.timeout, icmp_out.payload, icmp_out))
+
+            seq = self.increase_seq(seq)
+
+            if self.interval:
+                await asyncio.sleep(self.interval)

pythonping/aionetwork.py

@@ -0,0 +1,57 @@
+import time
+import select
+import asyncio
+from . import network
+
+class AsyncSocket(network.Socket):
+    def __init__(self, destination, protocol, options=(), buffer_size=2048, source=None):
+        """Creates a async network socket to exchange messages
+
+        :param destination: Destination IP address
+        :type destination: str
+        :param protocol: Name of the protocol to use
+        :type protocol: str
+        :param options: Options to set on the socket
+        :type options: tuple
+        :param source: Source IP to use - implemented in future releases
+        :type source: Union[None, str]
+        :param buffer_size: Size in bytes of the listening buffer for incoming packets (replies)
+        :type buffer_size: int"""
+        if options is None:
+            options = ()
+        super().__init__(destination, protocol, options, buffer_size, source)
+        # Nonblocking is required here to support asynchronous operations.
+        self.socket.setblocking(False)
+
+    async def send(self, packet):
+        """Sends a raw packet on the stream
+
+        :param packet: The raw packet to send
+        :type packet: bytes"""
+        loop = asyncio.get_running_loop()
+        if self.source:
+            self.socket.bind((self.source, 0))
+        await loop.sock_sendto(self.socket, packet, (self.destination, 0))
+
+    async def receive(self, timeout=2):
+        """Listen for incoming packets until timeout
+
+        :param timeout: Time after which stop listening
+        :type timeout: Union[int, float]
+        :return: The packet, the remote socket, and the time left before timeout
+        :rtype: (bytes, tuple, float)"""
+        loop = asyncio.get_running_loop()
+        start_time = time.perf_counter()
+        try:
+            response = await asyncio.wait_for(
+                loop.sock_recvfrom(
+                    self.socket,
+                    self.buffer_size
+                ),
+                timeout = timeout
+            )
+            packet, source = response
+        except asyncio.TimeoutError:
+            packet, source = b"", b""
+        end_time = time.perf_counter()
+        return packet, source, timeout - (end_time - start_time)

pythonping/aioping.py

@@ -0,0 +1,88 @@
+import sys
+from random import randint
+from . import aionetwork, aioexecutor, payload_provider
+from .utils import random_text
+
+
+# this needs to be available across all thread usages and will hold ints
+SEED_IDs = []
+
+
+async def aping(target,
+                timeout=2,
+                count=4,
+                size=1,
+                interval=0,
+                payload=None,
+                sweep_start=None,
+                sweep_end=None,
+                df=False,
+                verbose=False,
+                out=sys.stdout,
+                match=False,
+                source=None,
+                out_format='legacy'):
+    """Async ping to remote host handling responses
+
+    :param target: The remote hostname or IP address to ping
+    :type target: str
+    :param timeout: Time in seconds before considering each non-arrived reply permanently lost.
+    :type timeout: Union[int, float]
+    :param count: How many times to attempt the ping
+    :type count: int
+    :param size: Size of the entire packet to send
+    :type size: int
+    :param interval: Interval to wait between pings
+    :type interval: int
+    :param payload: Payload content, leave None if size is set to use random text
+    :type payload: Union[str, bytes]
+    :param sweep_start: If size is not set, initial size in a sweep of sizes
+    :type sweep_start: int
+    :param sweep_end: If size is not set, final size in a sweep of sizes
+    :type sweep_end: int
+    :param df: Don't Fragment flag value for IP Header
+    :type df: bool
+    :param verbose: Print output while performing operations
+    :type verbose: bool
+    :param out: Stream to which redirect the verbose output
+    :type out: stream
+    :param match: Do payload matching between request and reply (default behaviour follows that of Windows which is
+    by packet identifier only, Linux behaviour counts a non equivalent payload in reply as fail, such as when pinging
+    8.8.8.8 with 1000 bytes and reply is truncated to only the first 74 of request payload with packet identifiers
+    the same in request and reply)
+    :type match: bool
+    :param repr_format: How to __repr__ the response. Allowed: legacy, None
+    :type repr_format: str
+    :return: List with the result of each ping
+    :rtype: executor.ResponseList"""
+    provider = payload_provider.Repeat(b'', 0)
+    if sweep_start and sweep_end and sweep_end >= sweep_start:
+        if not payload:
+            payload = random_text(sweep_start)
+        provider = payload_provider.Sweep(payload, sweep_start, sweep_end)
+    elif size and size > 0:
+        if not payload:
+            payload = random_text(size)
+        provider = payload_provider.Repeat(payload, count)
+    options = ()
+    if df:
+        options = aionetwork.AsyncSocket.DONT_FRAGMENT
+
+    # Fix to allow for pythonping multithreaded usage;
+    # no need to protect this loop as no one will ever surpass 0xFFFF amount of threads
+    while True:
+        # seed_id needs to be less than or equal to 65535 (as original code was seed_id = getpid() & 0xFFFF)
+        seed_id = randint(0x1, 0xFFFF)
+        if seed_id not in SEED_IDs:
+            SEED_IDs.append(seed_id)
+            break
+
+
+    with aioexecutor.AsyncCommunicator(target, provider, timeout, interval, socket_options=options, verbose=verbose, output=out,
+                                       seed_id=seed_id, source=source, repr_format=out_format) as comm:
+
+        await comm.run(match_payloads=match)
+
+        SEED_IDs.remove(seed_id)
+
+        return comm.responses

pythonping/executor.py

@@ -290,7 +290,7 @@ def __init__(self, target, payload_provider, timeout, interval, socket_options=(
         :type output: file
         :param repr_format: How to __repr__ the response. Allowed: legacy, None
         :type repr_format: str"""
-        self.socket = network.Socket(target, 'icmp', options=socket_options, source=source)
+        self.socket = self._create_socket(target, 'icmp', options=socket_options, source=source)
         self.provider = payload_provider
         self.timeout = timeout
         self.interval = interval
@@ -300,6 +300,16 @@ def __init__(self, target, payload_provider, timeout, interval, socket_options=(
         # note that to make Communicator instances thread safe, the seed ID must be unique per thread
         if self.seed_id is None:
             self.seed_id = os.getpid() & 0xFFFF
+
+    @staticmethod
+    def _create_socket(
+        target,
+        protocol='icmp',
+        options=(),
+        buffer_size=2048,
+        source=None,
+    ):
+        return network.Socket(target, protocol, options, buffer_size, source)
 
     def __del__(self):
         pass
@@ -386,3 +396,13 @@ def run(self, match_payloads=False):
 
             if self.interval:
                 time.sleep(self.interval)
+
+    def close(self):
+        """Closes the socket"""
+        self.socket.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

pythonping/network.py

@@ -74,11 +74,21 @@ def receive(self, timeout=2):
                 return b'', '', time_left
             packet, source = self.socket.recvfrom(self.buffer_size)
             return packet, source, time_left
+
+    def close(self):
+        """Close the socket"""
+        self.socket.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
 
     def __del__(self):
         try:
             if hasattr(self, "socket") and self.socket:
-                self.socket.close()
+                self.close()
         except AttributeError:
             raise AttributeError("Attribute error because of failed socket init. Make sure you have the root privilege."
                                  " This error may also be caused from DNS resolution problems.")