#!/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 time

from vmcloak.agent import Agent
from vmcloak.misc import wait_for_host
from vmcloak.rand import random_string
from vmcloak.repository import Session, Image, Snapshot
from vmcloak.vm import VirtualBox
from vmcloak.winxp import WindowsXP
from vmcloak.win7 import Windows7

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

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("name", type=str, help="Name of the instance.")
    parser.add_argument("vmname", type=str, help="Virtual Machine name.")
    parser.add_argument("ipaddr", type=str, nargs="?", default="192.168.56.101", help="IP address.")
    parser.add_argument("--resolution", type=str, help="Screen resolution.")
    parser.add_argument("--ramsize", type=int, help="Amount of virtual memory to assign.")
    parser.add_argument("--cpus", type=int, help="Amount of CPUs to assign.")
    parser.add_argument("--hostname", type=str, default=random_string(8, 16), help="Hostname for this VM.")
    parser.add_argument("--adapter", type=str, help="Hostonly adapter for this VM.")
    parser.add_argument("--vm-visible", action="store_true", help="Start the Virtual Machine in GUI mode.")
    args = parser.parse_args()

    session = Session()

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

    m = VirtualBox(name=args.vmname)

    if image.osversion == "winxp":
        h = WindowsXP()
        ramsize = args.ramsize or 1024
    elif image.osversion == "win7" or image.osversion == "win7x64":
        h = Windows7()
        ramsize = args.ramsize or 2048

    # From now on this image is multiattach.
    image.mode = "multiattach"
    session.commit()

    m.create_vm()
    m.os_type(image.osversion)
    m.cpus(args.cpus or image.cpus)
    m.mouse("usbtablet")
    m.ramsize(ramsize)
    m.attach_hd(image.path, multi=True)
    m.hostonly(nictype=h.nictype, adapter=args.adapter or image.adapter)
    m.start_vm(visible=args.vm_visible)

    wait_for_host(image.ipaddr, image.port)
    a = Agent(image.ipaddr, image.port)
    a.ping()

    # Assign a new hostname.
    a.hostname(args.hostname)
    a.reboot()
    a.kill()

    # Wait for the reboot to kick in.
    time.sleep(10)
    wait_for_host(image.ipaddr, image.port)
    a.ping()

    if args.resolution:
        width, height = args.resolution.split("x")
        a.resolution(width, height)

    a.remove("C:\\vmcloak")
    a.static_ip(args.ipaddr, image.netmask, image.gateway)

    m.snapshot("vmcloak", "Snapshot created by VM Cloak.")
    m.stopvm()

    # Create a database entry for this snapshot.
    snapshot = Snapshot(image_id=image.id, vmname=args.vmname,
                        ipaddr=args.ipaddr, port=image.port,
                        hostname=args.hostname)
    session.add(snapshot)
    session.commit()

if __name__ == "__main__":
    main()
