#!/usr/bin/env python

# grab the first unformatted disk and format it, then label it and update to fstab
# this must be ran as root

import subprocess
import time
import os
import datetime
import sys


def cmd(script, safe=False):
    p = subprocess.Popen(script, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = p.communicate()

    if not safe and p.returncode != 0:
        print "Fail to execute command: %s" % script
        print stdout, stderr
        exit(1)

    return stdout, stderr, p.returncode


def inspect_disk(disk):
    # clear all block cache
    cmd("blkid -g")

    out, err, _ = cmd("blkid /dev/%s -o export" % disk, safe=True)
    rows = [r.split("=") for r in out.strip().split("\n") if r.strip()]
    return {r[0]: r[1] for r in rows}


def format_disk(disk, disk_type="ext4", label=None):
    cmd("mkfs -t %s /dev/%s" % (disk_type, disk))
    if label:
        cmd("e2label /dev/%s \"%s\"" % (disk, label))


def update_fstab(diskinfo, mount_point):
    assert diskinfo['UUID']

    # create mount point
    cmd("mkdir -p %s" % mount_point)

    # check if disk is in fstab
    has_mount = False
    fstab = open("/etc/fstab").readlines()
    for line in fstab:
        if diskinfo['UUID'] in line:
            has_mount = True
            break

    if not has_mount:
        print("Updating fstab ...")
        fstab.append("")
        fstab.append("\t".join([
            "UUID=%s" % diskinfo['UUID'],
            mount_point,
            "ext4",
            "defaults,nofail",
            "0",
            "2"
        ]))
        fstab.append("")

        # now write to fstab
        with open("/etc/fstab", "w") as f:
            f.write("\n".join(fstab))
            f.flush()


def find_target_disk(target_label):
    # use lsblk to list all disks
    disks = cmd('lsblk -d -n -r -o NAME')[0].strip().split("\n")

    target_disk = None
    for disk in disks:
        print("Inspecting %s" % disk)

        # check if disk is formatted, instance store will be pre-formated
        disk_format = cmd("file -s /dev/%s" % disk)[0].split(": ")[1].strip()
        label = inspect_disk(disk).get('LABEL', '')

        # debug information
        print "format: " + disk_format
        print "label: " + label
        print ""
        
        # disk is not yet formatted, let's pick this one
        if disk_format == "data":
            target_disk = disk
        # disk is already formatted with the right label
        elif target_label and label == target_label:
            target_disk = disk
            break

    return target_disk


if __name__ == "__main__":
    label = "ARIMO-DATA"
    mount_point = "/opt/arimo-data"
    timeout = 180

    now = datetime.datetime.now()
    
    while True:
        target_disk = find_target_disk(label)

        if not target_disk:
            # check timeout
            if (datetime.datetime.now() - now).seconds > timeout:
                break

            print("Waiting for EBS to be ready ...")
            time.sleep(10)
        else:
            break

    if target_disk:
        print("Target disk: %s" % target_disk)

        info = inspect_disk(target_disk)
        
        if not info.get('TYPE'):
            print("Formating disk ...")
            format_disk(target_disk, label=label)

        # grab the new information
        info = inspect_disk(target_disk)
        update_fstab(info, mount_point)

        # re-mount everything
        print("Updating mount ...")
        cmd("mount -a")
    else:
        print("Can not find any disk")