///////////////////////////////////////////////////////////////////////////////
//
//  ManageVirtualMachinesCmdlet.cs
//
//  Copyright 2018-2025, Josh Moyer <JMoyer@NODOMAIN.NET>. All rights reserved.
//
//

using System;
using System.Diagnostics;
using System.IO;
using System.Management.Automation;
using System.Security.Cryptography;
using System.Threading;

namespace NDDN.VirtualMachines
{
    [Cmdlet(VerbsCommon.Get, "VMs")]
    public sealed class Get_VMs : VirtualMachineCmdlet
    {
        protected override void ProcessRecord()
        {
            int i, n;

            Process Process;

            VirtualMachine[] VMs;

            VmInvocationProfile VmToolInvocationProfile = new VmInvocationProfile();
            VmToolInvocationProfile.VmOperation = VmOperation.List;
            VmToolInvocationProfile.Vmx = null;

            Process = InvokeCmd(VmToolInvocationProfile)[0];

            n = int.Parse(Process.StandardOutput.ReadLine().Split(':')[1].Trim());
            VMs = new VirtualMachine[n];

            for (i = 0; i < n; i++)
            {
                WriteObject(new VirtualMachine(Process.StandardOutput.ReadLine()));
            }
        }
    }

/*
    [Cmdlet(VerbsData.Backup, "VM", SupportsShouldProcess = true)]
    public sealed class Backup_VM : VirtualMachineCmdlet
    {
        [Parameter(
            Position = 1,
            ValueFromPipelineByPropertyName = true)]
        public Boolean Drive;

        [Parameter(
            Position = 1,
            ValueFromPipelineByPropertyName = true)]
        public Boolean Path;

        private Boolean WasVMRunning;

        // Dangerous assumption: all VMs files will be contained in the same directory.
        // how to manage how much to back up?
        protected override void ProcessRecord()
        {
            //if (ShouldProcess(Path.GetFileNameWithoutExtension(this.VirtualMachine.Specification)))
            {
                //WasVMRunning = InvokeCmd(Process, new VmInvocationProfile())
                // copy
                // resume if was running
            }
        }
    }
*/

    [Cmdlet(VerbsData.Update, "VM")]
    public sealed class Update_VM : VirtualMachineCmdlet
    {
        protected override void ProcessRecord()
        {
            //WriteDebug("Checking if " + Path.GetFileNameWithoutExtension(this.VirtualMachine.Specification) + " is running...");

            String OutputLine;

            Process Process;

            VmInvocationProfile VmToolInvocationProfile = new VmInvocationProfile();
            VmToolInvocationProfile.VmOperation = VmOperation.List;
            VmToolInvocationProfile.Vmx = this.VirtualMachine.Specification;

            Process = InvokeCmd(VmToolInvocationProfile)[0];

            while (!Process.StandardOutput.EndOfStream)
            {
                OutputLine = Process.StandardOutput.ReadLine();
                //WriteDebug("Found " + OutputLine + "...");
                if (OutputLine == this.VirtualMachine.Specification)
                {
                    this.VirtualMachine.State = VirtualMachineState.Running;
                    //WriteDebug(this.VirtualMachine.Specification + ": " + this.VirtualMachine.State);
                    return;
                }
            }

            this.VirtualMachine.State = VirtualMachineState.NotRunning;

            //WriteDebug(this.VirtualMachine.Specification + ": " + this.VirtualMachine.State);
        }
    }

    [Cmdlet("Pause", "VM", SupportsShouldProcess = true)]
    public sealed class Pause_VM : VirtualMachineCmdlet
    {
        protected override void ProcessRecord()
        {
            //WriteDebug("PauseVM()");
            if (ShouldProcess(Path.GetFileNameWithoutExtension(this.VirtualMachine.Specification)))
            {
                VmInvocationProfile VmToolInvocationProfile = new VmInvocationProfile();
                VmToolInvocationProfile.VmOperation = VmOperation.Pause;
                VmToolInvocationProfile.Vmx = this.VirtualMachine.Specification;

                //WriteDebug("Pausing ...");

                InvokeCmd(VmToolInvocationProfile);

                //WriteDebug("Process exited with " + Process.ExitCode.ToString());
            }
        }
    }

    [Cmdlet(VerbsLifecycle.Start, "VM", SupportsShouldProcess = true)]
    public sealed class Start_VM : VirtualMachineCmdlet
    {
        [Parameter(
            Position = 1,
            ValueFromPipeline = true,
            ValueFromPipelineByPropertyName = true)]
        public Byte Interactive;

        protected override void ProcessRecord()
        {
            //WriteDebug("StartVM()");
            if (ShouldProcess(Path.GetFileNameWithoutExtension(this.VirtualMachine.Specification)))
            {
                VmInvocationProfile VmToolInvocationProfile = new VmInvocationProfile();
                switch (Interactive)
                {
                    case 0:
                        VmToolInvocationProfile.VmType = VmType.Background;
                        goto default;
                    case 1:
                        VmToolInvocationProfile.VmType = VmType.Foreground;
                        goto default;
                    case 2:
                        VmToolInvocationProfile.VmType = VmType.Fullscreen;
                        goto default;
                    default:
                        VmToolInvocationProfile.VmOperation = VmOperation.Start;
                        break;
                }
                VmToolInvocationProfile.Vmx = this.VirtualMachine.Specification;

                InvokeCmd(VmToolInvocationProfile);
            }
        }
    }

    [Cmdlet(VerbsLifecycle.Suspend, "VM", SupportsShouldProcess = true)]
    public sealed class Suspend_VM : VirtualMachineCmdlet
    {
        protected override void ProcessRecord()
        {
            //WriteDebug("SuspendVM()");
            if (ShouldProcess(Path.GetFileNameWithoutExtension(this.VirtualMachine.Specification)))
            {
                VmInvocationProfile VmToolInvocationProfile = new VmInvocationProfile();
                VmToolInvocationProfile.VmOperation = VmOperation.Suspend;
                VmToolInvocationProfile.Vmx = this.VirtualMachine.Specification;
                VmToolInvocationProfile.VmOperationPressure = VmOperationPressure.Hard;

                //WriteDebug("Suspending ...");

                InvokeCmd(VmToolInvocationProfile);

                //WriteDebug("Process exited with " + Process.ExitCode.ToString());
            }
        }
    }

    [Cmdlet("Unpause", "VM", SupportsShouldProcess = true)]
    public sealed class Unpause_VM : VirtualMachineCmdlet
    {
        protected override void ProcessRecord()
        {
            //WriteDebug("UnpauseVM()");
            if (ShouldProcess(Path.GetFileNameWithoutExtension(this.VirtualMachine.Specification)))
            {
                VmInvocationProfile VmToolInvocationProfile = new VmInvocationProfile();
                VmToolInvocationProfile.VmOperation = VmOperation.Unpause;
                VmToolInvocationProfile.Vmx = this.VirtualMachine.Specification;

                //WriteDebug("Unpausing ...");

                InvokeCmd(VmToolInvocationProfile);

                //WriteDebug("Process exited with " + Process.ExitCode.ToString());
            }
        }
    }

    public class VirtualMachineCmdlet : Cmdlet
    {
        [Parameter(
            ValueFromPipeline = true,
            ValueFromPipelineByPropertyName = true,
            Position = 0,
            Mandatory = true)]
        [ValidateNotNullOrEmpty]
        public VirtualMachine VirtualMachine
        {
            get { return this.virtualMachine; }
            set { this.virtualMachine = value; }
        }
        private VirtualMachine virtualMachine;

        public VirtualMachineCmdlet()
        {
            //WriteDebug("VirtualMachineCmdlet()");
            this.VmWarePath = Environment.GetEnvironmentVariable("VmWarePath");
            //WriteDebug("$Env:VmWarePath = " + this.VmWarePath);
        }

        internal String VmRunExePath;

        internal String VmWareExePath;

        [Parameter(
            ValueFromPipeline = true,
            ValueFromPipelineByPropertyName = true)]
        [ValidateNotNullOrEmpty]
        public String VmWarePath
        {
            get { return vmwarePath; }
            set
            {
                //WriteDebug("setting VmWarePath to " + value + "...");
                vmwarePath = value;

                if (vmwarePath == null || vmwarePath == "" || vmwarePath == "NDDN.VirtualMachines.VirtualMachine")
                    vmwarePath = @"C:\Program Files (x86)\VMware\VMware Workstation";

                //WriteDebug("vmwarePath");
                if (!Directory.Exists(vmwarePath))
                    ThrowTerminatingError(new ErrorRecord(new FileNotFoundException("VmWarePath"), "ERROR-VmWarePath", ErrorCategory.ObjectNotFound, this));

                //WriteDebug("VmRunExePath");
                VmRunExePath = vmwarePath + "\\vmrun.exe";
                if (!File.Exists(VmRunExePath))
                    ThrowTerminatingError(new ErrorRecord(new FileNotFoundException("VmRunExePath"), "ERROR-VmRunExePath", ErrorCategory.ObjectNotFound, this));

                //WriteDebug("VmWareExePath");
                VmWareExePath = vmwarePath + "\\vmware.exe";
                if (!File.Exists(VmWareExePath))
                    ThrowTerminatingError(new ErrorRecord(new FileNotFoundException("VmWareExePath"), "ERROR-VmWareExePath", ErrorCategory.ObjectNotFound, this));
            }
        }
        internal String vmwarePath;

        internal Process[] InvokeCmd(VmInvocationProfile VmInvocationProfile)
        {
            //TODO: catch exceptions

            /* VmRun.exe VM start section
             * Always run.
             */

            Process VmRunProcess = new Process();
            VmRunProcess.StartInfo.CreateNoWindow = true;
            VmRunProcess.StartInfo.FileName = VmRunExePath;
            VmRunProcess.StartInfo.RedirectStandardOutput = true;
            VmRunProcess.StartInfo.UseShellExecute = false;

            // Hard-code for VMware Workstation for now
            VmRunProcess.StartInfo.Arguments += "-t ws ";

            byte[] vmEncryptedPassword = VmInvocationProfile.VmEncryptedPassword;
            if (vmEncryptedPassword != null)
            {
                string vmDecryptedPassword =
                    System.Text.Encoding.Unicode.GetString(
                        ProtectedData.Unprotect(
                            vmEncryptedPassword,
                            null,
                            DataProtectionScope.CurrentUser));
                VmRunProcess.StartInfo.Arguments +=
                    "-vp \"" + vmDecryptedPassword + "\" ";
                vmDecryptedPassword = null;
            }

            VmRunProcess.StartInfo.Arguments +=
                VmInvocationProfile.VmOperation.ToString() +
                " \"" + VmInvocationProfile.Vmx + "\" ";

            if (VmInvocationProfile.VmOperation == VmOperation.Suspend)
            {
                switch (VmInvocationProfile.VmOperationPressure)
                {
                    case VmOperationPressure.Hard:
                        VmRunProcess.StartInfo.Arguments += "hard ";
                        break;
                    case VmOperationPressure.Soft:
                        VmRunProcess.StartInfo.Arguments += "soft ";
                        break;
                }
            }

            if (VmInvocationProfile.VmOperation == VmOperation.Start)
            {
                VmRunProcess.StartInfo.Arguments += "nogui ";
            }

            WriteDebug("Will run " + VmRunProcess.StartInfo.FileName);
            WriteDebug("With args " + VmRunProcess.StartInfo.Arguments);
            VmRunProcess.Start();
            if (VmInvocationProfile.Sync)
            {
                VmRunProcess.WaitForExit();
                //WriteDebug("Process exited with " + VmRunProcess.ExitCode.ToString());
            }

            /* VMware.exe MKS start section
             * Only run during GUI start operations.
             */

            Process VmWareProcess = new Process();

            if (VmInvocationProfile.VmOperation == VmOperation.Start)
            {
                VmWareProcess.StartInfo.CreateNoWindow = true;
                VmWareProcess.StartInfo.FileName = VmWareExePath;
                VmRunProcess.StartInfo.RedirectStandardOutput = true;
                VmRunProcess.StartInfo.UseShellExecute = false;

                switch (VmInvocationProfile.VmType)
                {
                    case VmType.Background:
                        break;

                    case VmType.Foreground:
                        VmWareProcess.StartInfo.Arguments += "-n ";
                        goto default;

                    case VmType.Fullscreen:
                        VmWareProcess.StartInfo.Arguments += "-n -f ";
                        goto default;

                    default:
                        VmWareProcess.StartInfo.Arguments += " \"" + VmInvocationProfile.Vmx + "\" ";

                        WriteDebug("Will run " + VmWareProcess.StartInfo.FileName);
                        WriteDebug("With args " + VmWareProcess.StartInfo.Arguments);
                        VmWareProcess.Start();
                        /* Always seems to take exactly 30 seconds for the VM window to appear...
                         * since there's likely some synchronization requirement that causes that
                         * we sleep here for 35 seconds to prevent overlapping requests.
                         */
                        Thread.Sleep(35000);
                        //WriteDebug("Process exited with " + VmWareProcess.ExitCode.ToString());
                        break;
                }
            }

            return new Process[] { VmRunProcess, VmWareProcess };

        }
    }

    internal class VmInvocationProfile
    {
        private Boolean sync;
        private VmOperation vmOperation;
        private VmOperationPressure vmOperationPressure;
        private VmType vmType;
        private String vmx;

        public Boolean Sync
        {
            get { return sync; }
            set { sync = value; }
        }

        public byte[] VmEncryptedPassword
        {
            get
            {
                try
                {
                    return File.ReadAllBytes(
                        Path.GetFullPath(vmx).
                            Split(".".ToCharArray()[0])[0] +
                                ".pwbin");
                }
                catch (FileNotFoundException)
                {
                    return null;
                }
            }
        }

        public VmOperation VmOperation
        {
            get { return vmOperation; }
            set { vmOperation = value; }
        }
        public VmOperationPressure VmOperationPressure
        {
            get { return vmOperationPressure; }
            set { vmOperationPressure = value; }
        }

        public VmType VmType
        {
            get { return vmType; }
            set { vmType = value; }
        }

        public String Vmx
        {
            get { return vmx; }
            set { vmx = value; }
        }


        public VmInvocationProfile()
        {
            this.sync = true;
        }
    }

    internal enum VmType
    {
        Background,
        Foreground,
        Fullscreen
}

    internal enum VmOperation
    {
        List,
        Start,
        Suspend,
        Pause,
        Unpause
    }

    internal enum VmOperationPressure
    {
        Default,
        Hard,
        Soft
    }
}