#!/usr/bin/env python3

"""
Helper utility that lists and mounts AWS EFS filesystems.
"""

import os, sys, subprocess, argparse, logging, collections

import boto3, requests

logging.basicConfig(level=logging.ERROR)
logger = logging.getLogger(__name__)
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--timeout", type=int, default=30)
args = parser.parse_args()

def get_metadata(path):
    return requests.get("http://169.254.169.254/latest/meta-data/{}".format(path)).content.decode()

az = get_metadata("placement/availability-zone")
services_domain = get_metadata("services/domain")
#subnet = get_metadata("network/interfaces/macs/{}/subnet-id".format(get_metadata("mac")))

efs = boto3.Session(region_name=az[:-1]).client("efs")

mountpoints = {}

for filesystem in efs.describe_file_systems()["FileSystems"]:
    for tag in efs.describe_tags(FileSystemId=filesystem["FileSystemId"])["Tags"]:
        if tag["Key"] == "mountpoint":
            mountpoints[tag["Value"]] = filesystem
else:
    logger.info("No EFS filesystems found")

mount_procs = collections.OrderedDict()
for mountpoint, filesystem in mountpoints.items():
    logging.info("Mounting %s on %s", filesystem["FileSystemId"], mountpoint)
    fs_url = "{az}.{fs}.efs.{region}.{domain}:/".format(az=az,
                                                        fs=filesystem["FileSystemId"],
                                                        region=efs.meta.region_name,
                                                        domain=services_domain)
    mount_procs[mountpoint] = subprocess.Popen(["mount", "-t", "nfs4", fs_url, mountpoint])

for m, p in mount_procs.items():
    if not os.path.isdir(m):
        os.makedirs(m)
    try:
        p.wait(timeout=args.timeout)
    except Exception:
        p.terminate()
        logger.error("Error while mounting EFS filesystem")

exit(sum([getattr(p, "exitcode", 1) for p in mount_procs.values()]))
