#!/usr/bin/env bash

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

username=root
userDataServerPort=8080
configDriveLabel=config-2

function findPrimaryNetwork(){
    outputLog "Detecting primary network"
    if command -v ip &> /dev/null
    then
        primaryNet=$(ip -o -4 route show to default | awk '{print $5}')
    elif command -v netstat &> /dev/null
    then
        primaryNet=$(netstat -r4 | grep default | awk '{print $(NF)}')
    elif command -v route &> /dev/null
    then
        primaryNet=$(route -4 2> /dev/null | grep default | awk '{print $(NF)}')
        if [ -z "$primaryNet" ]
        then
            primaryNet=$(route get default 2> /dev/null | grep interface | tr -d ' ' | awk '{split($0,a,":"); print a[2]}')
        fi
    fi
    if [ -z "$primaryNet" ]
    then
        outputLog "Could not find primary network"
        return 1
    fi
    echo "$primaryNet"
    return 0
}

function findUserDataServer(){
    primaryNet=$1
    outputLog "Trying to find userdata server"
    if [ -z "$primaryNet" ]
    then
        outputLog "Unable to determine the userdata server, falling back to data-server"
        echo "data-server"
        return 0
    fi

    if command -v netplan &> /dev/null
    then
        outputLog "Operating System is using netplan"

        userDataServer=$(netplan ip leases "$primaryNet" | grep SERVER_ADDRESS | awk '{split($0,a,"="); print a[2]}')

        if [ -n "$userDataServer" ]
        then
            outputLog "Found userdata server IP $userDataServer in netplan config"
            echo "$userDataServer"
            return 0
        fi
    fi

    if command -v nmcli &> /dev/null
    then
        outputLog "Operating System is using NetworkManager"

        userDataServer=$(nmcli -t connection show "$(nmcli -t -f UUID,DEVICE connection | grep "$primaryNet" | awk '{split($0,a,":"); print a[1]}')" | grep next_server | tr -d ' ' |awk '{split($0,a,"="); print a[2]}')

        if [ -n "$userDataServer" ]
        then
            outputLog "Found userdata server IP $userDataServer in NetworkManager config"
            echo "$userDataServer"
            return 0
        fi
    fi

    if command -v wicked &> /dev/null
    then
        outputLog "Operating System is using wicked"

        userDataServer=$(grep SERVERID /run/wicked/leaseinfo."$primaryNet"* | tr -d "'" | awk '{split($0,a,"="); print a[2]}')

        if [ -n "$userDataServer" ]
        then
            outputLog "Found userdata server IP $userDataServer in wicked config"
            echo "$userDataServer"
            return 0
        fi
    fi

    if command -v udhcpc &> /dev/null
    then
        outputLog "Operating System is using udhcpc"

        userDataServer=$(< /run/dhcp-server-ip."$primaryNet")

        if [ -n "$userDataServer" ]
        then
            outputLog "Found userdata server IP $userDataServer in udhcpc"
            echo "$userDataServer"
            return 0
        fi
    fi

    outputLog "Searching for DHCP server in lease files"

    primaryLease=$(
        dhcpFolders="/var/lib/dhclient/* /var/lib/dhcp3/* /var/lib/dhcp/* /var/lib/NetworkManager/* /var/db/dhclient*"
        for files in $dhcpFolders
        do
            if [ -e "$files" ]
            then
                < "$files" tr -d '\n' | sed 's/  //g ; s/lease {//g ; s/}/\n/g' | grep 'option routers'
            fi
        done
    )

    serverList=$(
        IFS=$'\n'
        for line in $(echo -e "$primaryLease")
        do
            splitLine=$(echo "$line" | sed -e 's/;/\n/g')
            if date -j &> /dev/null
            then
                timestamp=$(date -j -f "%Y/%m/%d %H:%M:%S" "$(echo "$splitLine" | grep 'expire' | sed -r 's/.*expire [0-9]+ (.*)/\1/')" +"%s")
            else
                timestamp=$(date -d "$(echo "$splitLine" | grep 'expire' | sed -e 's/.*expire [0-9]\+ \(.*\)/\1/')" +"%s")
            fi
            interface=$(echo "$splitLine" | grep 'interface' | sed -e 's/.*interface "\(.*\)"/\1/')
            server=$(echo "$splitLine" | grep 'dhcp-server-identifier' | sed -e 's/.*dhcp-server-identifier \(.*\)/\1/')
            echo "$timestamp","$interface","$server"
        done
    )

    userDataServer=$(echo "$serverList" | grep "$primaryNet" | sort -n | tail -1 | awk '{split($0,a,","); print a[3]}')

    if [ -n "$userDataServer" ]
    then
        outputLog "Userdata server found: $userDataServer"
        echo "$userDataServer"
        return 0
    fi

    outputLog "Unable to determine the userdata server, falling back to data-server"
    echo "data-server"
    return 0
}

function getPasswordFromUserDataServer(){
    userDataServer=$1
    userDataServerPort=$2
    outputLog "Sending request to userdata server at $userDataServer to get the password"
    if ! response=$(curl --fail --silent --connect-timeout 20 --retry 3 --header "DomU_Request: send_my_password" http://"$userDataServer":"$userDataServerPort")
    then
        outputLog "Failed to send request to userdata server at $userDataServer"
        return 4
    fi
    outputLog "Got response from userdata server at $userDataServer"
    response=$(echo "$response" | tr -d '\r')
    case $response in
        "")
            outputLog "Userdata server at $userDataServer did not have any password for the VM"
            return 2
        ;;
        "bad_request")
            outputLog "VM sent an invalid request to userdata server at $userDataServer"
            return 3
        ;;
        "saved_password")
            outputLog "VM has already saved a password from the userdata server at $userDataServer"
            return 1
        ;;
        *)
            outputLog "VM got a valid password from server at $userDataServer"
            echo "$response"
            return 0
    esac
}

function findHomeDirectory(){
    username=$1
    getent passwd "$username"|awk -F ":" '{print $6}'
}

function setPassword(){
    username=$1
    homeDir=$2
    password=$3
    if command -v md5sum &> /dev/null
    then
        newMd5=$(echo "$password" | md5sum | awk '{print $1}')
    elif command -v md5 &> /dev/null
    then
        newMd5=$(echo "$password" | md5)
    else
        newMd5='N/A'
    fi
    if [ $newMd5 != 'N/A' ]
    then
        if [ -f "$homeDir"/.password.md5 ]
        then
            oldMd5=$(cat "$homeDir"/.password.md5)
        fi
        if [ "$newMd5" ==  "$oldMd5" ]
        then
            outputLog  "There is no update of VM password"
            return 0
        fi
    else
        outputLog "Cannot determine change of password"
    fi
    outputLog "Changing password for user $username"
    if command -v chpasswd &> /dev/null
    then
        echo "$username":"$password" | chpasswd
    elif command -v usermod &> /dev/null && command -v mkpasswd &> /dev/null
    then
        usermod -p "$(mkpasswd -m SHA-512 "$password")" "$username"
    elif command -v pw &> /dev/null
    then
        echo "$password" | pw mod user "$username" -h 0
    else
        outputLog "Failed to change password for user $username"
        return 1
    fi
    outputLog "Successfully changed password for user $username"
    if [ $newMd5 != 'N/A' ]
    then
        echo "$newMd5" > "$homeDir"/.password.md5
        chmod 600 "$homeDir"/.password.md5
        chown "$username": "$homeDir"/.password.md5
    fi
    return 0
}

function sendAckToUserDataServer(){
    userDataServer=$1
    userDataServerPort=$2
    outputLog "Sending acknowledgment to userdata server at $userDataServer"
    if ! curl --fail --silent --connect-timeout 20 --retry 3 --header "DomU_Request: saved_password" "$userDataServer":"$userDataServerPort" &> /dev/null
    then
        outputLog "Failed to sent acknowledgment to userdata server at $userDataServer"
        return 1
    fi
    outputLog "Successfully sent acknowledgment to userdata server at $userDataServer"
    return 0
}

function getPublicKeyFromUserDataServer(){
    userDataServer=$1
    outputLog "Sending request to userdata server at $userDataServer to get public key"
    if ! response=$(curl --fail --silent --connect-timeout 20 --retry 3 http://"$userDataServer"/latest/public-keys)
    then
        outputLog "Failed to get public key from userdata server"
        return 2
    fi
    outputLog "Got response from userdata server at $userDataServer"
    if [ -z "$response" ]
    then
        outputLog "Did not receive any public keys from userdata server"
        return 1
    fi
    outputLog "Successfully get public key from userdata server"
    echo "$response"
    return 0
}

function setPublicKey(){
    username=$1
    homeDir=$2
    publicKey=$3
    outputLog "Applying public key for $username"
    sshDir=$homeDir/.ssh
    authorizedKeysFile=$sshDir/authorized_keys

    if [ ! -d "$sshDir" ]
    then
        outputLog ".ssh directory for $username not found, creating .ssh directory"
        mkdir "$sshDir"
    fi

    if [ ! -f "$authorizedKeysFile" ]
    then
        outputLog "authorized_keys file for $username not found, creating authorized_keys file"
        touch "$authorizedKeysFile"
    fi
    if grep "$(echo "$publicKey" | awk '{print $2}')" "$authorizedKeysFile" > /dev/null
    then
        outputLog "No need to update authorized_keys file"
        return 0
    fi
    outputLog "Writing public key in authorized_keys file"
    sed -i "/ cloudstack@apache.org$/d" "$authorizedKeysFile"
    echo "$publicKey cloudstack@apache.org" >> "$authorizedKeysFile"
    chmod 600 "$authorizedKeysFile"
    chmod 700 "$sshDir"
    chown -R "$username": "$sshDir"
    which restorecon &> /dev/null && restorecon -R -v "$sshDir"
    return 0
}

function findConfigDrive(){
    configDriveLabel=$1
    outputLog "Searching for ConfigDrive"

    if [ -e /dev/disk/by-label/"$configDriveLabel" ]
    then
        outputLog "ConfigDrive found at /dev/disk/by-label/$configDriveLabel"
        echo "/dev/disk/by-label/$configDriveLabel"
        return 0
    fi

    if [ -e /dev/iso9660/"$configDriveLabel" ]
    then
        outputLog "ConfigDrive found at /dev/iso9660/$configDriveLabel"
        echo "/dev/iso9660/$configDriveLabel"
        return 0
    fi

    blockDevice=$(blkid -t LABEL="$configDriveLabel" /dev/hd? /dev/sd? /dev/xvd? /dev/vd? /dev/sr? -o device 2> /dev/null)
    if [ -n "$blockDevice" ]
    then
        outputLog "ConfigDrive found at $blockDevice"
        echo "$blockDevice"
        return 0
    fi
    outputLog "ConfigDrive not found"
    return 1
}

function mountConfigDrive(){
    disk=$1
    outputLog "Mounting ConfigDrive"
    mountDir=$(mktemp -d)
    if [ ! -e "$mountDir" ]
    then
        mkdir "$mountDir"
        chmod 700 "$mountDir"
    fi

    mounted=0
    if [ $mounted == 0 ] && mount -r "$disk" "$mountDir" &> /dev/null
    then
        mounted=1
    fi
    if [ $mounted == 0 ] && mount -r -t cd9660 "$disk" "$mountDir" &> /dev/null
    then
        mounted=1
    fi
    if [ $mounted == 0 ] && mount -r -t iso9660 "$disk" "$mountDir" &> /dev/null
    then
        mounted=1
    fi

    if [ $mounted == 1 ]
    then
        outputLog "$disk successfully mounted on $mountDir"
        echo "$mountDir"
        return 0
    fi

    outputLog "Failed mounting $disk on $mountDir"
    rm -rf "$mountDir"
    return 1
}

function unmountConfigDrive(){
    mountDir=$1
    outputLog "Unmounting ConfigDrive"
    if ! umount "$mountDir"
    then
        outputLog "Failed unmounting $mountDir"
        return 1
    fi
    rm -rf "$mountDir"
    outputLog "Successfully unmount $mountDir"
    return 0
}

function getPasswordFromConfigDrive(){
    mountDir=$1
    passwordFile=$mountDir/cloudstack/password/vm_password.txt
    if [ ! -f "$passwordFile" ]
    then
        outputLog "Password file not found in ConfigDrivee"
        return 3
    fi
    outputLog "Password file found in ConfigDrive"
    content=$(< "$passwordFile" tr -d '\r')

    case $content in

        "")
            outputLog "ConfigDrive did not have any password for the VM"
            return 2
        ;;

        "saved_password")
            outputLog "VM has already saved a password"
            return 1
        ;;

        *)
            outputLog "VM got a valid password"
            echo "$content"
            return 0
    esac
}

function getPublicKeyFromConfigDrive() {
    mountDir=$1
    publicKeyFile=$mountDir/cloudstack/metadata/public-keys.txt

    if [ ! -f "$publicKeyFile" ]
    then
        outputLog "Public key file not found in ConfigDrive"
        return 2
    fi
    content=$(< "$publicKeyFile" tr -d '\r')

    if [ -z "$content" ]
    then
        outputLog "Did not receive any public keys"
        return 1
    fi
    echo "$content"
    outputLog "Public key successfully received."
    return 0
}

function outputLog() {
    stderr=1
    logger=1
    message=$1
    if [ $stderr == 1 ]
    then
        echo "Cloud Password Manager: $message" 1>&2
    fi
    if [ $logger == 1 ]
    then
        logger -t "Cloud Password Manager" "$message"
    fi
}

publicKeyReceived=0
passwordReceived=0
dataSource=''

if disk=$(findConfigDrive "$configDriveLabel")
then
    if mountDir=$(mountConfigDrive "$disk")
    then
        dataSource='ConfigDrive'
        if publicKey=$(getPublicKeyFromConfigDrive "$mountDir")
        then
            publicKeyReceived=1
        fi
        if password=$(getPasswordFromConfigDrive "$mountDir")
        then
            passwordReceived=1
        fi
        unmountConfigDrive "$mountDir"
    fi
fi
if [ $publicKeyReceived == 0 ] || [ $passwordReceived == 0 ]
then
    primaryNet=$(findPrimaryNetwork)
    userDataServer=$(findUserDataServer "$primaryNet")
    if [ $publicKeyReceived == 0 ]
    then
        if publicKey=$(getPublicKeyFromUserDataServer "$userDataServer")
        then
            dataSource='UserDataServer'
            publicKeyReceived=1
        fi
    fi
    if [ $passwordReceived == 0 ]
    then
        if password=$(getPasswordFromUserDataServer "$userDataServer" "$userDataServerPort")
        then
            dataSource='UserDataServer'
            passwordReceived=1
        fi
    fi
fi
homeDir=$(findHomeDirectory "$username")
if [ $passwordReceived == 1 ]
then
    setPassword "$username" "$homeDir" "$password"
    if [ $dataSource == 'UserDataServer' ]
    then
        sendAckToUserDataServer "$userDataServer" "$userDataServerPort"
    fi
fi
if [ $publicKeyReceived == 1 ]
then
    setPublicKey "$username" "$homeDir" "$publicKey"
fi