#!/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.win7 import Windows7x64
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("--ramsize", type=int, help="Memory size")
    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
    if args.win7:
        h = Windows7()
        ramsize = 1024
        osversion = "win7"
    elif args.x64 or args.win7x64:
        h = Windows7x64()
        ramsize = 2048
        args.x64 = True
        osversion = "win7"
    else:
        log.error("Please provide either --winxp or --win7 or --win7x64.")
        exit(1)

    if args.ramsize is not None:
        ramsize = args.ramsize

    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, ramsize=ramsize))
    session.commit()

if __name__ == "__main__":
    main()
