Cyberax · December 19, 2019 20:14 · Dec 19, 2019
diff --git a/weaver.py b/weaver.py
@@ -0,0 +1,115 @@
+#!/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()
No results found