pssh -I tee /tmp/settings.yaml < $SETTINGS

pssh sudo apt-get install -y python-setuptools

pssh sudo easy_install pyyaml

pssh -I tee /tmp/postprep.py <<EOF
#!/usr/bin/env python
import os
import platform
import sys
import time
import urllib
import yaml

#################################

config = yaml.load(open("/tmp/settings.yaml"))
COMPOSE_VERSION = config["compose_version"]
MACHINE_VERSION = config["machine_version"]
SWARM_VERSION = config["swarm_version"]
CLUSTER_SIZE = config["clustersize"]
ENGINE_VERSION = config["engine_version"]

#################################

# This script will be run as ubuntu user, which has root privileges.
# docker commands will require sudo because the ubuntu user has no access to the docker socket.

STEP = 0
START = time.time()

def bold(msg):
    return "{} {} {}".format("$(tput smso)", msg, "$(tput rmso)")

def system(cmd):
    global STEP
    with open("/tmp/pp.status", "a") as f:
        t1 = time.time()
        f.write(bold("--- RUNNING [step {}] ---> {}...".format(STEP, cmd)))
        retcode = os.system(cmd)
        if retcode:
            retcode = bold(retcode)
        t2 = time.time()
        td = str(t2-t1)[:5]
        f.write("[{}] in {}s\n".format(retcode, td))
        STEP += 1
    with open("/home/ubuntu/.bash_history", "a") as f:
        f.write("{}\n".format(cmd))

# On EC2, the ephemeral disk might be mounted on /mnt.
# If /mnt is a mountpoint, place Docker workspace on it.
system("if mountpoint -q /mnt; then sudo mkdir /mnt/docker && sudo ln -s /mnt/docker /var/lib/docker; fi")

# Put our public IP in /tmp/ipv4
# ipv4_retrieval_endpoint = "http://169.254.169.254/latest/meta-data/public-ipv4"
ipv4_retrieval_endpoint = "http://myip.enix.org/REMOTE_ADDR"
system("curl --silent {} > /tmp/ipv4".format(ipv4_retrieval_endpoint))

ipv4 = open("/tmp/ipv4").read()

# Add a "docker" user with password "training"
system("sudo useradd -d /home/docker -m -s /bin/bash docker")
system("echo docker:training | sudo chpasswd")

# Helper for Docker prompt.
system("""sudo tee /usr/local/bin/docker-prompt <<SQRL
#!/bin/sh
case "\\\$DOCKER_HOST" in
*:3376)
  echo swarm
  ;;
*:2376)
  echo \\\$DOCKER_MACHINE_NAME
  ;;
*:2375)
  echo \\\$DOCKER_MACHINE_NAME
  ;;
*:55555)
  echo \\\$DOCKER_MACHINE_NAME
  ;;
"")
  echo local
  ;;
*)
  echo unknown
  ;;
esac
SQRL""")
system("sudo chmod +x /usr/local/bin/docker-prompt")

# Fancy prompt courtesy of @soulshake.
system("""sudo -u docker tee -a /home/docker/.bashrc <<SQRL
export PS1='\e[1m\e[31m[\h] \e[32m(\\\$(docker-prompt)) \e[34m\u@{}\e[35m \w\e[0m\n$ '
SQRL""".format(ipv4))

# Custom .vimrc
system("""sudo -u docker tee /home/docker/.vimrc <<SQRL
syntax on
set autoindent
set expandtab
set number
set shiftwidth=2
set softtabstop=2
SQRL""")

# add docker user to sudoers and allow password authentication
system("""sudo tee /etc/sudoers.d/docker <<SQRL
docker ALL=(ALL) NOPASSWD:ALL
SQRL""")

system("sudo sed -i 's/PasswordAuthentication no/PasswordAuthentication yes/' /etc/ssh/sshd_config")

system("sudo service ssh restart")
system("sudo apt-get -q update")
system("sudo apt-get -qy install git jq python-pip")

# increase the size of the conntrack table so we don't blow it up when going crazy with http load testing
system("echo 1000000 | sudo tee /proc/sys/net/nf_conntrack_max")

#######################
### DOCKER INSTALLS ###
#######################

# This will install the latest Docker.
system("curl --silent https://{}/ | grep -v '( set -x; sleep 20 )' | sudo sh".format(ENGINE_VERSION))

### Install docker-compose
#system("sudo pip install -U docker-compose=={}".format(COMPOSE_VERSION))
system("sudo curl -sSL -o /usr/local/bin/docker-compose https://github.com/docker/compose/releases/download/{}/docker-compose-{}-{}".format(COMPOSE_VERSION, platform.system(), platform.machine()))
system("sudo chmod +x /usr/local/bin/docker-compose")

### Install docker-machine
system("sudo curl -sSL -o /usr/local/bin/docker-machine https://github.com/docker/machine/releases/download/v{}/docker-machine-{}-{}".format(MACHINE_VERSION, platform.system(), platform.machine()))
system("sudo chmod +x /usr/local/bin/docker-machine")

system("sudo apt-get remove -y --purge dnsmasq-base")
system("sudo apt-get -qy install python-setuptools pssh apache2-utils httping htop unzip mosh")

### Wait for Docker to be up.
### (If we don't do this, Docker will not be responsive during the next step.)
system("while ! sudo -u docker docker version ; do sleep 2; done")

### Install Swarm
system("docker pull swarm:{}".format(SWARM_VERSION))
system("docker tag -f swarm:{} swarm".format(SWARM_VERSION))

### BEGIN CLUSTERING ###

addresses = list(l.strip() for l in sys.stdin)

assert ipv4 in addresses

def makenames(addrs):
    return [ "node%s"%(i+1) for i in range(len(addrs)) ]

while addresses:
    cluster = addresses[:CLUSTER_SIZE]
    addresses = addresses[CLUSTER_SIZE:]
    if ipv4 not in cluster:
        continue
    names = makenames(cluster)
    for ipaddr, name in zip(cluster, names):
        system("grep ^{} /etc/hosts || echo {} {} | sudo tee -a /etc/hosts"
                    .format(ipaddr, ipaddr, name))
    print(cluster)

    mynode = cluster.index(ipv4) + 1
    system("echo 'node{}' | sudo -u docker tee /tmp/node".format(mynode))
    system("sudo -u docker mkdir -p /home/docker/.ssh")
    system("sudo -u docker touch /home/docker/.ssh/authorized_keys")

    if ipv4 == cluster[0]:
        # If I'm node1 and don't have a private key, generate one (with empty passphrase)
        system("sudo -u docker [ -f /home/docker/.ssh/id_rsa ] || sudo -u docker ssh-keygen -t rsa -f /home/docker/.ssh/id_rsa -P ''")

FINISH = time.time()
duration = "Initial deployment took {}s".format(str(FINISH - START)[:5])
system("echo {}".format(duration))

EOF

IPS_FILE=ips.txt
if [ ! -s $IPS_FILE ]; then
    echo "ips.txt not found."
    exit 1
fi

pssh --timeout 900 --send-input "python /tmp/postprep.py >>/tmp/pp.out 2>>/tmp/pp.err" < $IPS_FILE

# If /home/docker/.ssh/id_rsa doesn't exist, copy it from node1
pssh "sudo -u docker [ -f /home/docker/.ssh/id_rsa ] || ssh -o StrictHostKeyChecking=no node1 sudo -u docker tar -C /home/docker -cvf- .ssh | sudo -u docker tar -C /home/docker -xf-"

# if 'docker@' doesn't appear in /home/docker/.ssh/authorized_keys, copy it there
pssh "grep docker@ /home/docker/.ssh/authorized_keys \
    || cat /home/docker/.ssh/id_rsa.pub \
        | sudo -u docker tee -a /home/docker/.ssh/authorized_keys"

# On node1, create and deploy TLS certs using Docker Machine
pssh "if grep -q node1 /tmp/node; then grep ' node' /etc/hosts | xargs -n2 sudo -H -u docker docker-machine create -d generic --generic-ssh-user docker --generic-ip-address; fi"

