#!/usr/bin/python 
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = '''
---
module: vagrant
short_description: create a local instance via vagrant
description:
  - creates VM instances via vagrant and optionally waits for it to be 'running'.
version_added: "1.1"
options:
  state:
    description:
      - Should the VMs be "present" or "absent."
  cmd:
    description:
      - vagrant subcommand to execute. Can be "up," "status," "config," "ssh," "halt," "destroy" or "clear." 
    required: false
    default: null
    aliases: ['command'] 
  box_name:
    description:
      - vagrant boxed image to start
    required: false
    default: null
    aliases: ['image']
  box_path:
    description:
      - path to vagrant boxed image to start
    required: false
    default: null
    aliases: []
  vm_name:
    description:
      - name to give an associated VM
    required: false
    default: null
    aliases: [] 
  count:
    description:
      - number of instances to launch
    required: False
    default: 1
    aliases: []
  forward_ports:
    description:
      - comma separated list of ports to forward to the host. If the port is under 1024, the host port will be the guest port +10000
    required: False
    aliases: []
  memory:
    description:
      - memory in MB
    required: False

examples:
   - code: 'local_action: vagrant cmd=up box_name=lucid32 vm_name=webserver'
     description:
requirements: [ "vagrant", "python-vagrant" ]
author: Rob Parrott
'''

VAGRANT_FILE = "./Vagrantfile"
VAGRANT_DICT_FILE = "./Vagrantfile.json"
VAGRANT_LOCKFILE = "./.vagrant-lock"

VAGRANT_FILE_HEAD = "Vagrant::Config.run do |config|\n"
VAGRANT_FILE_BOX_NAME = "  config.vm.box = \"%s\"\n"
VAGRANT_FILE_VM_STANZA_HEAD = """
  config.vm.define :%s do |%s_config|
    %s_config.vm.network :hostonly, "%s"
    %s_config.vm.box = "%s"
"""
VAGRANT_FILE_HOSTNAME_LINE     = "    %s_config.vm.host_name = \"%s\"\n"
VAGRANT_FILE_PORT_FORWARD_LINE = "    %s_config.vm.forward_port %s, %s\n"  
VAGRANT_FILE_MEMORY_LINE       = "    %s_config.vm.customize [\"modifyvm\", :id, \"--memory\", %s]\n" 
VAGRANT_FILE_VM_STANZA_TAIL="  end\n"

VAGRANT_FILE_TAIL = "\nend\n"

# If this is already a network on your machine, this may fail ... change it here.
VAGRANT_INT_IP = "192.168.179.%s"

DEFAULT_VM_NAME = "ansible"
DEFAULT_VM_RAM = 1024

import sys
import subprocess
#import time
import os.path
import json

try:
    import lockfile
except ImportError:
    print "Python module lockfile is not installed. Falling back to using flock(), which will fail on Windows."
    
try:
    import vagrant
except ImportError:
    print "failed=True msg='python-vagrant required for this module'"
    sys.exit(1)

class VagrantWrapper(object):

    def __init__(self):
        '''
        Wrapper around the python-vagrant module for use with ansible.
        Note that Vagrant itself is non-thread safe, as is the python-vagrant lib, so we need to lock on basically all operations ...
        '''
        # Get a lock
        self.lock = None

        try:
            self.lock = lockfile.FileLock(VAGRANT_LOCKFILE)
            self.lock.acquire()
        except:
            # fall back to using flock instead ...
            try:
                import fcntl
                self.lock = open(VAGRANT_LOCKFILE, 'w')
                fcntl.flock(self.lock, fcntl.LOCK_EX)
            except:
                print "failed=True msg='Could not get a lock for using vagrant. Install python module \"lockfile\" to use vagrant on non-POSIX filesytems.'"
                sys.exit(1)
            
        # Initialize vagrant and state files
        self.vg = vagrant.Vagrant()
        
        # operation will create a default data structure if none present
        self._deserialize()  
        self._serialize()
        
    def __del__(self):
        '''Clean up file locks'''
        try:
            self.lock.release()
        except:
            os.close(self.lock)
            os.unlink(self.lock)
        
    def prepare_box(self, box_name, box_path):
        '''Given a specified name and URL, import a Vagrant "box" for use.'''
        changed = False
        if box_name == None:
            raise Exception("You must specify a box_name with a box_path for vagrant.")
        boxes = self.vg.box_list()
        if not box_name in boxes:
            self.vg.box_add(box_name, box_path)
            changed = True
            
        return changed
             
    def up(self, box_name, vm_name=None, count=1, box_path=None, ports=[]):    
        '''Fire up a given VM and name it, using vagrant's multi-VM mode.'''

        changed = False
        if vm_name == None: 
            vm_name = DEFAULT_VM_NAME
        
        if box_name == None:
            raise Exception("You must specify a box name for Vagrant.")
        if box_path != None: 
            changed = self.prepare_box(box_name, box_path)
        
        for icount in range(int(count)):
            
            self._deserialize()
            
            this_instance_dict = self._get_instance(vm_name,icount)
            if not this_instance_dict.has_key('box_name'): 
                this_instance_dict['box_name'] = box_name   
                     
            this_instance_dict['forward_ports'] = ports  
            
            # Save our changes and run
            inst_array = self._instances()[vm_name]
            inst_array[icount] = this_instance_dict    
            
            self._serialize()
            
            # See if we need to fire it up ...
            vgn = this_instance_dict['vagrant_name']
            status = self.vg.status(vgn)
            if status != 'running':
                self.vg.up(False, this_instance_dict['vagrant_name'])
                changed =True
                
        ansible_instance_array = self._build_instance_array_for_ansible(vm_name)
        return (changed, ansible_instance_array)
    
    def status(self, vm_name = None, index = -1):
        '''Return the run status of the VM instance. If no instance N is given, returns first instance.'''
        vm_names = []
        if vm_name != None: vm_names = [vm_name]
        else:
            vm_names = self._instances().keys()
        
        statuses = {}
        for vmn in vm_names:
            stat_array = []
            instance_array = self.vg_data['instances'][vmn]
            if index >= 0:
                instance_array = [ self._get_instance(vmn,index) ]
            for inst in instance_array:
                vgn = inst['vagrant_name']
                stat_array.append(self.vg.status(vgn))
            statuses[vmn] = stat_array     
                    
        return (False, statuses)
            
    def config(self, vm_name, index = -1):
        '''Return info on SSH for the running instance.'''
        vm_names = []
        if vm_name != None: vm_names = [vm_name]
        else:
            vm_names = self._instances().keys()
        
        configs = {}
        for vmn in vm_names:
            conf_array = []
            instance_array = self.vg_data['instances'][vmn]
            if index >= 0:
                instance_array = [ self._get_instance(vmn,index) ]
            for inst in instance_array:
                cnf = self.vg.conf(None, inst['vagrant_name'])
                conf_array.append(cnf)
            configs[vmn] = conf_array     
                            
        return (False, configs)

    def halt(self, vm_name = None, index = -1):
        '''Shuts down a vm_name or all VMs.'''
        
        changed = False
        vm_names = []
        if vm_name != None: vm_names = [vm_name]
        else:
            vm_names = self._instances().keys()
        
        statuses = {}
        for vmn in vm_names:
            stat_array = []
            instance_array = self.vg_data['instances'][vmn]
            if index >= 0:
                instance_array = [ self.vg_data['instances'][vmn][index] ]
            for inst in instance_array:
                vgn = inst['vagrant_name']
                if self.vg.status(vgn) == 'running':
                    self.vg.halt(vgn)
                    changed = True
                stat_array.append(self.vg.status(vgn))
            statuses[vmn] = stat_array
            
        return (changed, statuses)         
            
    def destroy(self, vm_name=None, index = -1):
        '''Halt and remove data for a VM, or all VMs.'''
        
        self._deserialize()
        
        (changed, stats) = self.halt(vm_name, index)       
    
        self.vg.destroy(vm_name)
        if vm_name != None:
            self._instances().pop(vm_name)
        else:
            self.vg_data['instances'] = {}
            
        self._serialize()

        changed = True
        
        return changed
        
    def clear(self, vm_name=None):
        '''Halt and remove data for a VM, or all VMs. Also clear all state data.'''
        
        changed = self.vg.destroy(vm_name)
        
        if os.path.isfile(VAGRANT_FILE):
            os.remove(VAGRANT_FILE)
        if os.path.isfile(VAGRANT_DICT_FILE):
            os.remove(VAGRANT_DICT_FILE)            

        return changed           
#
# Helper Methods
#
    def _instances(self): 
        return self.vg_data['instances']
    
    def _get_instance(self, vm_name, index):
    
        instances = self._instances()

        inst_array = []        
        if instances.has_key(vm_name):
            inst_array = instances[vm_name]
            
        if len(inst_array) > index: 
            return inst_array[index]
        
        # 
        # otherwise create one afresh
        #
        this_instance_N = self.vg_data['num_inst']+1
        name_for_vagrant = "%s%d" % (vm_name.replace("-","_"),index)
        
        instance_dict = dict(
          n = index,
          N = this_instance_N,
          name = vm_name,
          vagrant_name = name_for_vagrant,
          internal_ip = VAGRANT_INT_IP % (255-this_instance_N),
          forward_ports = [],
          ram = DEFAULT_VM_RAM,
        )

        # Save this ...
        self.vg_data['num_inst'] = this_instance_N
        inst_array.append(instance_dict)
        
        self._instances()[vm_name] = inst_array
        
        return instance_dict          
    
    #
    # Serialize/Deserialize current state to a JSON representation, and 
    #  a file format for Vagrant.
    #
    # This is where we need to deal with file locking, since multiple threads/procs
    #  may be trying to operate on the same files 
    # 
    
    def _serialize(self):
        '''Save state to a JSON file, and write the Vagrantfile based on this.'''
        self._save_state()
        self._write_vagrantfile()
        
    def _deserialize(self): 
        '''Load in data from the JSON state file.'''
        self._load_state()
        
    
    #
    # Manage a JSON representation of vagrantfile for statefulness across invocations. 
    #
    def _load_state(self):        
        self.vg_data = dict(num_inst=0, instances = {})
        if os.path.isfile(VAGRANT_DICT_FILE):
            json_file=open(VAGRANT_DICT_FILE)
            self.vg_data = json.load(json_file)
            json_file.close()

#    def _state_as_string(self):
#        from StringIO import StringIO
#        io = StringIO()
#        json.dump(self.vg_data, io)
#        return io.getvalue()

    def _save_state(self):
        json_file=open(VAGRANT_DICT_FILE, 'w')
        json.dump(self.vg_data,json_file, sort_keys=True, indent=4, separators=(',', ': '))
        json_file.close()
      
    #
    # Translate the state dictionary into the Vagrantfile
    #    
    def _write_vagrantfile(self):

        vfile = open(VAGRANT_FILE, 'w')
        vfile.write(VAGRANT_FILE_HEAD)

        instances = self._instances()
        for vm_name in instances.keys():
            inst_array = instances[vm_name]
            for index in range(len(inst_array)):
                instance_dict = inst_array[index] 
                name     = instance_dict['vagrant_name']
                ip       = instance_dict['internal_ip']
                box_name = instance_dict['box_name']
                vfile.write(VAGRANT_FILE_VM_STANZA_HEAD % 
                            (name, name, name, ip, name, box_name) )
                if instance_dict.has_key('ram'):
                    vfile.write(VAGRANT_FILE_MEMORY_LINE  % (name, instance_dict['ram'])  )   
                vfile.write(VAGRANT_FILE_HOSTNAME_LINE  % (name, name.replace('_','-'))  )       
                if instance_dict.has_key('forward_ports'):
                    for port in instance_dict['forward_ports']:
                        port = int(port)
                        host_port = port
                        if port < 1024: 
                            host_port = port + 10000
                        vfile.write(VAGRANT_FILE_PORT_FORWARD_LINE % (name, port, host_port) )
                vfile.write(VAGRANT_FILE_VM_STANZA_TAIL)
  
        vfile.write(VAGRANT_FILE_TAIL)
        vfile.close()
        
    #
    # To be returned to ansible with info about instances
    #        
    def _build_instance_array_for_ansible(self, vmname=None):
    
        vm_names = []
        instances = self._instances()
        if vmname != None:
            vm_names = [vmname]
        else:
            vm_names = instances.keys()
        
        ans_instances = []   
        for vm_name in vm_names:
            for inst in instances[vm_name]:
                vagrant_name = inst['vagrant_name']
                cnf = self.vg.conf(None,vagrant_name) 
                vg_data = instances[vm_name]
                if cnf != None:
                    instance_dict = dict(
                      name            = vm_name,
                      vagrant_name    = vagrant_name, 
                      id              = cnf['Host'],
                      public_ip       = cnf['HostName'],
                      internal_ip     = inst['internal_ip'],
                      public_dns_name = cnf['HostName'],
                      port            = cnf['Port'],
                      username        = cnf['User'],
                      key             = cnf['IdentityFile'],
                      status          = self.vg.status(vagrant_name)           
                    )
                    ans_instances.append(instance_dict)

        return ans_instances
         
#--------
# MAIN
#--------
def main():
    
    module = AnsibleModule(
        argument_spec = dict(
            state=dict(),
            cmd=dict(required=False, aliases = ['command']),
            box_name=dict(required=False, aliases = ['image']),
            box_path=dict(),
            vm_name=dict(),
            forward_ports=dict(),
            count = dict(default='1'), 
       )
   )
    
    state = module.params.get('state')
    cmd = module.params.get('cmd')
    box_name = module.params.get('box_name')
    box_path = module.params.get('box_path')
    vm_name = module.params.get('vm_name')
    forward_ports = module.params.get('forward_ports')     

    if forward_ports != None:
        forward_ports=forward_ports.split(',')
    if forward_ports == None: 
        forward_ports=[]

    count = module.params.get('count') 
 
    # Initialize vagrant
    vgw = VagrantWrapper()
    
    #
    # Check if we are being invoked under an idempotency idiom of "state=present" or "state=absent"
    #
    try:
        if state != None:
            
            if state != 'present' and state != 'absent':
                module.fail_json(msg = "State must be \"present\" or \"absent\" in vagrant module.")
                 
            if state == 'present':
               
                changd, insts = vgw.up(box_name, vm_name, count, box_path, forward_ports)
                module.exit_json(changed = changd, instances = insts)
                 
            if state == 'absent':
                changd = vgw.halt(vm_name)    
                module.exit_json(changed = changd, status = vgw.status(vm_name))
                 
                 
        #
        # Main command tree for old style invocation
        #                     
            
        else:

            if cmd == 'up':
            
                if count == None: 
                    count = 1
                (changd, insts) = vgw.up(box_name, vm_name, count, box_path, forward_ports)
                module.exit_json(changed = changd, instances = insts)

            elif cmd == 'status':

#            if vm_name == None:
#                module.fail_json(msg = "Error: you must specify a vm_name when calling status." )
                
                (changd, result) = vgw.status(vm_name)
                module.exit_json(changed = changd, status = result)

            elif cmd == "config" or cmd == "conf":
            
                if vm_name == None:
                    module.fail_json(msg = "Error: you must specify a vm_name when calling config." )
                (changd, cnf) = vgw.config(vm_name)
                module.exit_json(changed = changd, config = cnf)

            elif cmd == 'ssh':
            
                if vm_name == None:
                    module.fail_json(msg = "Error: you must specify a vm_name when calling ssh." )             
                            
                (changd, cnf) = vgw.config(vm_name)
                sshcmd = "ssh -i %s -p %s %s@%s" % (cnf["IdentityFile"], cnf["Port"], cnf["User"], cnf["HostName"])
                sshmsg = "Execute the command \"vagrant ssh %s\"" % (vm_name)
                module.exit_json(changed = changd, msg = sshmsg, SshCommand = sshcmd)
           
            elif cmd == 'halt':

                (changd, stats) = vgw.halt(vm_name)
                module.exit_json(changed = changd, status = stats)

            elif cmd == 'destroy':

                changd = vgw.destroy(vm_name)    
                module.exit_json(changed = changd, status = vgw.status(vm_name))
            
            elif cmd == 'clear':
            
                changd = vgw.clear()          
                module.exit_json(changed = changd)
            
            else:
            
                module.fail_json(msg = "Unknown vagrant subcommand: \"%s\"." % (cmd))
            
    except subprocess.CalledProcessError as errer:     
        module.fail_json(msg = "Vagrant command failed: %s." % (errer))
    except Exception as errer:
        module.fail_json(msg = errer.__str__())
    module.exit_json(status = "success")
  

    

# this is magic, see lib/ansible/module_common.py
#<<INCLUDE_ANSIBLE_MODULE_COMMON>>

main()