#!/usr/bin/env python
# Copyright (C) 2014-2015 Jurriaan Bremer.
# This file is part of VMCloak - http://www.vmcloak.org/.
# See the file 'docs/LICENSE.txt' for copying permission.

import logging
import os
import pkg_resources
import shutil
import socket
import sys
import tempfile
import time

from vmcloak.conf import vboxmanage_path
from vmcloak.deps import DependencyManager, DependencyWriter
from vmcloak.exceptions import CommandError
from vmcloak.misc import add_bird, shared_parameters, register_cuckoo
from vmcloak.misc import wait_for_host, resolve_parameters
from vmcloak.os import WindowsXP, Windows7
from vmcloak.verify import valid_keyboard_layout
from vmcloak.vm import VirtualBox, initialize_vm

BIN_ROOT = os.path.abspath(os.path.dirname(__file__))

logging.basicConfig(level=logging.INFO)
log = logging.getLogger('vmcloak')


def main():
    parser, defaults, types = shared_parameters()
    parser.add_argument('--version', action='store_true', help='Display the VMCloak version.')
    parser.add_argument('--version-int', action='store_true', help='Display the VMCloak version as integer.')
    parser.add_argument('--delete', action='store_true', help='Completely delete a Virtual Machine and its associated files.')
    parser.add_argument('--winxp', action='store_true', default=None, help='Create a Windows XP Virtual Machine.')
    parser.add_argument('--win7', action='store_true', default=None, help='Create a Windows 7 Virtual Machine.')
    parser.add_argument('--win7x64', action='store_true', default=None, help='Shorthand for --win7 --x64.')
    parser.add_argument('--iso-mount', type=str, help='Mounted ISO Windows installer image.')
    parser.add_argument('--x64', action='store_true', default=None, help='Install 64-bit software in the Virtual Machine.')
    parser.add_argument('--serial-key', type=str, help='Windows Serial Key.')
    parser.add_argument('--dependencies', type=str, help='Comma-separated list of all dependencies in the Virtual Machine.')
    parser.add_argument('--bird', action='store_true', default=None, help='Make an immutable harddisk rather than a Virtual Machine.')
    parser.add_argument('--host-init-ip', type=str, help='Guest IP address to initialize initial connectivity.')
    parser.add_argument('--host-init-mask', type=str, help='Guest IP address mask to initialize initial connectivity.')
    parser.add_argument('--host-init-gateway', type=str, help='Guest IP address gateway to initialize initial connectivity.')

    defaults.update(dict(
        winxp=False,
        win7=False,
        win7x64=False,
        bird=False,
    ))

    args = parser.parse_args()

    if args.verbose:
        log.setLevel(logging.DEBUG)

    try:
        s = resolve_parameters(args, defaults, types)
    except Exception as e:
        sys.exit(e.message)

    # Not really sure if this is possible in a cleaner way, but if no
    # arguments have been given on the command line, then show the usage.
    if len(sys.argv) == 1:
        parser.print_help()
        exit(0)

    if args.version:
        print pkg_resources.require('VMCloak')[0].version
        exit(0)

    if args.version_int:
        version = pkg_resources.require('VMCloak')[0].version.split('.')
        print 1000 * int(version[0]) + 100 * int(version[1]) + int(version[2])
        exit(0)

    if not s.vmname:
        log.error('A name for the Virtual Machine is required.')
        exit(1)

    if s.register_cuckoo and not os.path.isdir(s.cuckoo):
        log.error('To register the Virtual Machine with Cuckoo '
                  'the Cuckoo directory has to be provided with --cuckoo.')
        log.info('To disable registering please provide --no-register-cuckoo '
                 'or specify register-cuckoo = false in the configuration.')
        exit(1)

    if s.vm.lower() == 'virtualbox':
        vboxmanage = vboxmanage_path(s)
        if not vboxmanage:
            exit(1)

        m = VirtualBox(s.vmname, s.vm_dir, s.data_dir,
                       vboxmanage=vboxmanage, temp_dir=s.temp_dirpath)
    else:
        log.error('Only the VirtualBox --vm is supported as of now.')
        exit(1)

    if not m.api_status():
        exit(1)

    if s.delete:
        log.info('Unregistering and deleting the VM and its associated '
                 'files.')
        m.delete_vm()
        exit(0)

    if s.win7x64:
        s.win7 = True
        s.x64 = True

    if s.winxp and s.win7:
        log.error('Please provide just --winxp or --win7, not both.')
        exit(1)

    if s.winxp:
        h = WindowsXP()
    elif s.win7:
        h = Windows7()
    else:
        log.error('Please provided either --winxp or --win7.')
        exit(1)

    if not os.path.isdir(s.iso_mount or h.mount):
        log.error('Please specify --iso-mount to a directory containing the '
                  'mounted Windows Installer ISO image.')
        log.info('Refer to the documentation on mounting an .iso image.')
        exit(1)

    if not h.set_serial_key(s.serial_key):
        exit(1)

    log.debug('Checking whether the keyboard layout is valid.')
    if not valid_keyboard_layout(s.keyboard_layout):
        log.error('Invalid keyboard layout: %s', s.keyboard_layout)
        log.info('Please use one provided as described in the documentation.')
        exit(1)

    if not os.path.isdir(os.path.join(os.getenv('HOME'), '.vmcloak')):
        os.mkdir(os.path.join(os.getenv('HOME'), '.vmcloak'))

    log.debug('Static Host IP: %s, Port %s', s.host_ip, s.host_port)
    log.debug('Static Guest hostonly IP: %s', s.hostonly_ip)
    log.debug('Static Guest bridged IP: %s', s.bridged_ip)

    h.configure(s)

    settings_bat = dict(
        BRIDGED='yes' if s.bridged else 'no',
        BRIDGEDIP=s.bridged_ip,
        BRIDGEDMASK=s.bridged_mask,
        BRIDGEDGATEWAY=s.bridged_gateway,
        DNSSERVER=s.dns_server,
        GUEST_IP=s.host_init_ip if s.bird else s.hostonly_ip,
        GUEST_MASK=s.host_init_mask if s.bird else s.hostonly_mask,
        GUEST_GATEWAY=s.host_init_gateway if s.bird else s.hostonly_gateway,
        VMMODE='bird' if s.bird else 'normal',
    )

    settings_py = dict(
        HOST_IP=s.host_ip,
        HOST_PORT=s.host_port,
        RESOLUTION=s.resolution,
        VMMODE='bird' if s.bird else 'normal',
    )

    if s.run_executable:
        settings_bat['RUNEXEC'] = os.path.basename(s.run_executable)
    else:
        settings_bat['RUNEXEC'] = ''

    bootstrap = tempfile.mkdtemp(dir=s.temp_dirpath)

    vmcloak_dir = os.path.join(bootstrap, 'vmcloak')
    os.mkdir(vmcloak_dir)

    # Write the configuration values for bootstrap.bat.
    with open(os.path.join(vmcloak_dir, 'settings.bat'), 'wb') as f:
        for key, value in settings_bat.items():
            print>>f, 'set %s=%s' % (key, value)

    # Write the configuration values for bootstrap.py.
    with open(os.path.join(vmcloak_dir, 'settings.py'), 'wb') as f:
        for key, value in settings_py.items():
            print>>f, '%s = %r' % (key.lower(), value)

    os.mkdir(os.path.join(vmcloak_dir, 'deps'))

    dm = DependencyManager(s.deps_directory, s.deps_repository)
    deps = DependencyWriter(dm, vmcloak_dir, bitsize_64=s.x64)

    if not deps.add('python27'):
        exit(1)

    for d in s.dependencies.split():
        if not d.strip():
            continue

        if not deps.add(d.strip()):
            exit(1)

    deps.write()

    # Write the auxiliary files.
    if s.auxiliary:
        aux_path = os.path.join(bootstrap, s.auxiliary_local)
        shutil.copytree(s.auxiliary, aux_path)

    # Write the run executable file.
    if s.run_executable:
        exec_fname = os.path.basename(s.run_executable)
        exec_path = os.path.join(vmcloak_dir, exec_fname)
        shutil.copy(s.run_executable, exec_path)

    # Create the ISO file.
    log.debug('Creating ISO file.')
    if not h.buildiso(s.iso_mount or h.mount, m.iso_path, bootstrap,
                      s.temp_dirpath):
        shutil.rmtree(bootstrap)
        exit(1)

    shutil.rmtree(bootstrap)

    initialize_vm(m, s, h, clone=False)

    log.info('Starting the Virtual Machine %r to install Windows.', s.vmname)
    m.start_vm(visible=s.vm_visible)

    log.debug('Setting the resolution to %s', s.resolution)

    log.info('Waiting for the Virtual Machine %r to connect back, this may '
             'take up to 30 minutes.', s.vmname)
    t = time.time()

    if not s.bird:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((s.host_ip, s.host_port))
        sock.listen(1)

        guest, _ = sock.accept()
        guest.close()
    else:
        # Wait until the Virtual Machine is shutdown.
        while True:
            try:
                if m.vminfo('VMState') == 'poweroff':
                    break
            except CommandError:
                pass

            time.sleep(1)

    log.debug('It took %d seconds to install Windows!', time.time() - t)

    if s.bird:
        # Wait a couple of seconds to really ensure the Virtual Machine
        # has shutdown or initialized.
        time.sleep(5)
    else:
        # Wait for the host to come up.
        wait_for_host(s.hostonly_ip)

    log.debug('Detaching the Windows Installation disk.')
    m.detach_iso()

    log.debug('Removing Windows Installation ISO file.')
    os.unlink(m.iso_path)

    if s.bird:
        log.debug('Removing the harddisk from the VM.')
        m.remove_hd()

        log.debug('Deleting the temporary Virtual Machine.')
        m.delete_vm()

        log.debug('Making the base harddisk immutable.')
        m.immutable_hd()

        log.debug('Registering this bird with the local VMCloak database.')
        add_bird(s.vmname, vmtype=h.name, hdd_path=m.hdd_path)
    else:
        # Give the system a little bit of time to fully initialize.
        time.sleep(10)

        log.debug('Taking a snapshot of the current state.')
        m.snapshot('vmcloak', 'Snapshot created by VMCloak.')

        log.debug('Powering off the virtual machine.')
        m.stopvm()

        if s.register_cuckoo:
            register_cuckoo(hostonly_ip=s.hostonly_ip, tags=s.tags,
                            vmname=s.vmname, cuckoo_dirpath=s.cuckoo)

    log.info('Virtual Machine %r created successfully.', s.vmname)

if __name__ == '__main__':
    main()
