VNFs/DPPD-PROX/helper-scripts/openstackrapid/prox_ctrl.py

   1 ##
   2 ## Copyright (c) 2010-2017 Intel Corporation
   3 ##
   4 ## Licensed under the Apache License, Version 2.0 (the "License");
   5 ## you may not use this file except in compliance with the License.
   6 ## You may obtain a copy of the License at
   7 ##
   8 ##     http://www.apache.org/licenses/LICENSE-2.0
   9 ##
  10 ## Unless required by applicable law or agreed to in writing, software
  11 ## distributed under the License is distributed on an "AS IS" BASIS,
  12 ## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 ## See the License for the specific language governing permissions and
  14 ## limitations under the License.
  15 ##
  16
  17 from __future__ import print_function
  18
  19 import os
  20 import subprocess
  21 import socket
  22
  23 class prox_ctrl(object):
  24     def __init__(self, ip, key=None, user=None):
  25         self._ip   = ip
  26         self._key  = key
  27         self._user = user
  28         self._children = []
  29         self._proxsock = []
  30
  31     def ip(self):
  32         return self._ip
  33
  34     def connect(self):
  35         """Simply try to run 'true' over ssh on remote system.
  36         On failure, raise RuntimeWarning exception when possibly worth
  37         retrying, and raise RuntimeError exception otherwise.
  38         """
  39         return self.run_cmd('true', True)
  40
  41     def close(self):
  42         """Must be called before program termination."""
  43         for prox in self._proxsock:
  44             prox.quit()
  45         children = len(self._children)
  46         if children == 0:
  47             return
  48         if children > 1:
  49             print('Waiting for %d child processes to complete ...' % children)
  50         for child in self._children:
  51             ret = os.waitpid(child[0], os.WNOHANG)
  52             if ret[0] == 0:
  53                 print("Waiting for child process '%s' to complete ..." % child[1])
  54                 ret = os.waitpid(child[0], 0)
  55             rc = ret[1]
  56             if os.WIFEXITED(rc):
  57                 if os.WEXITSTATUS(rc) == 0:
  58                     print("Child process '%s' completed successfully" % child[1])
  59                 else:
  60                     print("Child process '%s' returned exit status %d" % (
  61                             child[1], os.WEXITSTATUS(rc)))
  62             elif os.WIFSIGNALED(rc):
  63                 print("Child process '%s' exited on signal %d" % (
  64                         child[1], os.WTERMSIG(rc)))
  65             else:
  66                 print("Wait status for child process '%s' is 0x%04x" % (
  67                         child[1], rc))
  68
  69     def run_cmd(self, command, _connect=False):
  70         """Execute command over ssh on remote system.
  71         Wait for remote command completion.
  72         Return command output (combined stdout and stderr).
  73         _connect argument is reserved for connect() method.
  74         """
  75         cmd = self._build_ssh(command)
  76         try:
  77             return subprocess.check_output(cmd, stderr=subprocess.STDOUT)
  78         except subprocess.CalledProcessError as ex:
  79             if _connect and ex.returncode == 255:
  80                 raise RuntimeWarning(ex.output.strip())
  81             raise RuntimeError('ssh returned exit status %d:\n%s'
  82                     % (ex.returncode, ex.output.strip()))
  83
  84     def fork_cmd(self, command, name=None):
  85         """Execute command over ssh on remote system, in a child process.
  86         Do not wait for remote command completion.
  87         Return child process id.
  88         """
  89         if name is None:
  90             name = command
  91         cmd = self._build_ssh(command)
  92         pid = os.fork()
  93         if (pid != 0):
  94             # In the parent process
  95             self._children.append((pid, name))
  96             return pid
  97         # In the child process: use os._exit to terminate
  98         try:
  99             # Actually ignore output on success, but capture stderr on failure
 100             subprocess.check_output(cmd, stderr=subprocess.STDOUT)
 101         except subprocess.CalledProcessError as ex:
 102             raise RuntimeError("Child process '%s' failed:\n"
 103                     'ssh returned exit status %d:\n%s'
 104                     % (name, ex.returncode, ex.output.strip()))
 105         os._exit(0)
 106
 107     def prox_sock(self, port=8474):
 108         """Connect to the PROX instance on remote system.
 109         Return a prox_sock object on success, None on failure.
 110         """
 111         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 112         try:
 113             sock.connect((self._ip, port))
 114             prox = prox_sock(sock)
 115             self._proxsock.append(prox)
 116             return prox
 117         except:
 118             return None
 119
 120     def scp_put(self, src, dst):
 121         """Copy src file from local system to dst on remote system."""
 122         cmd = [ 'scp',
 123                 '-B',
 124                 '-oStrictHostKeyChecking=no',
 125                 '-oUserKnownHostsFile=/dev/null',
 126                 '-oLogLevel=ERROR' ]
 127         if self._key is not None:
 128             cmd.extend(['-i', self._key])
 129         cmd.append(src)
 130         remote = ''
 131         if self._user is not None:
 132             remote += self._user + '@'
 133         remote += self._ip + ':' + dst
 134         cmd.append(remote)
 135         try:
 136             # Actually ignore output on success, but capture stderr on failure
 137             subprocess.check_output(cmd, stderr=subprocess.STDOUT)
 138         except subprocess.CalledProcessError as ex:
 139             raise RuntimeError('scp returned exit status %d:\n%s'
 140                     % (ex.returncode, ex.output.strip()))
 141
 142     def _build_ssh(self, command):
 143         cmd = [ 'ssh',
 144                 '-oBatchMode=yes',
 145                 '-oStrictHostKeyChecking=no',
 146                 '-oUserKnownHostsFile=/dev/null',
 147                 '-oLogLevel=ERROR' ]
 148         if self._key is not None:
 149             cmd.extend(['-i', self._key])
 150         remote = ''
 151         if self._user is not None:
 152             remote += self._user + '@'
 153         remote += self._ip
 154         cmd.append(remote)
 155         cmd.append(command)
 156         return cmd
 157
 158 class prox_sock(object):
 159     def __init__(self, sock):
 160         self._sock = sock
 161         self._rcvd = b''
 162
 163     def quit(self):
 164         if self._sock is not None:
 165             self._send('quit')
 166             self._sock.close()
 167             self._sock = None
 168
 169     def start(self, cores):
 170         self._send('start %s' % ','.join(map(str, cores)))
 171
 172     def stop(self, cores):
 173         self._send('stop %s' % ','.join(map(str, cores)))
 174
 175     def speed(self, speed, cores, tasks=None):
 176         if tasks is None:
 177             tasks = [ 0 ] * len(cores)
 178         elif len(tasks) != len(cores):
 179             raise ValueError('cores and tasks must have the same len')
 180         for (core, task) in zip(cores, tasks):
 181             self._send('speed %s %s %s' % (core, task, speed))
 182
 183     def reset_stats(self):
 184         self._send('reset stats')
 185
 186     def core_stats(self, cores, task=0):
 187         rx = tx = drop = tsc = hz = 0
 188         self._send('core stats %s %s' % (','.join(map(str, cores)), task))
 189         for core in cores:
 190             stats = self._recv().split(',')
 191             rx += int(stats[0])
 192             tx += int(stats[1])
 193             drop += int(stats[2])
 194             tsc = int(stats[3])
 195             hz = int(stats[4])
 196         return rx, tx, drop, tsc, hz
 197
 198     def set_random(self, cores, task, offset, mask, length):
 199         self._send('set random %s %s %s %s %s' % (','.join(map(str, cores)), task, offset, mask, length))
 200
 201     def _send(self, cmd):
 202         """Append LF and send command to the PROX instance."""
 203         if self._sock is None:
 204             raise RuntimeError("PROX socket closed, cannot send '%s'" % cmd)
 205         self._sock.sendall(cmd.encode() + b'\n')
 206
 207     def _recv(self):
 208         """Receive response from PROX instance, and return it with LF removed."""
 209         if self._sock is None:
 210             raise RuntimeError("PROX socket closed, cannot receive anymore")
 211         pos = self._rcvd.find(b'\n')
 212         while pos == -1:
 213             self._rcvd += self._sock.recv(256)
 214             pos = self._rcvd.find(b'\n')
 215         rsp = self._rcvd[:pos]
 216         self._rcvd = self._rcvd[pos+1:]
 217         return rsp.decode()
 218