Skip to content

Instantly share code, notes, and snippets.

@Cyberax
Created December 19, 2019 20:14
Show Gist options
  • Save Cyberax/71030b46ebeb0ccda48e471664cd5540 to your computer and use it in GitHub Desktop.
Save Cyberax/71030b46ebeb0ccda48e471664cd5540 to your computer and use it in GitHub Desktop.
#!/usr/bin/python3
# Interweave two JSON log streams
import json
import os
import sys
import time
import urllib.request
from os import getenv
from subprocess import Popen, PIPE, STDOUT, check_output, check_call, call
from sys import stdout
from threading import Lock, Thread
AGENT = "/usr/bin/amazon-ssm-agent"
def weave(stream, output, lock):
while True:
# noinspection PyBroadException
try:
ln = stream.readline()
with lock:
output.write(ln)
except Exception:
break
def init_ssm():
# Get the optional task information
meta_url = getenv("ECS_CONTAINER_METADATA_URI", "")
creds_meta_url = getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", "")
if not meta_url or not creds_meta_url:
return False
contents = urllib.request.urlopen(meta_url + "/task").read()
meta = json.loads(contents.decode("utf-8"))
# E.g.: arn:aws:ecs:us-east-1:3412341234241:task\/e36711e4-0272-4f08-83e4-b80d36b9e180
taskId = meta['TaskARN'].split("/")[-1]
creds_contents = urllib.request.urlopen("http://169.254.170.2" + creds_meta_url).read()
creds_meta = json.loads(creds_contents.decode("utf-8"))
tags = [
"--tags",
"Key=Name,Value=%s" % taskId,
"Key=TaskId,Value=%s" % taskId,
"Key=Role,Value=Fargate",
"Key=Cluster,Value=%s" % meta['Cluster'],
"Key=TaskFamily,Value=%s" % meta['Family'],
"Key=TaskRevision,Value=%s" % meta['Revision'],
]
# arn:aws:iam::123412341234124:role/service-role/ApolloAlphaTask transformed
# to service-role/ApolloAlphaTask
ssm_role = "service-role/" + creds_meta["RoleArn"].split("/")[-1]
# E.g.: arn:aws:ecs:us-east-1:123412341234:task\/e36711e4-0272-4f08-83e4-b80d36b9e180
ssm_region = meta["TaskARN"].split(":")[3]
# arn:aws:ecs:us-east-1:1234123411313:cluster\/ApolloAlpha
cluster_name = meta["Cluster"].split("/")[-1]
ssm_desc = cluster_name + " Task"
act_out = check_output(["aws", "ssm", "create-activation", "--iam-role",
ssm_role, "--region", ssm_region, "--description", ssm_desc,
"--default-instance-name", "t-"+taskId] + tags)
act_data = json.loads(act_out.decode("utf-8"))
print('{"msg": "Received activation id %s"}' % (act_data["ActivationId"]))
# Register the instance in SSM
check_call([AGENT, "-register", "-y",
"-code", act_data["ActivationCode"], "-id", act_data["ActivationId"],
"-region", ssm_region, "-i", taskId])
print('{"msg": "Finished activation"}')
# Now we should have the machine-id, so we can link the task with it!
with open("/var/lib/amazon/ssm/registration") as fl:
reg_data = json.load(fl)
machine_id = reg_data["ManagedInstanceID"]
call(["aws", "ecs", "tag-resource", "--resource-arn", meta['TaskARN'],
"--tags", "key=MachineId,value=%s" % machine_id])
# Allow sudo access
with open("/etc/sudoers.d/ssm", "w") as fl:
fl.write('Defaults:ssm-user env_keep += *\n')
fl.write("ssm-user ALL=(ALL) NOPASSWD:ALL\n")
return True
def main():
mtx = Lock()
if not init_ssm():
# No SSM? Just launch the target directly
os.execlp(sys.argv[1], *sys.argv[1:])
ssm_process = Popen([AGENT], stdin=None, stdout=PIPE, stderr=STDOUT)
Thread(target=lambda: weave(ssm_process.stdout, stdout.buffer, mtx), name="SSMPump").start()
delegated = Popen(sys.argv[1:], stdin=sys.stdin, stdout=PIPE, stderr=sys.stderr)
Thread(target=lambda: weave(delegated.stdout, stdout.buffer, mtx), name="Delegated").start()
start = time.time()
delegated.wait()
# If the target has exited too fast (in less than 5 mins) then wait for
# 15 minutes before stopping the SSM to give us ability to log in and debug
# the issue (and to reduce the churn velocity).
if time.time() - start < 300:
time.sleep(900)
ssm_process.kill()
ssm_process.wait()
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment