#!/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 argparse
import logging
import os.path
import shutil
import tempfile

from vmcloak.dependencies import Python27
from vmcloak.winxp import WindowsXP
from vmcloak.win7 import Windows7
from vmcloak.repository import image_path, Session, Image
from vmcloak.vm import VirtualBox

logging.basicConfig()
log = logging.getLogger("vmcloak-init")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("name", type=str, help="Name of the new instance.")
    parser.add_argument("--winxp", action="store_true", help="This is a Windows XP instance.")
    parser.add_argument("--win7", action="store_true", help="This is a Windows 7 instance.")
    parser.add_argument("--x64", action="store_true", help="This is a 64-bit OS.")
    parser.add_argument("--win7x64", action="store_true", help="This is a Windows 7 64-bit instance.")
    parser.add_argument("--product", type=str, help="Windows 7 product version.")
    parser.add_argument("--vm", type=str, default="virtualbox", help="Virtual Machinery.")
    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("--ip", type=str, default="192.168.56.2", help="Guest IP address.")
    parser.add_argument("--port", type=int, default=8000, help="Port to run the Agent on.")
    parser.add_argument("--adapter", type=str, default="vboxnet0", help="Network adapter.")
    parser.add_argument("--netmask", type=str, default="255.255.255.0", help="Guest IP address.")
    parser.add_argument("--gateway", type=str, default="192.168.56.1", help="Guest IP address.")
    parser.add_argument("--dns", type=str, default="8.8.8.8", help="DNS Server.")
    parser.add_argument("--cpus", type=int, default=1, help="CPU count.")
    parser.add_argument("--tempdir", type=str, default=tempfile.gettempdir(), help="Temporary directory to build the ISO file.")
    parser.add_argument("--resolution", type=str, default="1024x768", help="Screen resolution.")
    parser.add_argument("--vm-visible", action="store_true", help="Start the Virtual Machine in GUI mode.")
    parser.add_argument("-d", "--debug", action="store_true", help="Install Virtual Machine in debug mode.")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logging.")
    args = parser.parse_args()

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

    session = Session()
    image = session.query(Image).filter_by(name=args.name).first()
    if image:
        log.error("Image already exists: %s", args.name)
        exit(1)

    if args.vm != "virtualbox":
        log.error("Only the VirtualBox Machinery is supported at this point.")
        exit(1)

    if args.winxp:
        h = WindowsXP()
        osversion = "winxp"
        ramsize = 1024
    elif args.win7 or args.win7x64:
        h = Windows7()
        ramsize = 2048
        if args.x64 or args.win7x64:
            osversion = "win7x64"
            args.x64 = True
        else:
            osversion = "win7"
    else:
        log.error("Please provide either --winxp or --win7.")
        exit(1)

    if not os.path.isdir(args.iso_mount or h.mount) or \
            not os.listdir(args.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(args.serial_key):
        exit(1)

    h.configure(args)

    reso_width, reso_height = args.resolution.split("x")

    settings = dict(
        GUEST_IP=args.ip,
        AGENT_PORT=args.port,
        GUEST_MASK=args.netmask,
        GUEST_GATEWAY=args.gateway,
        DNSSERVER=args.dns,
        DEBUG="yes" if args.debug else "no",
        RESO_WIDTH=reso_width,
        RESO_HEIGHT=reso_height,
    )

    bootstrap = tempfile.mkdtemp(dir=args.tempdir)

    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.items():
            print>>f, "set %s=%s" % (key, value)

    # Download the Python dependency and set it up for bootstrapping the VM.
    d = Python27(i=Image(osversion=osversion))
    d.download()
    shutil.copy(d.filepath, vmcloak_dir)

    iso_path = os.path.join(args.tempdir, "%s.iso" % args.name)
    hdd_path = os.path.join(image_path, "%s.vdi" % args.name)
    m = VirtualBox(name=args.name)

    if not h.buildiso(args.iso_mount or h.mount, iso_path, bootstrap,
                      args.tempdir):
        shutil.rmtree(bootstrap)
        exit(1)

    shutil.rmtree(bootstrap)

    m.create_vm()
    m.os_type(osversion)
    m.cpus(args.cpus)
    m.mouse("usbtablet")
    m.ramsize(ramsize)
    m.create_hd(hdd_path)
    m.attach_iso(iso_path)
    m.hostonly(nictype=h.nictype, adapter=args.adapter)

    log.info("Starting the Virtual Machine %r to install Windows.", args.name)
    m.start_vm(visible=args.vm_visible)

    m.wait_for_state(shutdown=True)

    m.detach_iso()
    os.unlink(iso_path)

    m.remove_hd()
    m.compact_hd(hdd_path)
    m.delete_vm()

    log.info("Added image %r to the repository.", args.name)
    session.add(Image(name=args.name, path=hdd_path, osversion=osversion,
                      servicepack="%s" % h.service_pack, mode="normal",
                      ipaddr=args.ip, port=args.port, adapter=args.adapter,
                      netmask=args.netmask, gateway=args.gateway,
                      cpus=args.cpus))
    session.commit()

if __name__ == "__main__":
    main()
