Added function to ping a host

This commit is contained in:
2024-04-07 14:24:43 +02:00
parent 439a7d545e
commit 87a3ee2886
1517 changed files with 278486 additions and 1 deletions

View File

@@ -0,0 +1,88 @@
import sys
from random import randint
from . import network, executor, payload_provider
from .utils import random_text
# this needs to be available across all thread usages and will hold ints
SEED_IDs = []
def ping(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'):
"""Pings a remote host and handles the 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 = network.Socket.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
comm = executor.Communicator(target, provider, timeout, interval, socket_options=options, verbose=verbose, output=out,
seed_id=seed_id, source=source, repr_format=out_format)
comm.run(match_payloads=match)
SEED_IDs.remove(seed_id)
return comm.responses

View File

@@ -0,0 +1,388 @@
"""Module that actually performs the ping, sending and receiving packets"""
import os
import sys
import time
from . import icmp
from . import network
# Python 3.5 compatibility
if sys.version_info[1] == 5:
from enum import IntEnum, Enum
class AutoNumber(Enum):
def __new__(cls):
value = len(cls.__members__) + 1
obj = object.__new__(cls)
obj._value = value
return obj
class SuccessOn(AutoNumber):
One = ()
Most = ()
All = ()
else:
from enum import IntEnum, auto
class SuccessOn(IntEnum):
One = auto()
Most = auto()
All = auto()
class Message:
"""Represents an ICMP message with destination socket"""
def __init__(self, target, packet, source):
"""Creates a message that may be sent, or used to represent a response
:param target: Target IP or hostname of the message
:type target: str
:param packet: ICMP packet composing the message
:type packet: icmp.ICMP
:param source: Source IP or hostname of the message
:type source: str"""
self.target = target
self.packet = packet
self.source = source
def send(self, source_socket):
"""Places the message on a socket
:param source_socket: The socket to place the message on
:type source_socket: network.Socket"""
source_socket.send(self.packet.packet)
def __repr__(self):
return repr(self.packet)
def represent_seconds_in_ms(seconds):
"""Converts seconds into human-readable milliseconds with 2 digits decimal precision
:param seconds: Seconds to convert
:type seconds: Union[int, float]
:return: The same time expressed in milliseconds, with 2 digits of decimal precision
:rtype: float"""
return round(seconds * 1000, 2)
class Response:
"""Represents a response to an ICMP message, with metadata like timing"""
def __init__(self, message, time_elapsed, source_request=None, repr_format=None):
"""Creates a representation of ICMP message received in response
:param message: The message received
:type message: Union[None, Message]
:param time_elapsed: Time elapsed since the original request was sent, in seconds
:type time_elapsed: float
:param source_request: ICMP packet represeting the request that originated this response
:type source_request: ICMP
:param repr_format: How to __repr__ the response. Allowed: legacy, None
:type repr_format: str"""
self.message = message
self.time_elapsed = time_elapsed
self.source_request = source_request
self.repr_format = repr_format
@property
def success(self):
return self.error_message is None
@property
def error_message(self):
if self.message is None:
return 'No response'
if self.message.packet.message_type == 0 and self.message.packet.message_code == 0:
# Echo Reply, response OK - no error
return None
if self.message.packet.message_type == 3:
# Destination unreachable, returning more details based on message code
unreachable_messages = [
'Network Unreachable',
'Host Unreachable',
'Protocol Unreachable',
'Port Unreachable',
'Fragmentation Required',
'Source Route Failed',
'Network Unknown',
'Host Unknown',
'Source Host Isolated',
'Communication with Destination Network is Administratively Prohibited',
'Communication with Destination Host is Administratively Prohibited',
'Network Unreachable for ToS',
'Host Unreachable for ToS',
'Communication Administratively Prohibited',
'Host Precedence Violation',
'Precedence Cutoff in Effect'
]
try:
return unreachable_messages[self.message.packet.message_code]
except IndexError:
# Should never generate IndexError, this serves as additional protection
return 'Unreachable'
# Error was not identified
return 'Network Error'
@property
def time_elapsed_ms(self):
return represent_seconds_in_ms(self.time_elapsed)
def legacy_repr(self):
if self.message is None:
return 'Request timed out'
if self.success:
return 'Reply from {0}, {1} bytes in {2}ms'.format(self.message.source,
len(self.message.packet.raw),
self.time_elapsed_ms)
# Not successful, but with some code (e.g. destination unreachable)
return '{0} from {1} in {2}ms'.format(self.error_message, self.message.source, self.time_elapsed_ms)
def __repr__(self):
if self.repr_format == 'legacy':
return self.legacy_repr()
if self.message is None:
return 'Timed out'
if self.success:
return 'status=OK\tfrom={0}\tms={1}\t\tbytes\tsnt={2}\trcv={3}'.format(
self.message.source,
self.time_elapsed_ms,
len(self.source_request.raw)+20,
len(self.message.packet.raw)
)
return 'status=ERR\tfrom={1}\terror="{0}"'.format(self.message.source, self.error_message)
class ResponseList:
"""Represents a series of ICMP responses"""
def __init__(self, initial_set=[], verbose=False, output=sys.stdout):
"""Creates a ResponseList with initial data if available
:param initial_set: Already existing responses
:type initial_set: list
:param verbose: Flag to enable verbose mode, defaults to False
:type verbose: bool
:param output: File where to write verbose output, defaults to stdout
:type output: file"""
self._responses = []
self.clear()
self.verbose = verbose
self.output = output
self.rtt_avg = 0
self.rtt_min = 0
self.rtt_max = 0
self.stats_packets_sent = 0
self.stats_packets_returned = 0
for response in initial_set:
self.append(response)
def success(self, option=SuccessOn.One):
"""Check success state of the request.
:param option: Sets a threshold for success sign. ( 1 - SuccessOn.One, 2 - SuccessOn.Most, 3 - SuccessOn.All )
:type option: int
:return: Whether this set of responses is successful
:rtype: bool
"""
result = False
success_list = [resp.success for resp in self._responses]
if option == SuccessOn.One:
result = True in success_list
elif option == SuccessOn.Most:
result = success_list.count(True) / len(success_list) > 0.5
elif option == SuccessOn.All:
result = False not in success_list
return result
@property
def packet_loss(self):
return self.packets_lost
@property
def rtt_min_ms(self):
return represent_seconds_in_ms(self.rtt_min)
@property
def rtt_max_ms(self):
return represent_seconds_in_ms(self.rtt_max)
@property
def rtt_avg_ms(self):
return represent_seconds_in_ms(self.rtt_avg)
def clear(self):
self._responses = []
self.stats_packets_sent = 0
self.stats_packets_returned = 0
def append(self, value):
self._responses.append(value)
self.stats_packets_sent += 1
if len(self) == 1:
self.rtt_avg = value.time_elapsed
self.rtt_max = value.time_elapsed
self.rtt_min = value.time_elapsed
else:
# Calculate the total of time, add the new value and divide for the new number
self.rtt_avg = ((self.rtt_avg * (len(self)-1)) + value.time_elapsed) / len(self)
if value.time_elapsed > self.rtt_max:
self.rtt_max = value.time_elapsed
if value.time_elapsed < self.rtt_min:
self.rtt_min = value.time_elapsed
if value.success:
self.stats_packets_returned += 1
if self.verbose:
print(value, file=self.output)
@property
def stats_packets_lost(self):
return self.stats_packets_sent - self.stats_packets_returned
@property
def stats_success_ratio(self):
return self.stats_packets_returned / self.stats_packets_sent
@property
def stats_lost_ratio(self):
return 1 - self.stats_success_ratio
@property
def packets_lost(self):
return self.stats_lost_ratio
def __len__(self):
return len(self._responses)
def __repr__(self):
ret = ''
for response in self._responses:
ret += '{0}\r\n'.format(response)
ret += '\r\n'
ret += 'Round Trip Times min/avg/max is {0}/{1}/{2} ms'.format(self.rtt_min_ms, self.rtt_avg_ms, self.rtt_max_ms)
return ret
def __iter__(self):
for response in self._responses:
yield response
class Communicator:
"""Instance actually communicating over the network, sending messages and handling responses"""
def __init__(self, target, payload_provider, timeout, interval, socket_options=(), seed_id=None,
verbose=False, output=sys.stdout, source=None, repr_format=None):
"""Creates an instance that can handle communication with the target device
:param target: IP or hostname of the remote device
:type target: str
:param payload_provider: An iterable list of payloads to send
:type payload_provider: PayloadProvider
:param timeout: Timeout that will apply to all ping messages, in seconds
:type timeout: Union[int, float]
:param interval: Interval to wait between pings, in seconds
:type interval: int
:param socket_options: Options to specify for the network.Socket
:type socket_options: tuple
:param seed_id: The first ICMP packet ID to use
:type seed_id: Union[None, int]
:param verbose: Flag to enable verbose mode, defaults to False
:type verbose: bool
:param output: File where to write verbose output, defaults to stdout
: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.provider = payload_provider
self.timeout = timeout
self.interval = interval
self.responses = ResponseList(verbose=verbose, output=output)
self.seed_id = seed_id
self.repr_format = repr_format
# 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
def __del__(self):
pass
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)
self.socket.send(i.packet)
return i
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 = 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 Response(Message('', response, source_socket[0]), timeout - time_left, source_request, repr_format=self.repr_format)
return Response(None, timeout, source_request, repr_format=self.repr_format)
@staticmethod
def increase_seq(sequence_number):
"""Increases an ICMP sequence number and reset if it gets bigger than 2 bytes
:param sequence_number: The sequence number to increase
:type sequence_number: int
:return: The increased sequence number of 1, in case an increase was not possible
:rtype: int"""
sequence_number += 1
if sequence_number > 0xFFFF:
sequence_number = 1
return sequence_number
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 = self.send_ping(identifier, seq, payload)
if not match_payloads:
self.responses.append(self.listen_for(identifier, self.timeout, None, icmp_out))
else:
self.responses.append(self.listen_for(identifier, self.timeout, icmp_out.payload, icmp_out))
seq = self.increase_seq(seq)
if self.interval:
time.sleep(self.interval)

View File

@@ -0,0 +1,221 @@
import os
import struct
def checksum(data):
"""Creates the ICMP checksum as in RFC 1071
:param data: Data to calculate the checksum ofs
:type data: bytes
:return: Calculated checksum
:rtype: int
Divides the data in 16-bits chunks, then make their 1's complement sum"""
subtotal = 0
for i in range(0, len(data)-1, 2):
subtotal += ((data[i] << 8) + data[i+1]) # Sum 16 bits chunks together
if len(data) % 2: # If length is odd
subtotal += (data[len(data)-1] << 8) # Sum the last byte plus one empty byte of padding
while subtotal >> 16: # Add carry on the right until fits in 16 bits
subtotal = (subtotal & 0xFFFF) + (subtotal >> 16)
check = ~subtotal # Performs the one complement
return ((check << 8) & 0xFF00) | ((check >> 8) & 0x00FF) # Swap bytes
class ICMPType:
"""Represents an ICMP type, as combination of type and code
ICMP Types should inherit from this class so that the code can identify them easily.
This is a static class, not meant to be instantiated"""
def __init__(self):
raise TypeError('ICMPType may not be instantiated')
class Types(ICMPType):
class EchoReply(ICMPType):
type_id = 0
ECHO_REPLY = (type_id, 0,)
class DestinationUnreachable(ICMPType):
type_id = 3
NETWORK_UNREACHABLE = (type_id, 0,)
HOST_UNREACHABLE = (type_id, 1,)
PROTOCOL_UNREACHABLE = (type_id, 2,)
PORT_UNREACHABLE = (type_id, 3,)
FRAGMENTATION_REQUIRED = (type_id, 4,)
SOURCE_ROUTE_FAILED = (type_id, 5,)
NETWORK_UNKNOWN = (type_id, 6,)
HOST_UNKNOWN = (type_id, 7,)
SOURCE_HOST_ISOLATED = (type_id, 8,)
NETWORK_ADMINISTRATIVELY_PROHIBITED = (type_id, 9,)
HOST_ADMINISTRATIVELY_PROHIBITED = (type_id, 10,)
NETWORK_UNREACHABLE_TOS = (type_id, 11,)
HOST_UNREACHABLE_TOS = (type_id, 12,)
COMMUNICATION_ADMINISTRATIVELY_PROHIBITED = (type_id, 13,)
HOST_PRECEDENCE_VIOLATION = (type_id, 14,)
PRECEDENCE_CUTOFF = (type_id, 15,)
class SourceQuench(ICMPType):
type_id = 4
SOURCE_QUENCH = (type_id, 0,)
class Redirect(ICMPType):
type_id = 5
FOR_NETWORK = (type_id, 0,)
FOR_HOST = (type_id, 1,)
FOR_TOS_AND_NETWORK = (type_id, 2,)
FOR_TOS_AND_HOST = (type_id, 3,)
class EchoRequest(ICMPType):
type_id = 8
ECHO_REQUEST = (type_id, 0,)
class RouterAdvertisement(ICMPType):
type_id = 9
ROUTER_ADVERTISEMENT = (type_id, 0,)
class RouterSolicitation(ICMPType):
type_id = 10
ROUTER_SOLICITATION = (type_id, 0)
# Aliases
ROUTER_DISCOVERY = ROUTER_SOLICITATION
ROUTER_SELECTION = ROUTER_SOLICITATION
class TimeExceeded(ICMPType):
type_id = 11
TTL_EXPIRED_IN_TRANSIT = (type_id, 0)
FRAGMENT_REASSEMBLY_TIME_EXCEEDED = (type_id, 1)
class BadIPHeader(ICMPType):
type_id = 12
POINTER_INDICATES_ERROR = (type_id, 0)
MISSING_REQUIRED_OPTION = (type_id, 1)
BAD_LENGTH = (type_id, 2)
class Timestamp(ICMPType):
type_id = 13
TIMESTAMP = (type_id, 0)
class TimestampReply(ICMPType):
type_id = 14
TIMESTAMP_REPLY = (type_id, 0)
class InformationRequest(ICMPType):
type_id = 15
INFORMATION_REQUEST = (type_id, 0)
class InformationReply(ICMPType):
type_id = 16
INFORMATION_REPLY = (type_id, 0)
class AddressMaskRequest(ICMPType):
type_id = 17
ADDRESS_MASK_REQUEST = (type_id, 0)
class AddressMaskReply(ICMPType):
type_id = 18
ADDRESS_MASK_REPLY = (type_id, 0)
class Traceroute(ICMPType):
type_id = 30
INFORMATION_REQUEST = (type_id, 30)
class ICMP:
LEN_TO_PAYLOAD = 41 # Ethernet, IP and ICMP header lengths combined
def __init__(self, message_type=Types.EchoReply, payload=None, identifier=None, sequence_number=1):
"""Creates an ICMP packet
:param message_type: Type of ICMP message to send
:type message_type: Union[ICMPType, (int, int), int]
:param payload: utf8 string or bytes payload
:type payload: Union[str, bytes]
:param identifier: ID of this ICMP packet
:type identifier: int"""
self.message_code = 0
if issubclass(message_type, ICMPType):
self.message_type = message_type.type_id
elif isinstance(message_type, tuple):
self.message_type = message_type[0]
self.message_code = message_type[1]
elif isinstance(message_type, int):
self.message_type = message_type
if payload is None:
payload = bytes('1', 'utf8')
elif isinstance(payload, str):
payload = bytes(payload, 'utf8')
self.payload = payload
if identifier is None:
identifier = os.getpid()
self.id = identifier & 0xFFFF # Prevent identifiers bigger than 16 bits
self.sequence_number = sequence_number
self.received_checksum = None
self.raw = None
@property
def packet(self):
"""The raw packet with header, ready to be sent from a socket"""
p = self._header(check=self.expected_checksum) + self.payload
if self.raw is None:
self.raw = p
return p
def _header(self, check=0):
"""The raw ICMP header
:param check: Checksum value
:type check: int
:return: The packed header
:rtype: bytes"""
# TODO implement sequence number
return struct.pack("BBHHH",
self.message_type,
self.message_code,
check,
self.id,
self.sequence_number)
def __repr__(self):
return ' '.join('{:02x}'.format(b) for b in self.raw)
@property
def is_valid(self):
"""True if the received checksum is valid, otherwise False"""
if self.received_checksum is None:
return True
return self.expected_checksum == self.received_checksum
@property
def expected_checksum(self):
"""The checksum expected for this packet, calculated with checksum field set to 0"""
return checksum(self._header() + self.payload)
@property
def header_length(self):
"""Length of the ICMP header"""
return len(self._header())
@staticmethod
def generate_from_raw(raw):
"""Creates a new ICMP representation from the raw bytes
:param raw: The raw packet including payload
:type raw: bytes
:return: An ICMP instance representing the packet
:rtype: ICMP"""
packet = ICMP()
packet.unpack(raw)
return packet
def unpack(self, raw):
"""Unpacks a raw packet and stores it in this object
:param raw: The raw packet, including payload
:type raw: bytes"""
self.raw = raw
self.message_type, \
self.message_code, \
self.received_checksum, \
self.id, \
self.sequence_number = struct.unpack("BBHHH", raw[20:28])
self.payload = raw[28:]

View File

@@ -0,0 +1,84 @@
import socket
import select
import time
class Socket:
DONT_FRAGMENT = (socket.SOL_IP, 10, 1) # Option value for raw socket
PROTO_LOOKUP = {"icmp": socket.IPPROTO_ICMP, "tcp": socket.IPPROTO_TCP, "udp": socket.IPPROTO_UDP,
"ip": socket.IPPROTO_IP, "raw": socket.IPPROTO_RAW}
def __init__(self, destination, protocol, options=(), buffer_size=2048, source=None):
"""Creates a 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"""
try:
self.destination = socket.gethostbyname(destination)
except socket.gaierror as e:
raise RuntimeError('Cannot resolve address "' + destination + '", try verify your DNS or host file')
self.protocol = Socket.getprotobyname(protocol)
self.buffer_size = buffer_size
self.socket = socket.socket(socket.AF_INET, socket.SOCK_RAW, self.protocol)
self.source = source
if options:
self.socket.setsockopt(*options)
# Implementing a version of socket.getprotobyname for this library since built-in is not thread safe
# for python 3.5, 3.6, and 3.7:
# https://bugs.python.org/issue30482
# This bug was causing failures as it would occasionally return a 0 (incorrect) instead of a 1 (correct)
# for the 'icmp' string, causing a OSError for "Protocol not supported" in multi-threaded usage:
# https://github.com/alessandromaggio/pythonping/issues/40
@staticmethod
def getprotobyname(name):
try:
return Socket.PROTO_LOOKUP[name.lower()]
except KeyError:
raise KeyError("'" + str(name) + "' is not in the list of supported proto types: "
+ str(list(Socket.PROTO_LOOKUP.keys())))
def send(self, packet):
"""Sends a raw packet on the stream
:param packet: The raw packet to send
:type packet: bytes"""
if self.source:
self.socket.bind((self.source, 0))
self.socket.sendto(packet, (self.destination, 0))
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)"""
time_left = timeout
while time_left > 0:
start_select = time.perf_counter()
data_ready = select.select([self.socket], [], [], time_left)
elapsed_in_select = time.perf_counter() - start_select
time_left -= elapsed_in_select
if not data_ready[0]:
# Timeout
return b'', '', time_left
packet, source = self.socket.recvfrom(self.buffer_size)
return packet, source, time_left
def __del__(self):
try:
if hasattr(self, "socket") and self.socket:
self.socket.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.")

View File

@@ -0,0 +1,90 @@
"""Module generating ICMP payloads (with no header)"""
class PayloadProvider:
def __init__(self):
raise NotImplementedError('Cannot create instances of PayloadProvider')
def __iter__(self):
raise NotImplementedError()
def __next__(self):
raise NotImplementedError()
class List(PayloadProvider):
def __init__(self, payload_list):
"""Creates a provider of payloads from an existing list of payloads
:param payload_list: An existing list of payloads
:type payload_list: list"""
self._payloads = payload_list
self._counter = 0
def __iter__(self):
self._counter = 0
return self
def __next__(self):
if self._counter < len(self._payloads):
ret = self._payloads[self._counter]
self._counter += 1
return ret
raise StopIteration
class Repeat(PayloadProvider):
def __init__(self, pattern, count):
"""Creates a provider of many identical payloads
:param pattern: The existing payload
:type pattern: Union[str, bytes]
:param count: How many payloads to generate
:type count: int"""
self.pattern = pattern
self.count = count
self._counter = 0
def __iter__(self):
self._counter = 0
return self
def __next__(self):
if self._counter < self.count:
self._counter += 1
return self.pattern
raise StopIteration
class Sweep(PayloadProvider):
def __init__(self, pattern, start_size, end_size):
"""Creates a provider of payloads of increasing size
:param pattern: The existing payload, may be cut or replicated to fit the size
:type pattern: Union[str, bytes]
:param start_size: The first payload size to start with, included
:type start_size: int
:param end_size: The payload size to end with, included
:type end_size: int"""
if start_size > end_size:
raise ValueError('end_size must be greater or equal than start_size')
if len(pattern) == 0:
raise ValueError('pattern cannot be empty')
self.pattern = pattern
self.start_size = start_size
self.end_size = end_size
# Extend the length of the pattern if needed
while not len(self.pattern) >= end_size:
self.pattern += pattern
self._current_size = self.start_size
def __iter__(self):
self._current_size = self.start_size
return self
def __next__(self):
if self._current_size <= self.end_size:
ret = self.pattern[0:self._current_size]
self._current_size += 1
return ret
raise StopIteration

View File

@@ -0,0 +1,14 @@
"""Module containing service classes and functions"""
import string
import random
def random_text(size):
"""Returns a random text of the specified size
:param size: Size of the random string, must be greater than 0
:type size int
:return: Random string
:rtype: str"""
return ''.join(random.SystemRandom().choice(string.ascii_uppercase + string.digits) for _ in range(size))