vstf/vstf/common/utils.py

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