shell.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. # -*- coding: utf-8 -*-
  2. # Licensed under the Apache License, Version 2.0 (the "License");
  3. # you may not use this file except in compliance with the License.
  4. # You may obtain a copy of the License at
  5. #
  6. # http://www.apache.org/licenses/LICENSE-2.0
  7. #
  8. # Unless required by applicable law or agreed to in writing, software
  9. # distributed under the License is distributed on an "AS IS" BASIS,
  10. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
  11. # implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import re
  15. import os
  16. import types
  17. import logging
  18. import subprocess
  19. from ..exceptions import ExecuteRuntimeError
  20. from ..exceptions import NetworkError
  21. from ..exceptions import ScriptRuntimeError
  22. from .strings import mask_string
  23. block_fmt = ("\n============= %(title)s ==========\n%(content)s\n"
  24. "======== END OF %(title)s ========")
  25. def execute(cmd, workdir=None, can_fail=True, mask_list=None,
  26. use_shell=False, log=True):
  27. """
  28. Runs shell command cmd. If can_fail is set to False
  29. ExecuteRuntimeError is raised if command returned non-zero return
  30. code. Otherwise
  31. """
  32. mask_list = mask_list or []
  33. repl_list = [("'", "'\\''")]
  34. if not isinstance(cmd, types.StringType):
  35. import pipes
  36. masked = ' '.join((pipes.quote(i) for i in cmd))
  37. else:
  38. masked = cmd
  39. masked = mask_string(masked, mask_list, repl_list)
  40. if log:
  41. logging.info("Executing command:\n%s" % masked)
  42. environ = os.environ
  43. environ['LANG'] = 'en_US.UTF8'
  44. proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
  45. stderr=subprocess.PIPE, cwd=workdir,
  46. shell=use_shell, close_fds=True,
  47. env=environ)
  48. out, err = proc.communicate()
  49. masked_out = mask_string(out, mask_list, repl_list)
  50. masked_err = mask_string(err, mask_list, repl_list)
  51. if log:
  52. logging.debug(block_fmt % {'title': 'STDOUT', 'content': masked_out})
  53. if proc.returncode:
  54. if log:
  55. logging.debug(block_fmt % {'title': 'STDERR',
  56. 'content': masked_err})
  57. if can_fail:
  58. msg = ('Failed to execute command, '
  59. 'stdout: %s\nstderr: %s' %
  60. (masked_out, masked_err))
  61. raise ExecuteRuntimeError(msg, stdout=out, stderr=err)
  62. return proc.returncode, out
  63. class ScriptRunner(object):
  64. _pkg_search = 'rpm -q --whatprovides'
  65. def __init__(self, ip=None):
  66. self.script = []
  67. self.ip = ip
  68. def append(self, s):
  69. self.script.append(s)
  70. def clear(self):
  71. self.script = []
  72. def execute(self, can_fail=True, mask_list=None, log=True):
  73. mask_list = mask_list or []
  74. repl_list = [("'", "'\\''")]
  75. script = "\n".join(self.script)
  76. masked = mask_string(script, mask_list, repl_list)
  77. if log:
  78. logging.info("[%s] Executing script:\n%s" %
  79. (self.ip or 'localhost', masked))
  80. _PIPE = subprocess.PIPE # pylint: disable=E1101
  81. if self.ip:
  82. cmd = ["ssh", "-o", "StrictHostKeyChecking=no",
  83. "-o", "UserKnownHostsFile=/dev/null",
  84. "root@%s" % self.ip, "bash -x"]
  85. else:
  86. cmd = ["bash", "-x"]
  87. environ = os.environ
  88. environ['LANG'] = 'en_US.UTF8'
  89. obj = subprocess.Popen(cmd, stdin=_PIPE, stdout=_PIPE, stderr=_PIPE,
  90. close_fds=True, shell=False, env=environ)
  91. script = "function t(){ exit $? ; } \n trap t ERR \n" + script
  92. out, err = obj.communicate(script)
  93. masked_out = mask_string(out, mask_list, repl_list)
  94. masked_err = mask_string(err, mask_list, repl_list)
  95. if log:
  96. logging.debug(block_fmt % {'title': 'STDOUT',
  97. 'content': masked_out})
  98. if obj.returncode:
  99. if log:
  100. logging.debug(block_fmt % {'title': 'STDERR',
  101. 'content': masked_err})
  102. if can_fail:
  103. pattern = (r'^ssh\:')
  104. if re.search(pattern, err):
  105. raise NetworkError(masked_err, stdout=out, stderr=err)
  106. else:
  107. msg = ('Failed to run remote script, '
  108. 'stdout: %s\nstderr: %s' %
  109. (masked_out, masked_err))
  110. raise ScriptRuntimeError(msg, stdout=out, stderr=err)
  111. return obj.returncode, out
  112. def template(self, src, dst, varsdict):
  113. with open(src) as fp:
  114. content = fp.read() % varsdict
  115. self.append("cat > %s <<- EOF\n%s\nEOF\n" % (dst, content))
  116. def if_not_exists(self, path, command):
  117. self.append("[ -e %s ] || %s" % (path, command))
  118. def if_exists(self, path, command):
  119. self.append("[ -e %s ] && %s" % (path, command))
  120. def if_installed(self, pkg, command):
  121. self.append("%s %s && %s" % (self._pkg_search, pkg, command))
  122. def if_not_installed(self, pkg, command):
  123. self.append("%s %s || %s" % (self._pkg_search, pkg, command))
  124. def chown(self, target, uid, gid):
  125. self.append("chown %s:%s %s" % (uid, gid, target))
  126. def chmod(self, target, mode):
  127. self.append("chmod %s %s" % (mode, target))