#!/usr/bin/env python
# Copyright (C) 2014 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 configure_winnt_sif, vboxmanage_path
from vmcloak.constants import VMCLOAK_ROOT
from vmcloak.deps import DependencyManager, DependencyWriter
from vmcloak.exceptions import CommandError
from vmcloak.iso import buildiso
from vmcloak.misc import add_bird, shared_parameters, register_cuckoo
from vmcloak.misc import wait_for_host, resolve_parameters
from vmcloak.verify import valid_serial_key, valid_keyboard_layout
from vmcloak.vm import VirtualBox, VBoxRPC, 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('--iso-mount', type=str, help='Mounted ISO Windows installer image.')
    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(
        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)
    elif s.vm.lower() == 'vboxrpc':
        m = VBoxRPC(s.vmname, url=s.vboxrpc_url,
                    auth=tuple(s.vboxrpc_auth.split(':', 1)),
                    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 not s.iso_mount or not os.path.isdir(s.iso_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 s.serial_key or not valid_serial_key(s.serial_key):
        log.error('The provided --serial-key is not encoded correctly.')
        log.info('Example encoding, AAAAA-BBBBB-CCCCC-DDDDD-EEEEE.')
        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)

    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)

    log.debug('Configuring WINNT.SIF')
    winnt_sif = os.path.join(VMCLOAK_ROOT, 'data', 'winnt.sif')
    buf = configure_winnt_sif(winnt_sif, s)

    # Write the WINNT.SIF file.
    _, winntsif = tempfile.mkstemp(suffix='.sif', dir=s.temp_dirpath)
    open(winntsif, 'wb').write(buf)

    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',
    )

    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)

    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)

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

    initialize_vm(m, s, 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:
        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, 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()
