deploy/ssh_client.py

   1 ###############################################################################
   2 # Copyright (c) 2015 Ericsson AB and others.
   3 # szilard.cserey@ericsson.com
   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
  11 import paramiko
  12 import scp
  13
  14 from common import (
  15     log,
  16     err,
  17 )
  18
  19 TIMEOUT = 600
  20
  21 class SSHClient(object):
  22
  23     def __init__(self, host, username, password):
  24         self.host = host
  25         self.username = username
  26         self.password = password
  27         self.client = None
  28
  29     def open(self, timeout=TIMEOUT):
  30         self.client = paramiko.SSHClient()
  31         self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  32         self.client.connect(self.host, username=self.username,
  33                             password=self.password, look_for_keys=False,
  34                             timeout=timeout)
  35
  36     def close(self):
  37         if self.client is not None:
  38             self.client.close()
  39             self.client = None
  40
  41     def __enter__(self):
  42         self.open()
  43         return self
  44
  45     def __exit__(self, type, value, traceback):
  46         self.close()
  47
  48     def exec_cmd(self, command, check=True, sudo=False, timeout=TIMEOUT):
  49         if sudo and self.username != 'root':
  50             command = "sudo -S -p '' %s" % command
  51         stdin, stdout, stderr = self.client.exec_command(command,
  52                                                          timeout=timeout)
  53         if sudo:
  54             stdin.write(self.password + '\n')
  55             stdin.flush()
  56         response = stdout.read().strip()
  57         error = stderr.read().strip()
  58
  59         if check:
  60             if error:
  61                 self.close()
  62                 raise Exception(error)
  63             else:
  64                 return response
  65         return response, error
  66
  67     def run(self, command):
  68         transport = self.client.get_transport()
  69         transport.set_keepalive(1)
  70         chan = transport.open_session()
  71         chan.exec_command(command)
  72         while not chan.exit_status_ready():
  73             if chan.recv_ready():
  74                 data = chan.recv(1024)
  75                 while data:
  76                     log(data.strip())
  77                     data = chan.recv(1024)
  78
  79             if chan.recv_stderr_ready():
  80                 error_buff = chan.recv_stderr(1024)
  81                 while error_buff:
  82                     log(error_buff.strip())
  83                     error_buff = chan.recv_stderr(1024)
  84         return chan.recv_exit_status()
  85
  86     def scp_get(self, remote, local='.', dir=False):
  87         try:
  88             with scp.SCPClient(self.client.get_transport()) as _scp:
  89                 _scp.get(remote, local, dir)
  90         except Exception as e:
  91             err(e)
  92
  93     def scp_put(self, local, remote='.', dir=False):
  94         try:
  95             with scp.SCPClient(self.client.get_transport()) as _scp:
  96                 _scp.put(local, remote, dir)
  97         except Exception as e:
  98             err(e)