#!/bin/bash
### BEGIN INIT INFO
# Provides:          ec2-vyos-init
# Required-Start:    vyos-router
# Required-Stop:
# Default-Start:     2 3 4 5
# Default-Stop:
# Short-Description: AWS EC2 instance init script to fetch and load ssh public key
# Description:       Retrieve user's public ssh key from EC2 instance metadata
#                    and load/set the key in config.boot
### END INIT INFO

# Author: hydrajump <wave@hydrajump.com>
#
# Based on http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/building-shared-amis.html#public-amis-install-credentials
#          https://github.com/andsens/bootstrap-vz/blob/master/providers/ec2/assets/init.d/ec2-get-credentials

. /lib/lsb/init-functions

# Are we running on AWS?
/opt/vyatta/sbin/ec2-check.pl
if [ $? != 0 ]; then
      exit 0
fi

# Hack for config permissions stuff
if [ $(groups | awk '{print $1}') != 'vyattacfg' ]; then
   sg vyattacfg $0
   exit
fi

: ${vyatta_env:=/etc/default/vyatta}
source $vyatta_env

# Configuration commands
SHELL_API=/bin/cli-shell-api
COMMIT=/opt/vyatta/sbin/my_commit
SAVE=/opt/vyatta/sbin/vyatta-save-config.pl
LOADKEY=/opt/vyatta/sbin/vyatta-load-user-key.pl
LOADCONFIG=/opt/vyatta/sbin/vyatta-load-config.pl

userdata_url=http://169.254.169.254/latest/user-data
public_key_url=http://169.254.169.254/latest/meta-data/public-keys/0/openssh-key
username='vyos'
ssh_dir="/home/$username/.ssh"
authorized_keys="$ssh_dir/authorized_keys"
group='vyattacfg'

# Obtain config session environment
session_env=$($SHELL_API getSessionEnv $PPID)
if [ $? -ne 0 ]; then
    echo "An error occured while obtaining session environment!"
    exit 0
fi

# Evaluate config environment string
eval $session_env

# Setup the config session
$SHELL_API setupSession
if [ $? -ne 0 ]; then
    echo "An error occured while setting up the configuration session!"
    exit 0
fi

load_user_data ()
{
    $LOADCONFIG $userdata_url
    $COMMIT
    $SAVE
}

load_ssh_public_key ()
{
    # Doesn't work.
    # if [ -x $vyatta_sbindir/vyatta-load-user-key.pl ]; then
    #     log_action_msg "Loaded ssh public key for user $username"
    #     sg ${group} -c "$vyatta_sbindir/vyatta-load-user-key.pl $username $public_key"
    # fi

    # Do this instead
    # Obtain session environment
    # Evaluate environment string
    # Setup the session
    # Commit and save config change
    # Tear down the session

    log_action_msg "EC2: Loaded ssh public key for user $username"
    $LOADKEY $username $public_key_url

    # Commit and save to config.boot
    $COMMIT
    $SAVE
}

# Try to load config from instance user-data
log_action_msg "EC2: -----BEGIN FETCH CONFIG-----"
log_action_msg "EC2: Requesting config from EC2 instance user-data"
if (curl --silent -f $userdata_url | grep 'vyatta-config-version' >/dev/null); then
    log_action_msg "EC2: Found Vyos config in EC2 instance user-data"
    load_user_data
else
    log_action_msg "EC2: No Vyos config found in EC2 instance user-data"
fi

log_action_msg "EC2: -----END FETCH CONFIG-----"
# Try to get the ssh public key from instance metadata
log_action_msg "EC2: -----BEGIN FETCH SSH PUBLIC KEY-----"
log_action_msg "EC2: Requesting ssh public key from EC2 instance metadata"
public_key=`/usr/bin/curl --silent -f $public_key_url`
if [ -n "$public_key" ]; then
    log_action_msg "EC2: Downloaded ssh public key from EC2 instance metadata"
    if [ ! -d $ssh_dir ]; then
        mkdir -m 700 $ssh_dir
        # chown $username:$username $ssh_dir
    fi

    # Check if the ssh public key is already loaded
    if ! grep -s -q "$public_key" $authorized_keys; then
        load_ssh_public_key
        # chmod 600 $authorized_keys
        # chown $username:$username $authorized_keys
    else
        log_action_msg "EC2: Already loaded ssh public key for user $username"
    fi
else
    log_action_msg "
    == WARNING ==
    No ssh public key found!
    If you launch an instance without specifying a keypair,
    you can't connect to the instance.
    Please terminate this instance and launch a new EC2 instance.

    == IMPORTANT ==
    Don't forget to create a keypair or select an existing one
    before you launch the new instance"
fi
log_action_msg "EC2: -----END FETCH SSH PUBLIC KEY-----"

# Tear down the config session
$SHELL_API teardownSession
if [ $? -ne 0 ]; then
    echo "An error occured while tearing down the session!"
    exit 0
fi
exit 0
