plainbox-provider-checkbox 0.3-2 / usr / lib / plainbox-providers-1 / checkbox / bin / network

#!/usr/bin/env python3
"""
Copyright (C) 2012 Canonical Ltd.

Authors
  Jeff Marcom <jeff.marcom@canonical.com>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License version 3,
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, see <http://www.gnu.org/licenses/>.
"""

from argparse import (
    ArgumentParser,
    RawTextHelpFormatter
)
import configparser
import fcntl
import ftplib
from ftplib import FTP
import logging
import os
import re
import shlex
import socket
import struct
import subprocess
from subprocess import (
    CalledProcessError,
    check_call,
    check_output
)
import sys
import time

logging.basicConfig(level=logging.DEBUG)


class IPerfPerformanceTest(object):
    """Measures performance of interface using iperf client
    and target. Calculated speed is measured against theorectical
    throughput of selected interface"""

    def __init__(
            self,
            interface,
            target,
            protocol="tcp",
            mbytes="1024M"):

        self.iface = Interface(interface)
        self.target = target
        self.protocol = protocol

        self.mbytes = mbytes

    def run(self):
        cmd = "timeout 180 iperf -c {} -n {}".format(self.target, self.mbytes)

        logging.debug(cmd)
        try:
            iperf_return = check_output(
                shlex.split(cmd), universal_newlines=True)
        except CalledProcessError as iperf_exception:
            if iperf_exception.returncode != 124:
                # timeout command will return 124 if iperf timed out, so any
                # other return value means something did fail
                logging.error("Failed executing iperf: %s",
                              iperf_exception.output)
                return iperf_exception.returncode
            else:
                # this is normal so we "except" this exception and we
                # "pass through" whatever output iperf did manage to produce.
                # When confronted with SIGTERM iperf should stop and output
                # a partial (but usable) result.
                logging.warning("iperf timed out - this should be OK")
                iperf_return = iperf_exception.output

        # 930 Mbits/sec\n'
        print(iperf_return)
        match = re.search(r'\d+\s([GM])bits', iperf_return)
        if match:
            throughput = match.group(0).split()[0]
            units = match.group(1)
            # self.iface.max_speed is always in mb/s, so we need to scale
            # throughput to match
            scaled_throughput = int(throughput)
            if units == 'G':
                scaled_throughput *= 1000
            if units == 'K':
                scaled_throughput /= 1000
            try:
                percent = scaled_throughput / int(self.iface.max_speed) * 100
            except ZeroDivisionError:
                # Catches a condition where the interface functions fine but
                # ethtool fails to properly report max speed. In this case
                # it's up to the reviewer to pass or fail.
                logging.error("Max Speed was not reported properly.  Run "
                              "ethtool and verify that the card is properly "
                              "reporting its capabilities.")
                percent = 0

            print("\nTransfer speed: {} {}b/s".format(throughput, units))
            print("%3.2f%% of " % percent, end="")
            print("theoretical max %sMb/s\n" % int(self.iface.max_speed))

            if percent < 40:
                logging.warn("Poor network performance detected")
                return 30

            logging.debug("Passed benchmark")
        else:
            print("Failed iperf benchmark")
            return 1


class FTPPerformanceTest(object):
    """Provides file transfer rate based information while
    using the FTP protocol and sending a file (DEFAULT=1GB)
    over the local or public network using a specified network
    interface on the host."""

    def __init__(
            self,
            target,
            username,
            password,
            interface,
            binary_size=1,
            file2send="ftp_performance_test"):

        self.target = target
        self.username = username
        self.password = password
        self.iface = Interface(interface)
        self.binary_size = binary_size
        self.file2send = file2send

    def _make_file2send(self):
        """
        Makes binary file to send over FTP.
        Size defaults to 1GB if not supplied.
        """

        logging.debug("Creating %sGB file", self.binary_size)

        file_size = (1024 * 1024 * 1024) * self.binary_size
        with open(self.file2send, "wb") as out:
            out.seek((file_size) - 1)
            out.write('\0'.encode())

    def send_file(self, filename=None):
        """
        Sends file over the network using FTP and returns the
        amount of bytes sent and delay between send and completed.
        """

        if filename is None:
            file = open(self.file2send, 'rb')
            filename = self.file2send

        send_time = time.time()

        try:
            logging.debug("Sending file")
            self.remote.storbinary("STOR " + filename, file, 1024)
        except (ftplib.all_errors) as send_failure:
            logging.error("Failed to send file to %s", self.target)
            logging.error("Reason: %s", send_failure)
            return 0, 0

        file.close()

        time_lapse = time.time() - send_time
        bytes_sent = os.stat(filename).st_size

        return bytes_sent, time_lapse

    def close_connection(self):
        """
        Close connection to remote FTP target
        """
        self.remote.close()

    def connect(self):
        """
        Connects to FTP target and set the current directory as /
        """

        logging.debug("Connecting to %s", self.target)
        try:
            self.remote = FTP(self.target)
            self.remote.set_debuglevel(2)
            self.remote.set_pasv(True)
        except socket.error as connect_exception:
            logging.error("Failed to connect to: %s", self.target)
            return False

        logging.debug("Logging in")
        logging.debug("{USER:%s, PASS:%s}", self.username, self.password)

        try:
            self.remote.login(self.username, self.password)
        except ftplib.error_perm as login_exception:
            logging.error("failed to log into target: %s", self.target)
            return False

        default_out_dir = ""
        self.remote.cwd(default_out_dir)
        return True

    def run(self):

        info = {
            "Interface": self.iface.interface,
            "HWAddress": self.iface.macaddress,
            "Duplex": self.iface.duplex_mode,
            "Speed": self.iface.max_speed,
            "Status": self.iface.status
        }

        logging.debug(info)

        if not os.path.isfile(self.file2send):
            self._make_file2send()

        # Connect to FTP target and send file
        connected = self.connect()

        if connected is False:
            return 3

        filesize, delay = self.send_file()

        # Remove created binary
        try:
            os.remove(self.file2send)
        except (IOError, OSError) as file_delete_error:
            logging.error("Could not remove previous ftp file")
            logging.error(file_delete_error)

        if connected and filesize > 0:

            logging.debug("Bytes sent (%s): %.2f seconds", filesize, delay)

            # Calculate transfer rate and determine pass/fail status
            mbs_speed = float(filesize / 131072) / float(delay)
            percent = (mbs_speed / int(info["Speed"])) * 100
            print("Transfer speed:")
            print("%3.2f%% of" % percent)
            print("theoretical max %smbs" % int(info["Speed"]))

            if percent < 40:
                logging.warn("Poor network performance detected")
                return 30

            logging.debug("Passed benchmark")
        else:
            print("Failed sending file via ftp")
            return 1

class StressPerformanceTest:

    def __init__(self, interface, target):
        self.interface = interface
        self.target = target

    def run(self):
        iperf_cmd = 'timeout 320 iperf -c {} -t 300'.format(self.target)
        print("Running iperf...")
        iperf = subprocess.Popen(shlex.split(iperf_cmd))

        ping_cmd = 'ping -I {} {}'.format(self.interface, self.target)
        ping = subprocess.Popen(shlex.split(ping_cmd), stdout=subprocess.PIPE)
        iperf.communicate()
        
        ping.terminate()
        (out, err) = ping.communicate()

        if iperf.returncode != 0:
            return iperf.returncode

        print("Running ping test...")
        result = 0
        time_re = re.compile('(?<=time=)[0-9]*')
        for line in out.decode().split('\n'):
            time = time_re.search(line)

            if time and int(time.group()) > 2000:
                print(line)
                print("ICMP packet was delayed by > 2000 ms.")
                result = 1
            if 'unreachable' in line.lower():
                print(line)
                result = 1
        
        return result
        
class Interface(socket.socket):
    """
    Simple class that provides network interface information.
    """

    def __init__(self, interface):

        super(Interface, self).__init__(
            socket.AF_INET, socket.IPPROTO_ICMP)

        self.interface = interface

        self.dev_path = os.path.join("/sys/class/net", self.interface)

    def _read_data(self, type):
        try:
            return open(os.path.join(self.dev_path, type)).read().strip()
        except OSError:
            print("{}: Attribute not found".format(type))

    @property
    def ipaddress(self):
        freq = struct.pack('256s', self.interface[:15].encode())

        try:
            nic_data = fcntl.ioctl(self.fileno(), 0x8915, freq)
        except IOError:
            logging.error("No IP address for %s", self.interface)
            return 1
        return socket.inet_ntoa(nic_data[20:24])

    @property
    def netmask(self):
        freq = struct.pack('256s', self.interface.encode())

        try:
            mask_data = fcntl.ioctl(self.fileno(), 0x891b, freq)
        except IOError:
            logging.error("No netmask for %s", self.interface)
            return 1
        return socket.inet_ntoa(mask_data[20:24])

    @property
    def max_speed(self):
        return self._read_data("speed")

    @property
    def macaddress(self):
        return self._read_data("address")

    @property
    def duplex_mode(self):
        return self._read_data("duplex")

    @property
    def status(self):
        return self._read_data("operstate")

    @property
    def device_name(self):
        return self._read_data("device/label")


def interface_test(args):
    if not "test_type" in vars(args):
        return

    # Check and make sure that interface is indeed connected
    try:
        cmd = "ip link set dev %s up" % args.interface
        check_call(shlex.split(cmd))
    except CalledProcessError as interface_failure:
        logging.error("Failed to use %s:%s", cmd, interface_failure)
        return 1

    # Give interface enough time to get DHCP address
    time.sleep(10)

    # Open Network config file
    DEFAULT_CFG = "/etc/checkbox.d/network.cfg"
    if not "config" in vars(args):
        config_file = DEFAULT_CFG
    else:
        config_file = args.config

    config = configparser.SafeConfigParser()
    config.readfp(open(config_file))

    # Set default network config options
    test_target = args.target
    test_user = args.username
    test_pass = args.password

    if test_target is None:
        # Set FTP parameters based on config file
        test_target = config.get("FTP", "Target")
        test_user = config.get("FTP", "User")
        test_pass = config.get("FTP", "Pass")

        if (args.test_type.lower() == "iperf" or
            args.test_type.lower() == "stress"):
            test_target = config.get("IPERF", "Target")

        if "example.com" in test_target:
            # Default values found in config file
            logging.error("Please supply target via: %s", config_file)
            sys.exit(1)


    result = 0
    # Stop all other interfaces
    extra_interfaces = \
        [iface for iface in os.listdir("/sys/class/net")
         if iface != "lo" and iface != args.interface]

    for iface in extra_interfaces:
        logging.debug("Shutting down interface:%s", iface)
        try:
            cmd = "ip link set dev %s down" % iface
            check_call(shlex.split(cmd))
        except CalledProcessError as interface_failure:
            logging.error("Failed to use %s:%s", cmd, interface_failure)
            result = 3

    if result == 0:
        # Execute FTP transfer benchmarking test
        if args.test_type.lower() == "ftp":
            ftp_benchmark = FTPPerformanceTest(
                test_target, test_user, test_pass, args.interface)

            if args.filesize:
                ftp_benchmark.binary_size = int(args.filesize)
            result = ftp_benchmark.run()

        elif args.test_type.lower() == "iperf":
            iperf_benchmark = IPerfPerformanceTest(args.interface, test_target)
            result = iperf_benchmark.run()

        elif args.test_type.lower() == "stress":
            stress_benchmark = StressPerformanceTest(args.interface, test_target)
            result = stress_benchmark.run()
    
    for iface in extra_interfaces:
        logging.debug("Restoring interface:%s", iface)
        try:
            cmd = "ip link set dev %s up" % iface
            check_call(shlex.split(cmd))
        except CalledProcessError as interface_failure:
            logging.error("Failed to use %s:%s", cmd, interface_failure)
            result = 3

    return result

def interface_info(args):

    info_set = ""
    if "all" in vars(args):
        info_set = args.all

    for key, value in vars(args).items():
        if value is True or info_set is True:
            key = key.replace("-", "_")
            try:
                print(
                    key + ":", getattr(Interface(args.interface), key),
                    file=sys.stderr)
            except AttributeError:
                pass


def main():

    intro_message = "Network module\n\nThis script provides benchmarking " \
    + "and information for a specified network interface.\n\n\n" \
    + "Example NIC information usage:\nnetwork info -i eth0 --max-speed " \
    + "\n\nFor running ftp benchmark test: \nnetwork test -i eth0 -t ftp " \
    + "--target 192.168.0.1 --username USERID --password PASSW0RD " \
    + "--filesize-2\n\nPlease note that this script can use configuration " \
    + "values supplied via a config file.\nExample config file:\n[FTP]\n" \
    + "Target: 192.168.1.23\nUser: FTPUser\nPass:PassW0Rd\n" \
    + "[IPERF]\nTarget: 192.168.1.45\n**NOTE**\nDefault config location " \
    + "is /etc/checkbox.d/network.cfg"


    parser = ArgumentParser(
        description=intro_message, formatter_class=RawTextHelpFormatter)
    subparsers = parser.add_subparsers()

    # Main cli options
    test_parser = subparsers.add_parser(
        'test', help=("Run network performance test"))
    info_parser = subparsers.add_parser(
        'info', help=("Gather network info"))

    # Sub test options
    test_parser.add_argument(
        '-i', '--interface', type=str, required=True)
    test_parser.add_argument(
        '-t', '--test_type', type=str, 
        choices=("ftp", "iperf", "stress"), default="ftp",
        help=("[FTP *Default*]"))
    test_parser.add_argument('--target', type=str)
    test_parser.add_argument(
        '--username', type=str, help=("For FTP test only"))
    test_parser.add_argument(
        '--password', type=str, help=("For FTP test only"))
    test_parser.add_argument(
        '--filesize', type=str,
        help="Size (GB) of binary file to send **Note** for FTP test only")
    test_parser.add_argument(
        '--config', type=str,
        default="/etc/checkbox.d/network.cfg",
        help="Supply config file for target/host network parameters")


    # Sub info options
    info_parser.add_argument(
        '-i', '--interface', type=str, required=True)
    info_parser.add_argument(
        '--all', default=False, action="store_true")
    info_parser.add_argument(
        '--duplex-mode', default=False, action="store_true")
    info_parser.add_argument(
        '--max-speed', default=False, action="store_true")
    info_parser.add_argument(
        '--ipaddress', default=False, action="store_true")
    info_parser.add_argument(
        '--netmask', default=False, action="store_true")
    info_parser.add_argument(
        '--device-name', default=False, action="store_true")
    info_parser.add_argument(
        '--macaddress', default=False, action="store_true")
    info_parser.add_argument(
        '--status', default=False, action="store_true",
        help=("displays connection status"))

    test_parser.set_defaults(func=interface_test)
    info_parser.set_defaults(func=interface_info)

    args = parser.parse_args()

    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())