Annotation of embedaddon/mtr/test/mtrpacket.py, revision 1.1

1.1     ! misho       1: #
        !             2: #   mtr  --  a network diagnostic tool
        !             3: #   Copyright (C) 2016  Matt Kimball
        !             4: #
        !             5: #   This program is free software; you can redistribute it and/or modify
        !             6: #   it under the terms of the GNU General Public License version 2 as
        !             7: #   published by the Free Software Foundation.
        !             8: #
        !             9: #   This program is distributed in the hope that it will be useful,
        !            10: #   but WITHOUT ANY WARRANTY; without even the implied warranty of
        !            11: #   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
        !            12: #   GNU General Public License for more details.
        !            13: #
        !            14: #   You should have received a copy of the GNU General Public License
        !            15: #   along with this program; if not, write to the Free Software
        !            16: #   Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
        !            17: #
        !            18: 
        !            19: '''Infrastructure for running tests which invoke mtr-packet.'''
        !            20: 
        !            21: import fcntl
        !            22: import os
        !            23: import select
        !            24: import socket
        !            25: import subprocess
        !            26: import sys
        !            27: import time
        !            28: import unittest
        !            29: 
        !            30: #
        !            31: #  typing is used for mypy type checking, but isn't required to run,
        !            32: #  so it's okay if we can't import it.
        !            33: #
        !            34: try:
        !            35:     # pylint: disable=locally-disabled, unused-import
        !            36:     from typing import Dict, List
        !            37: except ImportError:
        !            38:     pass
        !            39: 
        !            40: 
        !            41: IPV6_TEST_HOST = 'google-public-dns-a.google.com'
        !            42: 
        !            43: 
        !            44: class MtrPacketExecuteError(Exception):
        !            45:     "Exception raised when MtrPacketTest can't execute mtr-packet"
        !            46:     pass
        !            47: 
        !            48: 
        !            49: class ReadReplyTimeout(Exception):
        !            50:     'Exception raised by TestProbe.read_reply upon timeout'
        !            51: 
        !            52:     pass
        !            53: 
        !            54: 
        !            55: class WriteCommandTimeout(Exception):
        !            56:     'Exception raised by TestProbe.write_command upon timeout'
        !            57: 
        !            58:     pass
        !            59: 
        !            60: 
        !            61: class MtrPacketReplyParseError(Exception):
        !            62:     "Exception raised when MtrPacketReply can't parse the reply string"
        !            63: 
        !            64:     pass
        !            65: 
        !            66: 
        !            67: class PacketListenError(Exception):
        !            68:     'Exception raised when we have unexpected results from mtr-packet-listen'
        !            69: 
        !            70:     pass
        !            71: 
        !            72: 
        !            73: def set_nonblocking(file_descriptor):  # type: (int) -> None
        !            74:     'Put a file descriptor into non-blocking mode'
        !            75: 
        !            76:     flags = fcntl.fcntl(file_descriptor, fcntl.F_GETFL)
        !            77: 
        !            78:     # pylint: disable=locally-disabled, no-member
        !            79:     fcntl.fcntl(file_descriptor, fcntl.F_SETFL, flags | os.O_NONBLOCK)
        !            80: 
        !            81: 
        !            82: def check_for_local_ipv6():
        !            83:     '''Check for IPv6 support on the test host, to see if we should skip
        !            84:     the IPv6 tests'''
        !            85: 
        !            86:     addrinfo = socket.getaddrinfo(IPV6_TEST_HOST, 1, socket.AF_INET6)
        !            87:     if len(addrinfo):
        !            88:         addr = addrinfo[0][4]
        !            89: 
        !            90:     #  Create a UDP socket and check to see it can be connected to
        !            91:     #  IPV6_TEST_HOST.  (Connecting UDP requires no packets sent, just
        !            92:     #  a route present.)
        !            93:     sock = socket.socket(
        !            94:         socket.AF_INET6, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        !            95: 
        !            96:     connect_success = False
        !            97:     try:
        !            98:         sock.connect(addr)
        !            99:         connect_success = True
        !           100:     except socket.error:
        !           101:         pass
        !           102: 
        !           103:     sock.close()
        !           104: 
        !           105:     if not connect_success:
        !           106:         sys.stderr.write(
        !           107:             'This host has no IPv6.  Skipping IPv6 tests.\n')
        !           108: 
        !           109:     return connect_success
        !           110: 
        !           111: 
        !           112: HAVE_IPV6 = check_for_local_ipv6()
        !           113: 
        !           114: 
        !           115: # pylint: disable=locally-disabled, too-few-public-methods
        !           116: class MtrPacketReply(object):
        !           117:     'A parsed reply from mtr-packet'
        !           118: 
        !           119:     def __init__(self, reply):  # type: (unicode) -> None
        !           120:         self.token = 0  # type: int
        !           121:         self.command_name = None  # type: unicode
        !           122:         self.argument = {}  # type: Dict[unicode, unicode]
        !           123: 
        !           124:         self.parse_reply(reply)
        !           125: 
        !           126:     def parse_reply(self, reply):  # type (unicode) -> None
        !           127:         'Parses a reply string into members for the instance of this class'
        !           128: 
        !           129:         tokens = reply.split()  # type List[unicode]
        !           130: 
        !           131:         try:
        !           132:             self.token = int(tokens[0])
        !           133:             self.command_name = tokens[1]
        !           134:         except IndexError:
        !           135:             raise MtrPacketReplyParseError(reply)
        !           136: 
        !           137:         i = 2
        !           138:         while i < len(tokens):
        !           139:             try:
        !           140:                 name = tokens[i]
        !           141:                 value = tokens[i + 1]
        !           142:             except IndexError:
        !           143:                 raise MtrPacketReplyParseError(reply)
        !           144: 
        !           145:             self.argument[name] = value
        !           146:             i += 2
        !           147: 
        !           148: 
        !           149: class PacketListen(object):
        !           150:     'A test process which listens for a single packet'
        !           151: 
        !           152:     def __init__(self, *args):
        !           153:         self.process_args = list(args)  # type: List[unicode]
        !           154:         self.listen_process = None  # type: subprocess.Popen
        !           155:         self.attrib = None  # type: Dict[unicode, unicode]
        !           156: 
        !           157:     def __enter__(self):
        !           158:         try:
        !           159:             self.listen_process = subprocess.Popen(
        !           160:                 ['./mtr-packet-listen'] + self.process_args,
        !           161:                 stdin=subprocess.PIPE,
        !           162:                 stdout=subprocess.PIPE)
        !           163:         except OSError:
        !           164:             raise PacketListenError('unable to launch mtr-packet-listen')
        !           165: 
        !           166:         status = self.listen_process.stdout.readline().decode('utf-8')
        !           167:         if status != 'status listening\n':
        !           168:             raise PacketListenError('unexpected status')
        !           169: 
        !           170:         return self
        !           171: 
        !           172:     def __exit__(self, exc_type, exc_value, traceback):
        !           173:         self.wait_for_exit()
        !           174: 
        !           175:         self.attrib = {}
        !           176:         for line in self.listen_process.stdout.readlines():
        !           177:             tokens = line.decode('utf-8').split()
        !           178: 
        !           179:             if len(tokens) >= 2:
        !           180:                 name = tokens[0]
        !           181:                 value = tokens[1]
        !           182: 
        !           183:                 self.attrib[name] = value
        !           184: 
        !           185:         self.listen_process.stdin.close()
        !           186:         self.listen_process.stdout.close()
        !           187: 
        !           188:     def wait_for_exit(self):
        !           189:         '''Poll the subprocess for up to ten seconds, until it exits.
        !           190: 
        !           191:         We need to wait for its exit to ensure we are able to read its
        !           192:         output.'''
        !           193: 
        !           194:         wait_time = 10
        !           195:         wait_step = 0.1
        !           196: 
        !           197:         steps = int(wait_time / wait_step)
        !           198: 
        !           199:         exit_value = None
        !           200: 
        !           201:         # pylint: disable=locally-disabled, unused-variable
        !           202:         for i in range(steps):
        !           203:             exit_value = self.listen_process.poll()
        !           204:             if exit_value is not None:
        !           205:                 break
        !           206: 
        !           207:             time.sleep(wait_step)
        !           208: 
        !           209:         if exit_value is None:
        !           210:             raise PacketListenError('mtr-packet-listen timeout')
        !           211: 
        !           212:         if exit_value != 0:
        !           213:             raise PacketListenError('mtr-packet-listen unexpected error')
        !           214: 
        !           215: 
        !           216: class MtrPacketTest(unittest.TestCase):
        !           217:     '''Base class for tests invoking mtr-packet.
        !           218: 
        !           219:     Start a new mtr-packet subprocess for each test, and kill it
        !           220:     at the conclusion of the test.
        !           221: 
        !           222:     Provide methods for writing commands and reading replies.
        !           223:     '''
        !           224: 
        !           225:     def __init__(self, *args):
        !           226:         self.reply_buffer = None  # type: unicode
        !           227:         self.packet_process = None  # type: subprocess.Popen
        !           228:         self.stdout_fd = None  # type: int
        !           229: 
        !           230:         super(MtrPacketTest, self).__init__(*args)
        !           231: 
        !           232:     def setUp(self):
        !           233:         'Set up a test case by spawning a mtr-packet process'
        !           234: 
        !           235:         packet_path = os.environ.get('MTR_PACKET', './mtr-packet')
        !           236: 
        !           237:         self.reply_buffer = ''
        !           238:         try:
        !           239:             self.packet_process = subprocess.Popen(
        !           240:                 [packet_path],
        !           241:                 stdin=subprocess.PIPE,
        !           242:                 stdout=subprocess.PIPE)
        !           243:         except OSError:
        !           244:             raise MtrPacketExecuteError(packet_path)
        !           245: 
        !           246:         #  Put the mtr-packet process's stdout in non-blocking mode
        !           247:         #  so that we can read from it without a timeout when
        !           248:         #  no reply is available.
        !           249:         self.stdout_fd = self.packet_process.stdout.fileno()
        !           250:         set_nonblocking(self.stdout_fd)
        !           251: 
        !           252:         self.stdin_fd = self.packet_process.stdin.fileno()
        !           253:         set_nonblocking(self.stdin_fd)
        !           254: 
        !           255:     def tearDown(self):
        !           256:         'After a test, kill the running mtr-packet instance'
        !           257: 
        !           258:         self.packet_process.stdin.close()
        !           259:         self.packet_process.stdout.close()
        !           260: 
        !           261:         try:
        !           262:             self.packet_process.kill()
        !           263:         except OSError:
        !           264:             return
        !           265: 
        !           266:     def parse_reply(self, timeout=10.0):  # type: (float) -> MtrPacketReply
        !           267:         '''Read the next reply from mtr-packet and parse it into
        !           268:         an MtrPacketReply object.'''
        !           269: 
        !           270:         reply_str = self.read_reply(timeout)
        !           271: 
        !           272:         return MtrPacketReply(reply_str)
        !           273: 
        !           274:     def read_reply(self, timeout=10.0):  # type: (float) -> unicode
        !           275:         '''Read the next reply from mtr-packet.
        !           276: 
        !           277:         Attempt to read the next command reply from mtr-packet.  If no reply
        !           278:         is available withing the timeout time, raise ReadReplyTimeout
        !           279:         instead.'''
        !           280: 
        !           281:         start_time = time.time()
        !           282: 
        !           283:         #  Read from mtr-packet until either the timeout time has elapsed
        !           284:         #  or we read a newline character, which indicates a finished
        !           285:         #  reply.
        !           286:         while True:
        !           287:             now = time.time()
        !           288:             elapsed = now - start_time
        !           289: 
        !           290:             select_time = timeout - elapsed
        !           291:             if select_time < 0:
        !           292:                 select_time = 0
        !           293: 
        !           294:             select.select([self.stdout_fd], [], [], select_time)
        !           295: 
        !           296:             reply_bytes = None
        !           297: 
        !           298:             try:
        !           299:                 reply_bytes = os.read(self.stdout_fd, 1024)
        !           300:             except OSError:
        !           301:                 pass
        !           302: 
        !           303:             if reply_bytes:
        !           304:                 self.reply_buffer += reply_bytes.decode('utf-8')
        !           305: 
        !           306:             #  If we have read a newline character, we can stop waiting
        !           307:             #  for more input.
        !           308:             newline_ix = self.reply_buffer.find('\n')
        !           309:             if newline_ix != -1:
        !           310:                 break
        !           311: 
        !           312:             if elapsed >= timeout:
        !           313:                 raise ReadReplyTimeout()
        !           314: 
        !           315:         reply = self.reply_buffer[:newline_ix]
        !           316:         self.reply_buffer = self.reply_buffer[newline_ix + 1:]
        !           317:         return reply
        !           318: 
        !           319:     def write_command(self, cmd, timeout=10.0):
        !           320:         # type: (unicode, float) -> None
        !           321: 
        !           322:         '''Send a command string to the mtr-packet instance, timing out
        !           323:         if we are unable to write for an extended period of time.  The
        !           324:         timeout is to avoid deadlocks with the child process where both
        !           325:         the parent and the child are writing to their end of the pipe
        !           326:         and expecting the other end to be reading.'''
        !           327: 
        !           328:         command_str = cmd + '\n'
        !           329:         command_bytes = command_str.encode('utf-8')
        !           330: 
        !           331:         start_time = time.time()
        !           332: 
        !           333:         while True:
        !           334:             now = time.time()
        !           335:             elapsed = now - start_time
        !           336: 
        !           337:             select_time = timeout - elapsed
        !           338:             if select_time < 0:
        !           339:                 select_time = 0
        !           340: 
        !           341:             select.select([], [self.stdin_fd], [], select_time)
        !           342: 
        !           343:             bytes_written = 0
        !           344:             try:
        !           345:                 bytes_written = os.write(self.stdin_fd, command_bytes)
        !           346:             except OSError:
        !           347:                 pass
        !           348: 
        !           349:             command_bytes = command_bytes[bytes_written:]
        !           350:             if not len(command_bytes):
        !           351:                 break
        !           352: 
        !           353:             if elapsed >= timeout:
        !           354:                 raise WriteCommandTimeout()
        !           355: 
        !           356: 
        !           357: def check_running_as_root():
        !           358:     'Print a warning to stderr if we are not running as root.'
        !           359: 
        !           360:     # pylint: disable=locally-disabled, no-member
        !           361:     if sys.platform != 'cygwin' and os.getuid() > 0:
        !           362:         sys.stderr.write(
        !           363:             'Warning: many tests require running as root\n')

FreeBSD-CVSweb <freebsd-cvsweb@FreeBSD.org>