mirror of
https://github.com/traviscross/mtr.git
synced 2024-09-21 18:37:10 +00:00
364 lines
10 KiB
Python
364 lines
10 KiB
Python
#
|
|
# 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.,
|
|
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 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')
|