python amon_remote_execution.py

Posted 2021-05-09

import subprocess
import json
import logging
import os

logger = logging.getLogger(__name__)
from django.conf import settings

if not settings.DISABLED_REMOTE_ENGINE:
	import salt.client
else:
	salt = None

from django.template.loader import render_to_string


class AmonSaltClient(object):


	def __init__(self, minion_id=None):
		try:
			self.client = salt.client.LocalClient()
		except:
			self.client = None
			
		self.timeout = 120 	# 2 minutes default timeout
		self.minion_id = minion_id
		self.salt_path = '/salt/amon/'

		# Exit codes 0 for OK, 1 for WARNING, 2 for CRITICAL, and 3 or greater to indicate UNKNOWN or CUSTOM.
		self.exit_code = '; bash -c printf "\nexit_code: ${PIPESTATUS[@]}"'
		# self.exit_code = ""


	def copy_execute_cleanup_file(self, file=None, language=None, params=None):
		path_on_minion = self.copy_file_object(file=file)

		file_output = self.execute_file(filename=path_on_minion, language=language, params=params)
		# self.remove_file(filename=file.filename)

		return file_output


	def get_file(self, filename=None):
		result = self.client.cmd(self.minion_id, 'cp.get_file_str', [filename], timeout=self.timeout)
		execution_result = result.get(self.minion_id, False)

		return execution_result


	# Exit code in checks, not present in the remote console
	def execute_command(self, command=None, exit_code=None):
		if exit_code != None and exit_code != False:
			command_string = "{command} {exit_code}".format(command=command, exit_code=self.exit_code)
		else:
			command_string = command

		result = self.client.cmd(self.minion_id, 'cmd.run', [command_string], timeout=self.timeout, python_shell=True)

		command_result = result.get(self.minion_id, [])

		return command_result


	# Wrap language executions in a simple bash file, instead of hacky terminal pipes
	def _create_and_copy_command_wrapper(self, command=None):
		rendered = render_to_string('remote/execute_wrapper.sh', {'command': command})
		filename = "execute_wrapper.sh"
		file_path = '{salt_path}{filename}'.format(salt_path=self.salt_path, filename=filename)
		with open(file_path, 'w+') as f:
			f.write(rendered)

		source = "salt://{0}".format(filename)
		dest = '/tmp/{0}'.format(filename)

		result = self.client.cmd(self.minion_id, 'cp.get_file', [source, dest], timeout=self.timeout)


	def execute_file(self, filename=None,  language=None, params=None):
		command = "{language} {filename}".format(filename=filename, language=language)
		if params:
			command = "{command} {params} ".format(command=command, params=params)

		print params

		self._create_and_copy_command_wrapper(command=command)
		actual_command = "bash /tmp/execute_wrapper.sh"

		result = self.client.cmd(self.minion_id, 'cmd.run', [actual_command], timeout=self.timeout, python_shell=True)


		execution_result = result.get(self.minion_id, False)

		return execution_result


	def remove_file(self, filename=None, ):
		dest = '/tmp/{0}'.format(filename)
		command = "rm {0}".format(dest)

		result = self.client.cmd(self.minion_id, 'cmd.run', [command], timeout=self.timeout)
		cleanup_result = result.get(self.minion_id, False)

		try:
			os.remove('{salt_path}{filename}'.format(salt_path=self.salt_path, filename=filename))
		except:
			pass

		return cleanup_result

	# The file has been created elsewhere and already copied to /salt/amon
	def copy_file(self,  destination=None, filename=None):
		source = "salt://{0}".format(filename)

		result = self.client.cmd(self.minion_id, 'cp.get_file', [source, destination], timeout=self.timeout)
		copy_result = result.get(self.minion_id, False)

		return copy_result

	def copy_file_object(self, file=None, destination=None):
		filename = file.filename
		file_path = '{salt_path}{filename}'.format(salt_path=self.salt_path, filename=filename)
		with open(file_path, 'w+') as f:
			f.write(file.read())

		source = "salt://{0}".format(filename)

		dest = '/tmp/{0}'.format(filename)
		if destination != None:
			dest = destination

		result = self.client.cmd(self.minion_id, 'cp.get_file', [source, dest], timeout=self.timeout)
		copy_result = result.get(self.minion_id, False)

		return copy_result

	def ping_minion(self):
		result = self.client.cmd(self.minion_id, 'test.ping', timeout=10)
		ping_result = result.get(self.minion_id, False)

		return ping_result


	def check_minion_key_status(self):
		minion_key_accepted = False

		try:
			salt_keys = subprocess.Popen(['salt-key','-L', '--out=json'], stdout=subprocess.PIPE, close_fds=True).communicate()[0]
		except:
			salt_keys = False
			logger.exception("Can't process Salt Keys")

		try:
			keys_dict = json.loads(salt_keys)
		except:
			keys_dict = {}

		if salt_keys:
			minions_keys_list = keys_dict.get('minions')

			if self.minion_id in minions_keys_list:
				minion_key_accepted = True

		return minion_key_accepted


amon_salt_client = AmonSaltClient()