__init__.py 5.45 KB
Newer Older
1 2
from __future__ import print_function

3 4 5 6 7 8 9 10 11 12 13 14 15
import socket
import logging
import time
import logging.handlers
import subprocess
import re
import sys
import shutil
import netifaces
import random
import pycurl
import argparse
import json
16 17 18 19
from io import StringIO
from .ping import ping, ping6
from .dnsbench import resolve
from .http import get_curl, request
20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
import textwrap

class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter):

    def _get_help_string(self, action):
        return super(HelpFormatter, self)._get_help_string(action) \
            if action.default else action.help

    def _split_lines(self, text, width):
        """Preserves new lines in option descriptions"""
        lines = []
        for text in text.splitlines():
            lines += textwrap.wrap(text, width)
        return lines

    def _fill_text(self, text, width, indent):
        """Preserves new lines in other descriptions"""
        kw = dict(width=width, initial_indent=indent, subsequent_indent=indent)
        return '\n'.join(textwrap.fill(t, **kw) for t in text.splitlines())

botname = socket.gethostname()

date_reg_exp = re.compile('\d{4}[-/]\d{2}[-/]\d{2}')

def _get_network_gateway(self):
  return netifaces.gateways()["default"][netifaces.AF_INET][0]

def load_configuration(config_path):
  if config_path.startswith("http://") or \
      config_path.startswith("ftp://") or \
      config_path.startswith("https://") or \
      config_path.startswith("file://"):
    return download_external_configuration(config_path)

  with open(config_path, "r") as f:
    return json.load(f)

def download_external_configuration(url):
  buffer = StringIO()
  curl, _ = get_curl(buffer, url)

  response_code = curl.getinfo(pycurl.HTTP_CODE)
  
  if response_code in (200, 0):
    try:
      return json.loads(buffer.getvalue())
    except ValueError:
67
      print("Unable to parse external configuration, error:")
68 69 70
      import traceback
      traceback.print_exc(file=sys.stderr)
      sys.stderr.flush()
71
      print("Ignoring external configuration")
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182
    finally:
      curl.close()

  return {}

def is_rotate_log(log_file_path):
  try:
    log_file = open(log_file_path, 'r')
  except IOError:
    return False

  # Handling try-except-finally together.
  try:
    try:
      log_file.seek(0, 2)
      size = log_file.tell()
      log_file.seek(-min(size, 4096), 1)
      today = time.strftime("%Y-%m-%d")
      
      for line in reversed(log_file.read().split('\n')):
        if len(line.strip()):
          match_list = date_reg_exp.findall(line)
          if len(match_list):
            if match_list[0] != today:
              return ValueError(match_list[0])
            break

    except IOError:
      return False
  finally:
    log_file.close()

def rotate_logfile(handler, log_file):
  last_date = is_rotate_log(log_file)
  if last_date: 
    handler.doRollover()
    today = time.strftime("%Y-%m-%d")
    shutil.move("%s.%s" % (log_file, today),
                "%s.%s" % (log_file, last_date))
    # The simpler the better
    sp = subprocess.Popen("gzip %s.%s" % (log_file, last_date),
               stdout=subprocess.PIPE, stdin=subprocess.PIPE, shell=True)
    sp.communicate()

def create_logger(name, log_folder, verbose):
  new_logger = logging.getLogger(name)

  new_logger.setLevel(logging.DEBUG)
  log_file = '%s/network_bench.%s.log' % (log_folder, name)
  handler = logging.handlers.TimedRotatingFileHandler(
                     log_file, when="D",
                     backupCount=1000)

  rotate_logfile(handler, log_file)

  format = "%%(asctime)-16s;%s;%%(message)s" % botname
  handler.setFormatter(logging.Formatter(format))
  new_logger.addHandler(handler)

  if verbose:
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(logging.Formatter(format))
    new_logger.addHandler(handler)
  return new_logger

def run_all(config_dict, log_folder, verbose):
  dns_logger = create_logger("dns", log_folder, verbose)
  name_dict = config_dict.get("dns", {})
  for name in name_dict:
    expected = name_dict[name].get("expected")
    dns_logger.info(';'.join(str(x) for x in resolve(name, expected)))

  ping_logger = create_logger("ping", log_folder, verbose)
  for host in config_dict.get("ping",[]):
    ping_logger.info(';'.join(str(x) for x in ping(host)))

  ping6_logger = create_logger("ping6", log_folder, verbose)
  for host in config_dict.get("ping6", []):
    ping6_logger.info(';'.join(str(x) for x in ping6(host)))

  http_logger = create_logger("http", log_folder, verbose)
  url_dict = config_dict.get("url", {})
  for url in url_dict:
    http_logger.info(';'.join(str(x) for x in request(url, url_dict[url])))

def main():
  parser = argparse.ArgumentParser(
        description="Run network benchmarch.",
        )
  _ = parser.add_argument
  _('-l', '--logdir', default=".",
                      help="Directory where the logs are going to be placed.")
  _('-c', '--conf', help="Path to the configuration json file.")
  _('-v', '--verbose', action='store_true', 
                       help="Show the results on stdout.")
  _('-d', '--delay', default=random.randint(0, 30), 
                     help="Delay before start to run," \
                          "as this script can be called on cron.")

  config = parser.parse_args()

  print("Downloading %s..." % config.conf.strip())
  config_dict = load_configuration(config.conf)

  print("Waiting %s before start..." % config.delay)
  time.sleep(float(config.delay))

  run_all(config_dict, 
          log_folder=config.logdir, 
          verbose=config.verbose)