deploy/ssh_client.py

   1 ###############################################################################
   2 # Copyright (c) 2015 Ericsson AB and others.
   3 # szilard.cserey@ericsson.com
   4 # peter.barabas@ericsson.com
   5 # All rights reserved. This program and the accompanying materials
   6 # are made available under the terms of the Apache License, Version 2.0
   7 # which accompanies this distribution, and is available at
   8 # http://www.apache.org/licenses/LICENSE-2.0
   9 ###############################################################################
  10
  11
  12 import paramiko
  13 import scp
  14
  15 from common import (
  16     log,
  17     err,
  18 )
  19
  20 TIMEOUT = 600
  21
  22 class SSHClient(object):
  23
  24     def __init__(self, host, username, password):
  25         self.host = host
  26         self.username = username
  27         self.password = password
  28         self.client = None
  29
  30     def open(self, timeout=TIMEOUT):
  31         self.client = paramiko.SSHClient()
  32         self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  33         self.client.connect(self.host, username=self.username,
  34                             password=self.password, look_for_keys=False,
  35                             timeout=timeout)
  36
  37     def close(self):
  38         if self.client is not None:
  39             self.client.close()
  40             self.client = None
  41
  42     def __enter__(self):
  43         self.open()
  44         return self
  45
  46     def __exit__(self, type, value, traceback):
  47         self.close()
  48
  49     def exec_cmd(self, command, check=True, sudo=False, timeout=TIMEOUT):
  50         if sudo and self.username != 'root':
  51             command = "sudo -S -p '' %s" % command
  52         stdin, stdout, stderr = self.client.exec_command(command,
  53                                                          timeout=timeout)
  54         if sudo:
  55             stdin.write(self.password + '\n')
  56             stdin.flush()
  57         response = stdout.read().strip()
  58         error = stderr.read().strip()
  59
  60         if check:
  61             if error:
  62                 self.close()
  63                 raise Exception(error)
  64             else:
  65                 return response
  66         return response, error
  67
  68     def run(self, command):
  69         transport = self.client.get_transport()
  70         transport.set_keepalive(1)
  71         chan = transport.open_session()
  72         chan.exec_command(command)
  73         while not chan.exit_status_ready():
  74             if chan.recv_ready():
  75                 data = chan.recv(1024)
  76                 while data:
  77                     log(data.strip())
  78                     data = chan.recv(1024)
  79
  80             if chan.recv_stderr_ready():
  81                 error_buff = chan.recv_stderr(1024)
  82                 while error_buff:
  83                     log(error_buff.strip())
  84                     error_buff = chan.recv_stderr(1024)
  85         return chan.recv_exit_status()
  86
  87     def scp_get(self, remote, local='.', dir=False):
  88         try:
  89             with scp.SCPClient(self.client.get_transport(),
  90                                sanitize=lambda x: x,
  91                                socket_timeout=TIMEOUT) as _scp:
  92                 _scp.get(remote, local, dir)
  93         except Exception as e:
  94             err(e)
  95
  96     def scp_put(self, local, remote='.', dir=False):
  97         try:
  98             with scp.SCPClient(self.client.get_transport(),
  99                                sanitize=lambda x: x,
 100                                socket_timeout=TIMEOUT) as _scp:
 101                 _scp.put(local, remote, dir)
 102         except Exception as e:
 103             err(e)