#!/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 argparse
import lockfile
import logging
import os.path
import pkg_resources
import shutil
import socket
import subprocess
import sys
import tempfile
import time

from vmcloak.conf import Configuration, configure_winnt_sif, vboxmanage_path
from vmcloak.constants import VMCLOAK_ROOT
from vmcloak.deps import DependencyManager, DependencyWriter
from vmcloak.iso import buildiso
from vmcloak.misc import DummyLock
from vmcloak.verify import valid_serial_key, valid_keyboard_layout
from vmcloak.vm import VirtualBox, VBoxRPC

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

logging.basicConfig(level=logging.INFO)
logging.getLogger('requests').setLevel(logging.WARNING)
logging.getLogger('urllib3').setLevel(logging.WARNING)
log = logging.getLogger('vmcloak')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('vmname', type=str, nargs='?', help='Name of the Virtual Machine.')
    parser.add_argument('-d', '--debug', action='store_true', help='Display debug messages.')
    parser.add_argument('-v', '--version', action='store_true', help='Display the VMCloak version.')
    parser.add_argument('--cuckoo', type=str, help='Directory where Cuckoo is located.')
    parser.add_argument('--vm-dir', type=str, help='Base directory for the virtual machine and its associated files.')
    parser.add_argument('--data-dir', type=str, help='Base directory for the virtual machine harddisks and images.')
    parser.add_argument('--vm', type=str, help='Virtual Machine Software (VirtualBox.)')
    parser.add_argument('--vboxrpc-url', type=str, help='URL to VBoxRPC instance.')
    parser.add_argument('--vboxrpc-auth', type=str, help='Credentials to VBoxRPC instance.')
    parser.add_argument('--delete', action='store_true', help='Completely delete a Virtual Machine and its associated files.')
    parser.add_argument('--ramsize', help='Available virtual memory (in MB) for this virtual machine.')
    parser.add_argument('--resolution', type=str, help='Virtual Machine resolution.')
    parser.add_argument('--hdsize', help='Maximum size (in MB) of the dynamically allocated harddisk.')
    parser.add_argument('--iso-mount', type=str, help='Mounted ISO Windows installer image.')
    parser.add_argument('--host-ip', type=str, help='Static IP address to bind to on the Host.')
    parser.add_argument('--hostonly-ip', type=str, help='Static IP address to use on the Guest for the hostonly network.')
    parser.add_argument('--hostonly-mask', type=str, help='Static IP address mask to use on the Guest for the hostonly network.')
    parser.add_argument('--hostonly-gateway', type=str, help='Static IP address gateway to use on the Guest for the hostonly network.')
    parser.add_argument('--hostonly-macaddr', type=str, help='Mac address for the hostonly interface.')
    parser.add_argument('--bridged', type=str, help='Network interface for the bridged network.')
    parser.add_argument('--bridged-ip', type=str, help='Static IP address to use on the Guest for the bridged network.')
    parser.add_argument('--bridged-mask', type=str, help='Static IP address mask to use on the Guest for the bridged network.')
    parser.add_argument('--bridged-gateway', type=str, help='Static IP address gateway to use on the Guest for the bridged network.')
    parser.add_argument('--bridged-macaddr', type=str, help='Mac address for the bridged interface.')
    parser.add_argument('--nat', action='store_true', help='Name of the NAT network to attach to.')
    parser.add_argument('--dns-server', type=str, help='Address of DNS server to be used.')
    parser.add_argument('--hwvirt', action='store_true', default=None, help='Explicitly enable Hardware Virtualization.')
    parser.add_argument('--no-hwvirt', action='store_false', default=None, dest='hwvirt', help='Explicitly disable Hardware Virtualization.')
    parser.add_argument('--serial-key', type=str, help='Windows Serial Key.')
    parser.add_argument('--tags', type=str, help='Cuckoo Tags for the Virtual Machine.')
    parser.add_argument('--no-register-cuckoo', action='store_false', default=None, dest='register_cuckoo', help='Explicitly disable registering the Virtual Machine with Cuckoo upon completion.')
    parser.add_argument('--vboxmanage', type=str, help='Path to VBoxManage.')
    parser.add_argument('--dependencies', type=str, help='Comma-separated list of all dependencies in the Virtual Machine.')
    parser.add_argument('--vm-visible', action='store_true', default=None, help='Start the Virtual Machine in GUI mode.')
    parser.add_argument('--keyboard-layout', type=str, help='Keyboard Layout within the Virtual Machine.')
    parser.add_argument('--cpu-count', help='Number of CPUs to use with this Virtual Machine.')
    parser.add_argument('--vrde', action='store_true', help='Enable VRDE support in VirtualBox.')
    parser.add_argument('--lock-dirpath', type=str, help='Path to directory for creating an inter-process lock.')
    parser.add_argument('--hwconfig-profile', type=str, help='Take a particular hardware profile.')
    parser.add_argument('--auxiliary', type=str, help='Path to a directory containing auxiliary files that should be shipped to the Virtual Machine.')
    parser.add_argument('--auxiliary-local', type=str, help='Overwrite the directory path to the auxiliary files in the Virtual Machine.')
    parser.add_argument('-u', '--unlock', action='store_true', default=None, help='Unlock the Virtual Machine mutex/lock.')
    parser.add_argument('--temp-dirpath', type=str, help='Directory where to put temporary files.')
    parser.add_argument('-s', '--settings', type=str, default=[], action='append', help='Configuration file with various settings.')
    parser.add_argument('-r', '--recommended-settings', action='store_true', help='Use the recommended settings.')
    parser.add_argument('--deps-directory', type=str, help='Dependency directory.')
    parser.add_argument('--deps-repository', type=str, help='Dependency repository.')

    defaults = dict(
        vm='virtualbox',
        vboxrpc_url='http://localhost:9002/',
        vboxrpc_auth='root:toor',
        ramsize=1024,
        resolution='1024x768',
        hdsize=256*1024,
        host_ip='192.168.56.1',
        hostonly_ip='192.168.56.101',
        hostonly_mask='255.255.255.0',
        hostonly_gateway='192.168.56.1',
        bridged_mask='255.255.255.0',
        dns_server='8.8.8.8',
        tags='',
        vboxmanage='/usr/bin/VBoxManage',
        vm_visible=False,
        keyboard_layout='US',
        cpu_count=1,
        register_cuckoo=True,
        dependencies='',
        auxiliary_local='auxiliary',
        deps_directory=os.path.join(os.getenv('HOME'), '.vmcloak', 'deps'),
        deps_repository='https://raw.githubusercontent.com/jbremer/vmcloak-deps/master/',
        vrde=False
    )

    types = dict(
        ramsize=int,
        hdsize=int,
        cpu_count=int,
    )

    args = parser.parse_args()

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

    s = Configuration()

    if args.recommended_settings:
        s.from_file(os.path.join(VMCLOAK_ROOT, 'data', 'recommended.ini'))

    for settings in args.settings:
        s.from_file(settings)

    s.from_args(args)
    s.from_defaults(defaults)

    for k, v in s.conf.items():
        if k in types:
            try:
                s.conf[k] = types[k](v)
            except ValueError:
                log.critical('Flag %r should be an %r.', k, types[k].__name__)
                exit(1)

    # 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:
        version = pkg_resources.require('VMCloak')[0].version
        print 'VMCloak version', version
        exit(0)

    if args.unlock:
        # Default lock dirpath.
        default_lock_dirpath = os.path.join(os.getenv('HOME'),
                                            '.vmcloak', 'lock')

        path = '%s.lock' % (s.lock_dirpath or default_lock_dirpath)
        if os.path.isdir(path):
            shutil.rmtree(path)
        exit(0)

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

    if not s.cuckoo and s.register_cuckoo:
        log.error('To register the Virtual Machine with Cuckoo '
                  'the Cuckoo directory has to be provided.')
        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(':')),
                    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'))

    if s.lock_dirpath:
        lock = lockfile.MkdirFileLock(s.lock_dirpath)
    else:
        lock = DummyLock()

    try:
        log.debug('Awaiting the opportunity to obtain the vmcloak lock.')
        lock.acquire(timeout=1)
    except lockfile.LockTimeout:
        log.warning('If no other instances of vmcloak are running then '
                    'unlock the vmcloak lock by running the following '
                    'command: vmcloak --unlock')

        lock.acquire()

    log.debug('Static Host IP: %s', s.host_ip)
    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.bind((s.host_ip, 0))
    sock.listen(1)
    _, port = sock.getsockname()

    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(
        HOSTONLYIP=s.hostonly_ip,
        HOSTONLYMASK=s.hostonly_mask,
        HOSTONLYGATEWAY=s.hostonly_gateway,
        BRIDGED='yes' if s.bridged else 'no',
        BRIDGEDIP=s.bridged_ip,
        BRIDGEDMASK=s.bridged_mask,
        BRIDGEDGATEWAY=s.bridged_gateway,
        DNSSERVER=s.dns_server,
    )

    settings_py = dict(
        HOST_IP=s.host_ip,
        HOST_PORT=port,
        RESOLUTION=s.resolution,
    )

    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, 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'):
        lock.release()
        exit(1)

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

        if not deps.add(d.strip()):
            lock.release()
            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):
        lock.release()
        exit(1)

    log.debug('Creating VM.')
    log.debug(m.create_vm())

    m.ramsize(s.ramsize)
    m.os_type(os='xp', sp=3)

    log.debug('Creating Harddisk.')
    m.create_hd(s.hdsize)

    log.debug('Temporarily attaching DVD-Rom unit for the ISO installer.')
    m.attach_iso(m.iso_path)

    log.debug('Randomizing Hardware.')
    m.init_vm(profile=s.hwconfig_profile)

    log.debug('Setting CPU count to %d', s.cpu_count)
    m.cpus(s.cpu_count)

    log.debug('Checking VirtualBox hostonly network.')
    if m.hostonly(macaddr=s.hostonly_macaddr) is False:
        lock.release()
        exit(1)

    if s.nat:
        m.nat()

    if s.bridged:
        m.bridged(s.bridged, macaddr=s.bridged_macaddr)

    if s.hwvirt is not None:
        if s.hwvirt:
            log.debug('Enabling Hardware Virtualization.')
        else:
            log.debug('Disabling Hardware Virtualization.')
        m.hwvirt(s.hwvirt)

    if args.vrde:
        m.vrde(True)

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

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

    lock.release()

    guest, _ = sock.accept()
    log.debug('It took %d seconds to install Windows!', time.time() - t)

    lock.acquire()

    log.debug('Setting the resolution to %s', s.resolution)
    if ord(guest.recv(1)):
        log.debug('Resolution was set successfully.')
    else:
        log.error('Error setting the resolution.')

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

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

    lock.release()

    # Give the system a little bit of time to fully initialize.
    time.sleep(10)

    lock.acquire()

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

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

    if s.register_cuckoo:
        log.debug('Registering the Virtual Machine with Cuckoo.')
        try:
            machine_py = os.path.join(s.cuckoo, 'utils', 'machine.py')
            subprocess.check_call([machine_py, '--add',
                                   '--ip', s.hostonly_ip,
                                   '--platform', 'windows',
                                   '--tags', s.tags,
                                   '--snapshot', 'vmcloak',
                                   s.vmname],
                                  cwd=s.cuckoo)
        except OSError as e:
            log.error('Is $CUCKOO/utils/machine.py executable? -> %s', e)
            lock.release()
            exit(1)
        except subprocess.CalledProcessError as e:
            log.error('Error registering the VM: %s.', e)
            lock.release()
            exit(1)

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

if __name__ == '__main__':
    main()
