/boto-2.5.2/boto/manage/cmdshell.py

# · Python · 241 lines · 149 code · 29 blank · 63 comment · 15 complexity · f72bda4c9b7aecfa15b861cfc9d3e1af MD5 · raw file

  1. # Copyright (c) 2006-2009 Mitch Garnaat http://garnaat.org/
  2. #
  3. # Permission is hereby granted, free of charge, to any person obtaining a
  4. # copy of this software and associated documentation files (the
  5. # "Software"), to deal in the Software without restriction, including
  6. # without limitation the rights to use, copy, modify, merge, publish, dis-
  7. # tribute, sublicense, and/or sell copies of the Software, and to permit
  8. # persons to whom the Software is furnished to do so, subject to the fol-
  9. # lowing conditions:
  10. #
  11. # The above copyright notice and this permission notice shall be included
  12. # in all copies or substantial portions of the Software.
  13. #
  14. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
  15. # OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
  16. # ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
  17. # SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  18. # WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  19. # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
  20. # IN THE SOFTWARE.
  21. from boto.mashups.interactive import interactive_shell
  22. import boto
  23. import os
  24. import time
  25. import shutil
  26. import StringIO
  27. import paramiko
  28. import socket
  29. import subprocess
  30. class SSHClient(object):
  31. def __init__(self, server,
  32. host_key_file='~/.ssh/known_hosts',
  33. uname='root', ssh_pwd=None):
  34. self.server = server
  35. self.host_key_file = host_key_file
  36. self.uname = uname
  37. self._pkey = paramiko.RSAKey.from_private_key_file(server.ssh_key_file,
  38. password=ssh_pwd)
  39. self._ssh_client = paramiko.SSHClient()
  40. self._ssh_client.load_system_host_keys()
  41. self._ssh_client.load_host_keys(os.path.expanduser(host_key_file))
  42. self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  43. self.connect()
  44. def connect(self):
  45. retry = 0
  46. while retry < 5:
  47. try:
  48. self._ssh_client.connect(self.server.hostname,
  49. username=self.uname,
  50. pkey=self._pkey)
  51. return
  52. except socket.error, (value, message):
  53. if value == 61 or value == 111:
  54. print 'SSH Connection refused, will retry in 5 seconds'
  55. time.sleep(5)
  56. retry += 1
  57. else:
  58. raise
  59. except paramiko.BadHostKeyException:
  60. print "%s has an entry in ~/.ssh/known_hosts and it doesn't match" % self.server.hostname
  61. print 'Edit that file to remove the entry and then hit return to try again'
  62. raw_input('Hit Enter when ready')
  63. retry += 1
  64. except EOFError:
  65. print 'Unexpected Error from SSH Connection, retry in 5 seconds'
  66. time.sleep(5)
  67. retry += 1
  68. print 'Could not establish SSH connection'
  69. def open_sftp(self):
  70. return self._ssh_client.open_sftp()
  71. def get_file(self, src, dst):
  72. sftp_client = self.open_sftp()
  73. sftp_client.get(src, dst)
  74. def put_file(self, src, dst):
  75. sftp_client = self.open_sftp()
  76. sftp_client.put(src, dst)
  77. def open(self, filename, mode='r', bufsize=-1):
  78. """
  79. Open a file on the remote system and return a file-like object.
  80. """
  81. sftp_client = self.open_sftp()
  82. return sftp_client.open(filename, mode, bufsize)
  83. def listdir(self, path):
  84. sftp_client = self.open_sftp()
  85. return sftp_client.listdir(path)
  86. def isdir(self, path):
  87. status = self.run('[ -d %s ] || echo "FALSE"' % path)
  88. if status[1].startswith('FALSE'):
  89. return 0
  90. return 1
  91. def exists(self, path):
  92. status = self.run('[ -a %s ] || echo "FALSE"' % path)
  93. if status[1].startswith('FALSE'):
  94. return 0
  95. return 1
  96. def shell(self):
  97. """
  98. Start an interactive shell session on the remote host.
  99. """
  100. channel = self._ssh_client.invoke_shell()
  101. interactive_shell(channel)
  102. def run(self, command):
  103. """
  104. Execute a command on the remote host. Return a tuple containing
  105. an integer status and a two strings, the first containing stdout
  106. and the second containing stderr from the command.
  107. """
  108. boto.log.debug('running:%s on %s' % (command, self.server.instance_id))
  109. status = 0
  110. try:
  111. t = self._ssh_client.exec_command(command)
  112. except paramiko.SSHException:
  113. status = 1
  114. std_out = t[1].read()
  115. std_err = t[2].read()
  116. t[0].close()
  117. t[1].close()
  118. t[2].close()
  119. boto.log.debug('stdout: %s' % std_out)
  120. boto.log.debug('stderr: %s' % std_err)
  121. return (status, std_out, std_err)
  122. def run_pty(self, command):
  123. """
  124. Execute a command on the remote host with a pseudo-terminal.
  125. Returns a string containing the output of the command.
  126. """
  127. boto.log.debug('running:%s on %s' % (command, self.server.instance_id))
  128. channel = self._ssh_client.get_transport().open_session()
  129. channel.get_pty()
  130. channel.exec_command(command)
  131. return channel
  132. def close(self):
  133. transport = self._ssh_client.get_transport()
  134. transport.close()
  135. self.server.reset_cmdshell()
  136. class LocalClient(object):
  137. def __init__(self, server, host_key_file=None, uname='root'):
  138. self.server = server
  139. self.host_key_file = host_key_file
  140. self.uname = uname
  141. def get_file(self, src, dst):
  142. shutil.copyfile(src, dst)
  143. def put_file(self, src, dst):
  144. shutil.copyfile(src, dst)
  145. def listdir(self, path):
  146. return os.listdir(path)
  147. def isdir(self, path):
  148. return os.path.isdir(path)
  149. def exists(self, path):
  150. return os.path.exists(path)
  151. def shell(self):
  152. raise NotImplementedError('shell not supported with LocalClient')
  153. def run(self):
  154. boto.log.info('running:%s' % self.command)
  155. log_fp = StringIO.StringIO()
  156. process = subprocess.Popen(self.command, shell=True, stdin=subprocess.PIPE,
  157. stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  158. while process.poll() == None:
  159. time.sleep(1)
  160. t = process.communicate()
  161. log_fp.write(t[0])
  162. log_fp.write(t[1])
  163. boto.log.info(log_fp.getvalue())
  164. boto.log.info('output: %s' % log_fp.getvalue())
  165. return (process.returncode, log_fp.getvalue())
  166. def close(self):
  167. pass
  168. class FakeServer(object):
  169. """
  170. A little class to fake out SSHClient (which is expecting a
  171. :class`boto.manage.server.Server` instance. This allows us
  172. to
  173. """
  174. def __init__(self, instance, ssh_key_file):
  175. self.instance = instance
  176. self.ssh_key_file = ssh_key_file
  177. self.hostname = instance.dns_name
  178. self.instance_id = self.instance.id
  179. def start(server):
  180. instance_id = boto.config.get('Instance', 'instance-id', None)
  181. if instance_id == server.instance_id:
  182. return LocalClient(server)
  183. else:
  184. return SSHClient(server)
  185. def sshclient_from_instance(instance, ssh_key_file,
  186. host_key_file='~/.ssh/known_hosts',
  187. user_name='root', ssh_pwd=None):
  188. """
  189. Create and return an SSHClient object given an
  190. instance object.
  191. :type instance: :class`boto.ec2.instance.Instance` object
  192. :param instance: The instance object.
  193. :type ssh_key_file: str
  194. :param ssh_key_file: A path to the private key file used
  195. to log into instance.
  196. :type host_key_file: str
  197. :param host_key_file: A path to the known_hosts file used
  198. by the SSH client.
  199. Defaults to ~/.ssh/known_hosts
  200. :type user_name: str
  201. :param user_name: The username to use when logging into
  202. the instance. Defaults to root.
  203. :type ssh_pwd: str
  204. :param ssh_pwd: The passphrase, if any, associated with
  205. private key.
  206. """
  207. s = FakeServer(instance, ssh_key_file)
  208. return SSHClient(s, host_key_file, user_name, ssh_pwd)