#!/usr/bin/ruby

require 'rubygems'
require 'rdialog'
require 'sqlite3'
require 'fastercsv'
require 'pathname'

@DB_NAME = "cloud.db"

# The RDialog library left out the --form option, which is unacceptable,
# so I've decided to add this functionality myself. Perhaps later on
# I can fork it and fix it up in the actual library, but this is easier.
RDialog.class_eval do
    def form(text, items, height=0, width=0, fheight=0)
        temp_file = Tempfile.new('tmp')

        command = option_string() + "--form \ " + "\"#{text.to_s}\"" + "\ "
        command += height.to_i.to_s + " " + width.to_i.to_s + " " + fheight.to_i.to_s + "\ "

        items.each do |item|
            command += "\"#{item[0].to_s}\" #{item[1].to_i.to_s} #{item[2].to_i.to_s} "
            command += "\"#{item[3].to_s}\" #{item[4].to_i.to_s} #{item[5].to_i.to_s} "
            command += "#{item[6].to_i.to_s} #{item[7].to_i.to_s} \ "
        end

        command += "2> #{temp_file.path}"

        success = system(command)

        results = Array.new
        if success
            results = temp_file.read.split("\n")
            temp_file.close
            return results
        else
            temp_file.close
            return success
        end
    end
end

# This sets up a few dialog options we want to stick around
@dialog = RDialog.new
@dialog.backtitle = "VM Manager"
@dialog.shadow = false

# Our main menu! Everything goes through here.
def menu_win
    # Build up the menu items...
    items = Array.new
    items.push(["Create New VM",    ""])
    items.push(["Modify VM",        ""])
    items.push(["About",            ""])
    items.push(["Quit",             ""])

    # Display them, and figure out what to execute
    selected = @dialog.menu("VM Manager", items)

    case selected
        when "Create New VM"
            create_win
        when "Modify VM"
            modify_win
        when "About"
            about_win
        else
            exit 0
    end
end

# From this menu the user will be able to bring online a new VM
# with whatever parameters they desire! Huzzah!
def create_win

    # TODO: fix behavior associated with VMs from existing images

    # Ask for the host to bring the VM up on...
    # We've got a nifty list of hosts in a CSV file, so lets read it
    nodes = Hash.new
    FasterCSV.foreach('hosts.csv') do |row|
        nodes[row[0]] = row[1]
    end

    # Pressing CTRL+c exits this dialog but keeps the program going,
    # which is no good. This is just a simple defense against that.
    node = nil
    while(node == nil)
        # Now we prompt the user with our list of nodes
        sorted_nodes = nodes.to_a.sort { |a,b| a[1] <=> b[1] }
        node_result = @dialog.menu("Select a node to bring up VM on:", sorted_nodes, 30, 60)

        if(node_result == false)
            menu_win
        end

        # Keep in mind, node_result is just the hostname, not the IP:
        node = nodes[node_result].strip
    end

    # We have one last thing to prompt the user for:
    # whether or not they want to clone the HD or just
    # use the one passed in by the path.
    clone = nil
    while(clone == nil)
        prompt = "Is this VM going to be cloned from a template disk image, or should I not clone it?"
        clone_result = @dialog.menu(prompt, [["Yes, clone it", ""], ["No, don't clone it", ""]], 10, 60)
        case clone_result
            when "Yes, clone it"
                clone = true
            when "No, don't clone it"
                clone = false
            else
                menu_win
        end
    end

    # If we're cloning the HD, follow this path:
    if(clone)
        clone_win(node)
    else
        # We're not cloning the HD! Take this path.
        import_win(node)
    end

    menu_win
end

# This gets called when the user wants to import an OVA/OVF file
def import_win(node)
    items = Array.new
    items += [["Operating System" , 1, 1, "", 1, 25, 30, 100]]
    items += [["OVA/OVF Path"     , 2, 1, "", 2, 25, 30, 100]]
    items += [["Network Interface", 3, 1, "", 3, 25, 30, 100]]
    items += [["VM Number (0-99)" , 4, 1, "", 4, 25, 30, 100]]

    results = @dialog.form("Import a VM", items, 25, 60, 16)

    if(results == 0)
        menu_win
    end

# BEGIN ERROR CHECKING ==========================================================================
    # Check the validity of the VM number
    # We'll need to see which VM numbers are already taken first...
    db = SQLite3::Database.new(@DB_NAME)
    existing_nums = db.execute("select id from vms;").map { |i| i[0] }
    db.close

    if(results[3].match(/^\d{2}$/) == nil)
        @dialog.msgbox("ERROR\n\tInvalid VM number. Please Enter a number between 0 and 99.")
        create_win
    else if(existing_nums.include?(results[3]))
        @dialog.msgbox("ERROR\n\tA VM with #{results[3]} already exists. Please choose a different number.")
        create_win
    end

    # Make sure the OS type is a valid VBox OS type
    # First, grab a list of valid OS types from Virtual Box...
    os_types = `vboxmanage list ostypes`.scan(/^ID:\s*(\w*)$/).map { |i| i[0] }
    if(!os_types.include?(results[0]))
        @dialog.msgbox("ERROR\n\tInvalid OS type. Please check the list of valid VBox OS types with 'vboxmanage list ostypes' first.")
        create_win
    end

    # Make sure the path to the OVA/OVF file exists on the node
    disk_exists = system("ssh root@#{node} 'ls #{results[1]}'")
    if(!disk_exists)
        @dialog.msgbox("ERROR\n\tSpecified OVA/OVF file does not exist on #{node}")
        create_win
    end

    mac = "00C5CE3150#{results[3].to_s}"
    vm_name = Pathname.new(results[1]).basename.to_s

    command = "vboxmanage import #{results[1]};"
    command += "vboxmanage modifyvm --macaddress1 \"#{mac}\" --nic1 bridged --bridgeadapter1 \"#{results[2]}\";"
    command += "vboxheadless --startvm #{vm_name}"

    success = system("ssh root@#{node} '#{command}'")

    if(!success)
        @dialog.msgbox("ERROR\n\tCould not connect to host and/or modify/start the VM.")
        create_win
    else
        @dialog.msgbox("Success! #{vm_name} is now running on #{node}")
        menu_win
    end
end

# This gets called when we want to clone a VM from an existing disk
def clone_win(node)
    # Here we build up our giant form:
    items = Array.new
    items += [["Hostname",          1, 1, "",       1, 25, 30, 100]]
    items += [["Username",          2, 1, "",       2, 25, 30, 100]]
    items += [["Operating System",  3, 1, "Ubuntu", 3, 25, 30, 100]]
    items += [["RAM (MB)",          4, 1, "512",    4, 25, 30, 100]]
    items += [["VRAM (MB)",         5, 1, "32",     5, 25, 30, 100]]
    items += [["Network Interface", 6, 1, "br0",    6, 25, 30, 100]]
    items += [["Template VDI Path", 7, 1, "/home/cloud315/VMs/lamp_template.vdi", 7, 25, 30, 100]]
    items += [["New Disk Location", 8, 1, "/home/cloud315/VMs/", 8, 25, 30, 100]]
    items += [["Setup Script Path", 9, 1, "config_lamp.sh",       9, 25, 30, 100]]
    items += [["VM Number (0-99)", 10, 1, "",      10, 25, 30, 100]]

    # And now it gets sent off to my awesome form thing!
    results = @dialog.form("Set up a new VM", items, 25, 60, 16)

    # If they chose "cancel" then put 'em back at the beginning
    if(results == 0)
        menu_win
    end

# BEGIN ERROR CHECKING ==========================================================================
    # Check the validity of the VM number
    # We'll need to see which VM numbers are already taken first...
    db = SQLite3::Database.new(@DB_NAME)
    existing_nums = db.execute("select id from vms;").map { |i| i[0] }
    db.close

    if(results[9].match(/^\d{2}$/) == nil)
        @dialog.msgbox("ERROR\n\tInvalid VM number. Please Enter a number between 0 and 99.")
        create_win
    else if(existing_nums.include?(results[9]))
        @dialog.msgbox("ERROR\n\tA VM with #{results[9]} already exists. Please choose a different number.")
        create_win
    end

    # Check the hostname (can only contain a-z, 0-9, and dashes)
    if(results[0].match(/[^a-zA-Z\-0-9]+/) != nil)
            @dialog.msgbox("ERROR\n\tInvalid hostname. Hostnames can only contain characters a-z, digits 0-9, and dashes.")
            create_win
    end

    # Make sure the OS type is a valid VBox OS type
    # First, grab a list of valid OS types from Virtual Box...
    os_types = `vboxmanage list ostypes`.scan(/^ID:\s*(\w*)$/).map { |i| i[0] }
    if(!os_types.include?(results[2]))
        @dialog.msgbox("ERROR\n\tInvalid OS type. Please check the list of valid VBox OS types with 'vboxmanage list ostypes' first.")
        create_win
    end

    # Check the RAM and VRAM
    if(results[3].match(/[^0-9]/) != nil)
        @dialog.msgbox("ERROR\n\tInvalid RAM parameter. RAM must be specified as the number of MB of RAM.")
        create_win
    end

    if(results[4].match(/[^0-9]/) != nil)
        @dialog.msgbox("ERROR\n\tInvalid VRAM parameter. VRAM must be specified as the number of MB of VRAM.")
        create_win
    end

    # Make sure the path to the template disk image exists on the host
    template_exists = system("ssh root@#{node} 'ls #{results[6]}'")
    if(!template_exists)
        @dialog.msgbox("ERROR\n\tTemplate image does not exist on #{node}! Make sure your path is correct.")
        create_win
    end

    # Make sure the directory we're putting the cloned image in exists
    clone_dir_exists = system("ssh root@#{node} 'ls #{results[7]}'")
    if(!clone_dir_exists)
        @dialog.msgbox("ERROR\n\tDirectory specified for cloned disk image does not exist! Make sure the path you entered is correct.")
        create_win
    end

    # Verify that the setup script exists
    setup_script_exists = system("ls #{results[8]}")
    if(!setup_script_exists)
        @dialog.msgbox("ERROR\n\tSpecified setup script does not exist!")
        create_win
    end

# END ERROR CHECKING ==========================================================================

    mac = "00C5CE3150#{results[9].to_s}"

    # Now that we have the info we need, we build up a giant command:
    command = "ssh root@#{node} /home/cloud315/bin/vm_manage/create_vm.sh "
    command += "-h \"#{results[0]}\" -u \"#{results[1]}\" "
    command += "-o \"#{results[2]}\" -M \"#{results[3]}\" -V \"#{results[4]}\" "
    command += "-n \"#{results[5]}\" -v \"#{results[6]}\" -d \"#{results[7]}/\" "
    command += "-m \"#{mac}\""

    command += "; exit 0"

    # ...and execute said giant command
    puts node
    puts command
    result = system(command)

    # NOTE: Commented out until further notice, since command doesn't return true
    # If things didn't work out, let the user know
    # if(!result)
    #   @dialog.msgbox("ERROR\n\tUnable to reach host #{node}")
    #   create_win
    # end

    # Now we run the config script on the VM.
    # We HAVE to do it from the root node, because the template disk
    # images need to have the keys of the root node in the authorized
    # keys file of the root user so we can ssh in without a password.

    # scp over the config script
    password = `pwgen`.strip
    puts "PASSWORD GENERATED: #{password}"
    puts "scp #{results[8]} root@192.168.1.2#{results[9]}:/root/"
    success = system("scp #{results[8]} root@192.168.1.2#{results[9]}:/root/")

    # Inform the user if something went wrong
    if(!success)
        @dialog.msgbox("ERROR\n\tUnable to SCP config script to VM at 192.168.1.2#{results[9]}")
        create_win
    end

    # Execute the config script...
    ssh_cmd = "ssh root@192.168.1.2#{results[9]} /root/#{results[8]}"
    ssh_cmd += " -h \"#{results[0]}\" -u \"#{results[1]}\" -p \"#{password}\""
    puts ssh_cmd
    success = system(ssh_cmd)

    # Inform the user if something went wrong
    if(!success)
        @dialog.msgbox("ERROR\n\tUnable to execute config script on VM at 192.168.1.2#{results[9]}")
        create_win
    end

    @dialog.msgbox("Congrats! Your new VM is ready!
    Hostname: #{results[0]}\nUsername: #{results[1]}\nPassword: #{password}")
end

# You'll be able to change settings on currently running VMs
# from this window. Also a work in progress.
def modify_win
    # First, we need to build a list of VMs to display for the user...
    db = SQLite::Database.new(@DB_NAME)
    vms = db.execute("SELECT hostname, host_ip, status from vms;")
    vms.map! { |vm| "#{vm[0]} - on #{vm[1]} (#{status.gsub(/\(.*\)/, "").strip})" }
    menu_win
end

def about_win
    @dialog.msgbox("About:\n\tThe VM Manager was written by Eric Wood and Robert Schumacher for Dr. Liu's CSCE 315 class.")
    menu_win
end

menu_win