awsr

#!/usr/bin/python
# Licensed under http://www.apache.org/licenses/LICENSE-2.0

import os
import sys
import tempfile

rd = tempfile.gettempdir() + "/awsr_rd"
wr = tempfile.gettempdir() + "/awsr_wr"


def run_client():
	with open(rd, "w") as out:
		out.write(" ".join(sys.argv))
		out.write("\n")
		if "-" in sys.argv:
			while True:
				data = sys.stdin.read(4096)
				if data == "":
					break
				out.write(data)

	last = None
	with open(wr, "r") as inp:
		while True:
			result = inp.read(4096)
			if len(result) == 0:
				break
			if last is not None:
				sys.stdout.write(last)
			sys.stdout.write(result[:-1])
			last = result[-1:]
	return ord(last)


def run_daemon():
	from awscli.clidriver import CLIOperationCaller, LOG, create_clidriver, HISTORY_RECORDER

	def patchedInit(self, session):
		self._session = session
		self._client = None

	def patchedInvoke(self, service_name, operation_name, parameters, parsed_globals):
		if self._client is None:
			LOG.debug("Creating new %s client" % service_name)
			self._client = self._session.create_client(
				service_name, region_name=parsed_globals.region,
				endpoint_url=parsed_globals.endpoint_url,
				verify=parsed_globals.verify_ssl)
		client = self._client

		response = self._make_client_call(
			client, operation_name, parameters, parsed_globals)
		self._display_response(operation_name, response, parsed_globals)
		return 0

	CLIOperationCaller.__init__ = patchedInit
	CLIOperationCaller.invoke = patchedInvoke

	driver = create_clidriver()
	while True:
		with open(rd, "r") as inp:
			os.dup2(inp.fileno(), 0)
			args = bytes()
			while True:
				ch = os.read(inp.fileno(), 1)
				if ch == b"\n" or ch == b"":
					break
				args += ch
			args = args.decode("utf-8").split(" ")[1:]

			if len(args) > 0 and args[0] == "exit":
				sys.exit(0)

			sys.stdin = inp
			old_stdout = sys.stdout
			try:
				with open(wr, "w") as sys.stdout:
					try:
						rc = driver.main(args)
					except SystemExit as e:
						rc = 2
					# encode return code in the last byte of stdout
					sys.stdout.write(chr(rc))

					HISTORY_RECORDER.record('CLI_RC', rc, 'CLI')
			except BrokenPipeError:
				pass
			sys.stdout = old_stdout


def main():
	if not os.access(rd, os.R_OK | os.W_OK):
		os.mkfifo(rd)
	if not os.access(wr, os.R_OK | os.W_OK):
		os.mkfifo(wr)

	if os.environ.get("AWSR_CLIENT") == "True":
		rc = run_client()
		sys.exit(rc)
	elif os.environ.get("AWSR_DAEMON") == "True":
		try:
			run_daemon()
		finally:
			os.remove(rd)
			os.remove(wr)
		return

	# fork if awsr daemon is not already running
	import psutil
	import subprocess
	ps = psutil.process_iter(attrs=["cmdline"])
	procs = 0
	for p in ps:
		cmd = p.info["cmdline"]
		if cmd and len(cmd) > 1 and cmd[0].endswith("python") and cmd[1] == sys.argv[0]:
			procs += 1
	if procs < 2:
		sys.stderr.write("Forking new awsr background process\n")
		with open(os.devnull, 'r+b', 0) as DEVNULL:
			# new instance will see env var, and run itself as daemon
			p = subprocess.Popen(sys.argv, stdin=DEVNULL, stdout=DEVNULL, stderr=DEVNULL, close_fds=True, env={
				"AWSR_DAEMON": "True",
				"AWS_PROFILE": os.environ.get("AWS_PROFILE", "default"),
			})
			run_client()

	else:
		run_client()

if __name__ == "__main__":
	main()