using System;
using System.Linq;
using System.IO;
using System.Collections.Generic;
using Microsoft.Management.Infrastructure;
using Microsoft.Management.Infrastructure.Options;

namespace ActivationApp
{
    class Program
    {
        static void Main(String[] args)
        {
            if (args.Length == 0)
            {
                MI.LogMessageToFile("Missing command arguments, exiting...");
                System.Environment.Exit(1);
            }

            MI.LogMessageToFile("=================================");
            MI.LogMessageToFile("Windows Activation app started...");

            // Windows version WMI query
            string osVersion = "select Caption from Win32_OperatingSystem";

            // Query custom MDT WMI class to verify MDT based deployments
            string mdtClass = "SELECT BuildVersion FROM Cust_BuildVer";

            //Office App ID
            string officeAppId = "0ff1ce15-a989-479d-af46-f275c6370663";

            // Windows license WMI query
            string licenseStatus = "SELECT LicenseStatus,Description FROM SoftwareLicensingProduct" +
                " WHERE PartialProductKey IS NOT NULL and Description LIKE 'Windows%'";

            // Office license WMI query
            string officeLicenseStatus = "SELECT LicenseStatus,Description,ID,ApplicationID,Name,PartialProductKey FROM SoftwareLicensingProduct" +
                " WHERE PartialProductKey IS NOT NULL and ApplicationId = '" + officeAppId + "'";

            string licenseKey = null;
            string licensingService = "SoftwareLicensingService";
            string rootNS = "root\\cimv2";
            string mdtNS = "root\\DataCollect";

            // All keys here are placeholders, replace with the actual keys before compiling
            string serverKey = "XXXXX-XXXXX-XXXXX-XXXXX-XXXXX";
            string clientKey = "XXXXX-XXXXX-XXXXX-XXXXX-XXXXX";
            string officeKey = "XXXXX-XXXXX-XXXXX-XXXXX-XXXXX";

            MI.LogMessageToFile("Validating installation.");
            CimInstance mdtTattooCim = MI.QueryInstance(mdtClass, mdtNS);
            if (mdtTattooCim == null)
            {
                MI.LogMessageToFile("Failed to validate MDT installation, exiting...");
                System.Environment.Exit(1);
            }
            else
            {
                MI.LogMessageToFile("Installation validated.");
            }

            foreach (string arg in args.Distinct())
            {
                if (arg == "windows")
                {
                    MI.LogMessageToFile("Will try to activate Windows...");
                    MI.LogMessageToFile("Determining Windows version...");
                    CimInstance winVersionCim = MI.QueryInstance(osVersion, rootNS);
                    string winVersion = winVersionCim.CimInstanceProperties["Caption"].Value.ToString();

                    if (winVersion.Contains("Microsoft Windows Server 2012 R2"))
                    {
                        MI.LogMessageToFile("Windows Server 2012 family detected.");
                        licenseKey = serverKey;
                    }
                    else if (winVersion.Contains("Microsoft Windows 10"))
                    {
                        MI.LogMessageToFile("Windows 10 detected.");
                        licenseKey = clientKey;
                    }
                    else
                    {
                        MI.LogMessageToFile("Unsupported Windows version detected, skipping...");
                        break;
                    }

                    MI.LogMessageToFile("Checking the current Windows license status.");
                    bool isActivated = MI.ValidateActivation(licenseStatus, rootNS, 3);

                    if (!isActivated)
                    {
                        MI.LogMessageToFile("Activating Windows.");
                        MI.ActivateProduct(licensingService, "windows", licenseKey, licenseStatus, rootNS);

                        MI.LogMessageToFile("Checking the current Windows license status after activation attempt.");
                        MI.ValidateActivation(licenseStatus, rootNS, 0);
                    }
                    else
                    {
                        MI.LogMessageToFile("Windows already activated.");
                    }
                }
                else if (arg == "office")
                {
                    MI.LogMessageToFile("Will try to activate Office...");
                    MI.LogMessageToFile("Checking if Office is installed...");
                    CimInstance officeWmiResult = MI.QueryInstance(officeLicenseStatus, rootNS);

                    if (officeWmiResult != null)
                    {
                        MI.LogMessageToFile("Checking the current Office license status.");
                        bool isActivated = MI.ValidateActivation(officeLicenseStatus, rootNS, 3);

                        if (!isActivated)
                        {
                            MI.LogMessageToFile("Activating Office.");
                            MI.ActivateProduct(licensingService, "office", officeKey, officeLicenseStatus, rootNS);

                            MI.LogMessageToFile("Checking the current Office license status after activation attempt.");
                            MI.ValidateActivation(officeLicenseStatus, rootNS, 0);
                        }
                        else
                        {
                            MI.LogMessageToFile("Office already activated.");
                        }
                    }
                    else
                    {
                        MI.LogMessageToFile("Office is not installed.");
                    }
                }
                else
                {
                    MI.LogMessageToFile("Unknown argument detected, skipping...");
                }
            }
            MI.LogMessageToFile("Windows Activation app finished...");
            MI.LogMessageToFile("=================================");
        }
    }
}

public partial class MI
{
    public static void ActivateProduct(String className, String productName, String productKey, String query, String ns)
    {
        try
        {
            CimMethodParametersCollection productKeyCollection = new CimMethodParametersCollection();
            productKeyCollection.Add(CimMethodParameter.Create("ProductKey", productKey, CimType.String, CimFlags.In));

            CimSession cimSession = CreateCimSession();
            var cimInstance = cimSession.EnumerateInstances(ns, className).First();

            LogMessageToFile("Applying MAK product key...");
            cimSession.InvokeMethod(ns, cimInstance, "InstallProductKey", productKeyCollection);

            LogMessageToFile("Refreshing license status...");
            if (productName == "windows")
            {
            cimSession.InvokeMethod(ns, cimInstance, "RefreshLicenseStatus", null);
            }
            else if (productName == "office")
            {
            var licenseResult = cimSession.QueryInstances(@ns, "WQL", query).First();
            cimSession.InvokeMethod(ns, licenseResult, "Activate", null);
            }
        }
        catch (CimException ex)
        {
            LogMessageToFile("Failed to activate product with the following error:");
            LogMessageToFile(ex.Message);
        }
    }

    private enum LicenseStatusMap
    {
        Unlicensed,
        Licensed,
        OOBGrace,
        OOTGrace,
        NonGenuineGrace,
        Notification,
        ExtendedGrace
    }

     private static bool GetOSActivation(CimInstance cimInstance)
    {
        bool result = false;
        try
        {
            // Find active license by querying for PartialProductKey
            string activeLicenseDescription = cimInstance.CimInstanceProperties["Description"].Value.ToString();
            LogMessageToFile("Current active license is:");
            LogMessageToFile(activeLicenseDescription);

            string licenseStatusString = cimInstance.CimInstanceProperties["LicenseStatus"].Value.ToString();
            LicenseStatusMap licenseStatusMessage = (LicenseStatusMap)Convert.ToInt32(licenseStatusString);
            LogMessageToFile("Current license status is:");
            LogMessageToFile(licenseStatusMessage.ToString());

            // Check license status for MAK
            if (activeLicenseDescription.Contains("VOLUME_MAK"))
            {
                if(licenseStatusString == "1")
                {
                    result = true;
                }
                else
                {
                    result = false;
                }
                return result;
            }
            // Check license status for RETAIL(MAK)
            if (activeLicenseDescription.Contains("RETAIL(MAK)"))
            {
                if(licenseStatusString == "1")
                {
                    result = true;
                }
                else
                {
                    result = false;
                }
                return result;
            }
            // Check license status for trial license
            else if (activeLicenseDescription.Contains("TIMEBASED_EVAL"))
            {
                result = false;
                return result;
            }
            // Check license status for KMS
            else if (activeLicenseDescription.Contains("KMS"))
            {
                if(licenseStatusString == "1")
                {
                    result = true;
                }
                else
                {
                    result = false;
                }
                return result;
            }
            else
            {
                LogMessageToFile("Unsupported license type detected, exiting...");
                System.Environment.Exit(1);
            }
        }
        catch (CimException ex)
        {
            LogMessageToFile(ex.Message);
            LogMessageToFile("Exiting..");
            System.Environment.Exit(1);
        }
        return result;
    }

    private static CimSession CreateCimSession()
    {
        CimSession cimSession = null;
        try
        {
            var sessionOptions = new DComSessionOptions
            {
                Timeout = TimeSpan.FromSeconds(30)
            };

            cimSession = CimSession.Create("localhost", sessionOptions);
        }
        catch (CimException ex)
        {
            LogMessageToFile(ex.Message);
            LogMessageToFile("Exiting...");
            System.Environment.Exit(1);
        }
        return cimSession;
    }

    public static CimInstance QueryInstance(String query, String ns)
    {
        CimInstance result = null;
        try
        {
            CimSession cimSession = CreateCimSession();

            IEnumerable<CimInstance> queryInstances =
              cimSession.QueryInstances(@ns, "WQL", query);
            foreach (CimInstance cimInstance in queryInstances)
            {
                result = cimInstance;
            }
        }
        catch (CimException ex)
        {
            LogMessageToFile(ex.Message);
            LogMessageToFile("Exiting...");
            System.Environment.Exit(1);
        }
        return result;
    }

    public static void DeleteClass(String query)
    {
        try
        {
            CimSession cimSession = CimSession.Create("localhost");
            IEnumerable<CimInstance> queryInstances =
              cimSession.QueryInstances(@"root\cimv2", "WQL", query);

            foreach (CimInstance cimInstance in queryInstances)
            {
                cimSession.DeleteInstance(cimInstance);
                LogMessageToFile(cimInstance.ToString());
                LogMessageToFile("WMI record has been deleted.");
            }
        }
        catch (CimException ex)
        {
            LogMessageToFile(ex.Message);
            LogMessageToFile("Exiting..");
        }
    }

    public static bool ValidateActivation(String query, String ns, Int32 attemptNumber)
    {
        bool result = false;
        try
        {
            CimInstance licenseResult = MI.QueryInstance(query, ns);
            if (licenseResult != null)
            {
                bool activated = GetOSActivation(licenseResult);

                if (activated)
                {
                    LogMessageToFile("Product is activated.");
                    result = true;
                }
                else if (!activated && attemptNumber < 3)
                {
                    LogMessageToFile("Product is not activated, waiting 5 seconds to check one more time...");
                    System.Threading.Thread.Sleep(5000);
                    ValidateActivation(query, ns, ++attemptNumber);
                }
                else
                {
                    LogMessageToFile("Product is not activated.");
                }
            }
            else
            {
                LogMessageToFile("Cannot find product license information.");
                result = false;
            }
        }
        catch (CimException ex)
        {
            LogMessageToFile(ex.Message);
            LogMessageToFile("Exiting...");
        }
        return result;
    }

    private static string GetTempPath()
    {
        string setupPath = Environment.GetEnvironmentVariable("windir") + "\\EngLogs\\Deploy";
        string scriptPath = Environment.GetEnvironmentVariable("windir") + "\\EngLogs\\Scripts";
        string path = null;

        if (Directory.Exists(setupPath))
        {
            path = setupPath; 
        }
        else if (Directory.Exists(scriptPath))
        {
            path = scriptPath;
        }
        else
        {
            Console.WriteLine("Cannot find valid log directory, exiting...");
            System.Environment.Exit(1);
        }

        if (!path.EndsWith("\\")) path += "\\";
        return path;
    }

    public static void LogMessageToFile(String msg)
    {
        System.IO.StreamWriter sw = System.IO.File.AppendText(
            GetTempPath() + "DEPLOY-CT-WINACTIVATION-R1.log");
        try
        {
            string logLine = System.String.Format(
                "{0:G}: {1}", System.DateTime.Now, msg);
            sw.WriteLine(logLine);
        }
        finally
        {
            sw.Close();
        }
    }
}