diff --git a/pythonping/__init__.py b/pythonping/__init__.py index 2b3c29e..c5ea7f4 100644 --- a/pythonping/__init__.py +++ b/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 diff --git a/pythonping/aioexecutor.py b/pythonping/aioexecutor.py new file mode 100644 index 0000000..893aa0b --- /dev/null +++ b/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) diff --git a/pythonping/aionetwork.py b/pythonping/aionetwork.py new file mode 100644 index 0000000..f9b37d9 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/pythonping/aioping.py b/pythonping/aioping.py new file mode 100644 index 0000000..f2b0b44 --- /dev/null +++ b/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 diff --git a/pythonping/executor.py b/pythonping/executor.py index 085c6b1..270e658 100644 --- a/pythonping/executor.py +++ b/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() \ No newline at end of file diff --git a/pythonping/network.py b/pythonping/network.py index 080b9b0..55d2707 100644 --- a/pythonping/network.py +++ b/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.")