#
# mtr -- a network diagnostic tool
# Copyright (C) 2016 Matt Kimball
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#
'''Infrastructure for running tests which invoke mtr-packet.'''
import fcntl
import os
import select
import socket
import subprocess
import sys
import time
import unittest
#
# typing is used for mypy type checking, but isn't required to run,
# so it's okay if we can't import it.
#
try:
# pylint: disable=locally-disabled, unused-import
from typing import Dict, List
except ImportError:
pass
IPV6_TEST_HOST = 'google-public-dns-a.google.com'
class MtrPacketExecuteError(Exception):
"Exception raised when MtrPacketTest can't execute mtr-packet"
pass
class ReadReplyTimeout(Exception):
'Exception raised by TestProbe.read_reply upon timeout'
pass
class WriteCommandTimeout(Exception):
'Exception raised by TestProbe.write_command upon timeout'
pass
class MtrPacketReplyParseError(Exception):
"Exception raised when MtrPacketReply can't parse the reply string"
pass
class PacketListenError(Exception):
'Exception raised when we have unexpected results from mtr-packet-listen'
pass
def set_nonblocking(file_descriptor): # type: (int) -> None
'Put a file descriptor into non-blocking mode'
flags = fcntl.fcntl(file_descriptor, fcntl.F_GETFL)
# pylint: disable=locally-disabled, no-member
fcntl.fcntl(file_descriptor, fcntl.F_SETFL, flags | os.O_NONBLOCK)
def check_for_local_ipv6():
'''Check for IPv6 support on the test host, to see if we should skip
the IPv6 tests'''
addrinfo = socket.getaddrinfo(IPV6_TEST_HOST, 1, socket.AF_INET6)
if len(addrinfo):
addr = addrinfo[0][4]
# Create a UDP socket and check to see it can be connected to
# IPV6_TEST_HOST. (Connecting UDP requires no packets sent, just
# a route present.)
sock = socket.socket(
socket.AF_INET6, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
connect_success = False
try:
sock.connect(addr)
connect_success = True
except socket.error:
pass
sock.close()
if not connect_success:
sys.stderr.write(
'This host has no IPv6. Skipping IPv6 tests.\n')
return connect_success
HAVE_IPV6 = check_for_local_ipv6()
# pylint: disable=locally-disabled, too-few-public-methods
class MtrPacketReply(object):
'A parsed reply from mtr-packet'
def __init__(self, reply): # type: (unicode) -> None
self.token = 0 # type: int
self.command_name = None # type: unicode
self.argument = {} # type: Dict[unicode, unicode]
self.parse_reply(reply)
def parse_reply(self, reply): # type (unicode) -> None
'Parses a reply string into members for the instance of this class'
tokens = reply.split() # type List[unicode]
try:
self.token = int(tokens[0])
self.command_name = tokens[1]
except IndexError:
raise MtrPacketReplyParseError(reply)
i = 2
while i < len(tokens):
try:
name = tokens[i]
value = tokens[i + 1]
except IndexError:
raise MtrPacketReplyParseError(reply)
self.argument[name] = value
i += 2
class PacketListen(object):
'A test process which listens for a single packet'
def __init__(self, *args):
self.process_args = list(args) # type: List[unicode]
self.listen_process = None # type: subprocess.Popen
self.attrib = None # type: Dict[unicode, unicode]
def __enter__(self):
try:
self.listen_process = subprocess.Popen(
['./mtr-packet-listen'] + self.process_args,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE)
except OSError:
raise PacketListenError('unable to launch mtr-packet-listen')
status = self.listen_process.stdout.readline().decode('utf-8')
if status != 'status listening\n':
raise PacketListenError('unexpected status')
return self
def __exit__(self, exc_type, exc_value, traceback):
self.wait_for_exit()
self.attrib = {}
for line in self.listen_process.stdout.readlines():
tokens = line.decode('utf-8').split()
if len(tokens) >= 2:
name = tokens[0]
value = tokens[1]
self.attrib[name] = value
self.listen_process.stdin.close()
self.listen_process.stdout.close()
def wait_for_exit(self):
'''Poll the subprocess for up to ten seconds, until it exits.
We need to wait for its exit to ensure we are able to read its
output.'''
wait_time = 10
wait_step = 0.1
steps = int(wait_time / wait_step)
exit_value = None
# pylint: disable=locally-disabled, unused-variable
for i in range(steps):
exit_value = self.listen_process.poll()
if exit_value is not None:
break
time.sleep(wait_step)
if exit_value is None:
raise PacketListenError('mtr-packet-listen timeout')
if exit_value != 0:
raise PacketListenError('mtr-packet-listen unexpected error')
class MtrPacketTest(unittest.TestCase):
'''Base class for tests invoking mtr-packet.
Start a new mtr-packet subprocess for each test, and kill it
at the conclusion of the test.
Provide methods for writing commands and reading replies.
'''
def __init__(self, *args):
self.reply_buffer = None # type: unicode
self.packet_process = None # type: subprocess.Popen
self.stdout_fd = None # type: int
super(MtrPacketTest, self).__init__(*args)
def setUp(self):
'Set up a test case by spawning a mtr-packet process'
packet_path = os.environ.get('MTR_PACKET', './mtr-packet')
self.reply_buffer = ''
try:
self.packet_process = subprocess.Popen(
[packet_path],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE)
except OSError:
raise MtrPacketExecuteError(packet_path)
# Put the mtr-packet process's stdout in non-blocking mode
# so that we can read from it without a timeout when
# no reply is available.
self.stdout_fd = self.packet_process.stdout.fileno()
set_nonblocking(self.stdout_fd)
self.stdin_fd = self.packet_process.stdin.fileno()
set_nonblocking(self.stdin_fd)
def tearDown(self):
'After a test, kill the running mtr-packet instance'
self.packet_process.stdin.close()
self.packet_process.stdout.close()
try:
self.packet_process.kill()
except OSError:
return
def parse_reply(self, timeout=10.0): # type: (float) -> MtrPacketReply
'''Read the next reply from mtr-packet and parse it into
an MtrPacketReply object.'''
reply_str = self.read_reply(timeout)
return MtrPacketReply(reply_str)
def read_reply(self, timeout=10.0): # type: (float) -> unicode
'''Read the next reply from mtr-packet.
Attempt to read the next command reply from mtr-packet. If no reply
is available withing the timeout time, raise ReadReplyTimeout
instead.'''
start_time = time.time()
# Read from mtr-packet until either the timeout time has elapsed
# or we read a newline character, which indicates a finished
# reply.
while True:
now = time.time()
elapsed = now - start_time
select_time = timeout - elapsed
if select_time < 0:
select_time = 0
select.select([self.stdout_fd], [], [], select_time)
reply_bytes = None
try:
reply_bytes = os.read(self.stdout_fd, 1024)
except OSError:
pass
if reply_bytes:
self.reply_buffer += reply_bytes.decode('utf-8')
# If we have read a newline character, we can stop waiting
# for more input.
newline_ix = self.reply_buffer.find('\n')
if newline_ix != -1:
break
if elapsed >= timeout:
raise ReadReplyTimeout()
reply = self.reply_buffer[:newline_ix]
self.reply_buffer = self.reply_buffer[newline_ix + 1:]
return reply
def write_command(self, cmd, timeout=10.0):
# type: (unicode, float) -> None
'''Send a command string to the mtr-packet instance, timing out
if we are unable to write for an extended period of time. The
timeout is to avoid deadlocks with the child process where both
the parent and the child are writing to their end of the pipe
and expecting the other end to be reading.'''
command_str = cmd + '\n'
command_bytes = command_str.encode('utf-8')
start_time = time.time()
while True:
now = time.time()
elapsed = now - start_time
select_time = timeout - elapsed
if select_time < 0:
select_time = 0
select.select([], [self.stdin_fd], [], select_time)
bytes_written = 0
try:
bytes_written = os.write(self.stdin_fd, command_bytes)
except OSError:
pass
command_bytes = command_bytes[bytes_written:]
if not len(command_bytes):
break
if elapsed >= timeout:
raise WriteCommandTimeout()
def check_running_as_root():
'Print a warning to stderr if we are not running as root.'
# pylint: disable=locally-disabled, no-member
if sys.platform != 'cygwin' and os.getuid() > 0:
sys.stderr.write(
'Warning: many tests require running as root\n')