vstf/vstf/common/utils.py

   1 ##############################################################################
   2 # Copyright (c) 2015 Huawei Technologies Co.,Ltd and others.
   3 #
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9
  10 import re
  11 import logging
  12 import subprocess
  13 import random
  14 import os
  15 import signal
  16 import time
  17 from StringIO import StringIO
  18
  19 LOG = logging.getLogger(__name__)
  20
  21
  22 def info():
  23     def _deco(func):
  24         def __deco(*args, **kwargs):
  25             if "shell" in kwargs and kwargs["shell"]:
  26                 LOG.info(args[0])
  27             else:
  28                 LOG.info(' '.join(args[0]))
  29             return func(*args, **kwargs)
  30         return __deco
  31     return _deco
  32
  33
  34 @info()
  35 def call(cmd, shell=False):
  36     ret = subprocess.call(cmd, shell=shell)
  37     if ret != 0:
  38         LOG.info("warning: %s not success.", cmd)
  39
  40
  41 @info()
  42 def check_call(cmd, shell=False):
  43     subprocess.check_call(cmd, shell=shell)
  44
  45
  46 @info()
  47 def check_output(cmd, shell=False):
  48     return subprocess.check_output(cmd, shell=shell)
  49
  50
  51 @info()
  52 def my_popen(cmd, shell=False, stdout=None, stderr=None):
  53     return subprocess.Popen(cmd, shell=shell, stdout=stdout, stderr=stderr)
  54
  55
  56 def ping(ip):
  57     cmd = "ping -w2 -c1 %s" % ip
  58     p = my_popen(cmd, shell=True)
  59     return 0 == p.wait()
  60
  61
  62 def get_device_name(bdf):
  63     path = '/sys/bus/pci/devices/0000:%s/net/' % bdf
  64     path1 = '/sys/bus/pci/devices/0000:%s/virtio*/net/' % bdf
  65     if os.path.exists(path):
  66         device = check_output("ls " + path, shell=True).strip()
  67         return device
  68     else:  # virtio driver
  69         try:
  70             device = check_output("ls " + path1, shell=True).strip()
  71             return device
  72         except Exception:
  73             return None
  74
  75
  76 def my_sleep(delay):
  77     LOG.info('sleep %s' % delay)
  78     time.sleep(delay)
  79
  80
  81 def my_mkdir(filepath):
  82     try:
  83         LOG.info("mkdir -p %s" % filepath)
  84         os.makedirs(filepath)
  85     except OSError, e:
  86         if e.errno == 17:
  87             LOG.info("! %s already exists" % filepath)
  88         else:
  89             raise
  90
  91
  92 def get_eth_by_bdf(bdf):
  93     bdf = bdf.replace(' ', '')
  94     path = '/sys/bus/pci/devices/0000:%s/net/' % bdf
  95     if os.path.exists(path):
  96         device = check_output("ls " + path, shell=True).strip()
  97     else:
  98         raise Exception("cann't get device name of bdf:%s" % bdf)
  99     return device
 100
 101
 102 def check_and_kill(process):
 103     cmd = "ps -ef | grep -v grep | awk '{print $8}' | grep -w %s | wc -l" % process
 104     out = check_output(cmd, shell=True)
 105     if int(out):
 106         check_call(['killall', process])
 107
 108
 109 def list_mods():
 110     return check_output("lsmod | sed 1,1d | awk '{print $1}'", shell=True).split()
 111
 112
 113 def check_and_rmmod(mod):
 114     if mod in list_mods():
 115         check_call(['rmmod', mod])
 116
 117
 118 def kill_by_name(process):
 119     out = check_output(['ps', '-A'])
 120     for line in out.splitlines():
 121         values = line.split()
 122         pid, name = values[0], values[3]
 123         if process == name:
 124             pid = int(pid)
 125             os.kill(pid, signal.SIGKILL)
 126             LOG.info("os.kill(%s)" % pid)
 127
 128
 129 def ns_cmd(ns, cmd):
 130     netns_exec_str = "ip netns exec %s "
 131     if ns in (None, 'null', 'None', 'none'):
 132         pass
 133     else:
 134         cmd = (netns_exec_str % ns) + cmd
 135     return cmd
 136
 137
 138 def randomMAC():
 139     mac = [0x00, 0x16, 0x3e,
 140            random.randint(0x00, 0x7f),
 141            random.randint(0x00, 0xff),
 142            random.randint(0x00, 0xff)]
 143     return ':'.join(map(lambda x: "%02x" % x, mac))
 144
 145
 146 class IPCommandHelper(object):
 147     def __init__(self, ns=None):
 148         self.devices = []
 149         self.macs = []
 150         self.device_mac_map = {}
 151         self.mac_device_map = {}
 152         self.bdf_device_map = {}
 153         self.device_bdf_map = {}
 154         self.mac_bdf_map = {}
 155         self.bdf_mac_map = {}
 156         cmd = "ip link"
 157         if ns:
 158             cmd = "ip netns exec %s " % ns + cmd
 159         buf = check_output(cmd, shell=True)
 160         sio = StringIO(buf)
 161         for line in sio:
 162             m = re.match(r'^\d+:(.*):.*', line)
 163             if m and m.group(1).strip() != "lo":
 164                 device = m.group(1).strip()
 165                 self.devices.append(device)
 166                 mac = self._get_mac(ns, device)
 167                 self.macs.append(mac)
 168         for device, mac in zip(self.devices, self.macs):
 169             self.device_mac_map[device] = mac
 170             self.mac_device_map[mac] = device
 171
 172         cmd = "ethtool -i %s"
 173         if ns:
 174             cmd = "ip netns exec %s " % ns + cmd
 175         for device in self.devices:
 176             buf = check_output(cmd % device, shell=True)
 177             bdfs = re.findall(r'^bus-info: \d{4}:(\d{2}:\d{2}\.\d*)$', buf, re.MULTILINE)
 178             if bdfs:
 179                 self.bdf_device_map[bdfs[0]] = device
 180                 self.device_bdf_map[device] = bdfs[0]
 181                 mac = self.device_mac_map[device]
 182                 self.mac_bdf_map[mac] = bdfs[0]
 183                 self.bdf_mac_map[bdfs[0]] = mac
 184
 185     @staticmethod
 186     def _get_mac(ns, device):
 187         cmd = "ip addr show dev %s" % device
 188         if ns:
 189             cmd = "ip netns exec %s " % ns + cmd
 190         buf = check_output(cmd, shell=True)
 191         macs = re.compile(r"[A-F0-9]{2}(?::[A-F0-9]{2}){5}", re.IGNORECASE | re.MULTILINE)
 192         for mac in macs.findall(buf):
 193             if mac.lower() not in ('00:00:00:00:00:00', 'ff:ff:ff:ff:ff:ff'):
 194                 return mac
 195         return None
 196
 197     def get_device_verbose(self, identity):
 198         if identity in self.device_mac_map:
 199             device = identity
 200         elif identity in self.bdf_device_map:
 201             device = self.bdf_device_map[identity]
 202         elif identity in self.mac_device_map:
 203             device = self.mac_device_map[identity]
 204         else:
 205             raise Exception("cann't find the device by identity:%s" % identity)
 206         detail = {
 207             'bdf': self.device_bdf_map[device] if device in self.device_bdf_map else None,
 208             'iface': device,
 209             'mac': self.device_mac_map[device] if device in self.device_mac_map else None,
 210         }
 211         return detail
 212
 213
 214 class AttrDict(dict):
 215     """A dictionary with attribute-style access. It maps attribute access to
 216     the real dictionary.  """
 217
 218     def __init__(self, init={}):
 219         dict.__init__(self, init)
 220
 221     def __getstate__(self):
 222         return self.__dict__.items()
 223
 224     def __setstate__(self, items):
 225         for key, val in items:
 226             self.__dict__[key] = val
 227
 228     def __repr__(self):
 229         return "%s(%s)" % (self.__class__.__name__, dict.__repr__(self))
 230
 231     def __setitem__(self, key, value):
 232         return super(AttrDict, self).__setitem__(key, value)
 233
 234     def __getitem__(self, name):
 235         return super(AttrDict, self).__getitem__(name)
 236
 237     def __delitem__(self, name):
 238         return super(AttrDict, self).__delitem__(name)
 239
 240     __getattr__ = __getitem__
 241     __setattr__ = __setitem__
 242
 243     def copy(self):
 244         ch = AttrDict(self)
 245         return ch
 246
 247
 248 if __name__ == "__main__":
 249     ipcmd = IPCommandHelper()
 250     print ipcmd.device_mac_map
 251     print ipcmd.mac_device_map
 252     print ipcmd.bdf_device_map
 253     print ipcmd.device_bdf_map
 254     print ipcmd.mac_bdf_map
 255     print ipcmd.bdf_mac_map
 256     print ipcmd.get_device_verbose("tap0")