//-----------------------------------------------------------------------
// 
//  Copyright (C) Microsoft Corporation.  All rights reserved.
// 
// THIS CODE AND INFORMATION ARE PROVIDED AS IS WITHOUT WARRANTY OF ANY
// KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//-----------------------------------------------------------------------

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Text;

namespace CustomActions
{
    [System.Security.Permissions.PermissionSetAttribute(System.Security.Permissions.SecurityAction.Demand, Name = "FullTrust")]
    internal sealed class CaspolSecurityPolicyCreator
    {
        private CaspolSecurityPolicyCreator() {}

        internal static void AddSecurityPolicy(
            bool   machinePolicyLevel,
            string solutionCodeGroupName,
            string solutionCodeGroupDescription,
            string assemblyPath,
            string assemblyCodeGroupName,
            string assemblyCodeGroupDescription)
        {
            string frameworkFolder = GetFrameworkFolder();

            string solutionInstallationLocation = Path.GetDirectoryName(assemblyPath);
            string solutionInstallationUrl = Path.Combine(solutionInstallationLocation, "*");

            string policyLevel;
            string parentCodeGroup;
            if (machinePolicyLevel)
            {
                policyLevel = "-m"; // Use Machine-level policy.
                parentCodeGroup = "My_Computer_Zone"; // Use My_Computer_Zone for assemblies installed on the computer.
            }
            else
            {
                policyLevel = "-u"; // Use User-level policy.
                parentCodeGroup = "All_Code";
            }

            // Add the solution code group. Grant no permission at this level.
            string arguments = policyLevel + " -q -ag " + parentCodeGroup + " -url \"" + solutionInstallationUrl + "\" Nothing -n \"" + solutionCodeGroupName + "\" -d \"" + solutionCodeGroupDescription + "\"";
            try
            {
                RunCaspolCommand(frameworkFolder, arguments);
            }
            catch (Exception ex)
            {
                string error = String.Format("Cannot create the security code group '{0}'.", solutionCodeGroupName);
                throw new Exception(error, ex);
            }

            // Add the assembly code group. Grant FullTrust permissions to the main assembly.
            try
            {
                // Use the assembly strong name as the membership condition.
                // Ensure that the assembly is strong-named to give it full trust.
                AssemblyName assemblyName = Assembly.LoadFile(assemblyPath).GetName();
                arguments = policyLevel + " -q -ag \"" + solutionCodeGroupName + "\" -strong -file \"" + assemblyPath + "\" \"" + assemblyName.Name + "\" \"" + assemblyName.Version.ToString(4) + "\" FullTrust -n \"" + assemblyCodeGroupName + "\" -d \"" + assemblyCodeGroupDescription + "\"";

                RunCaspolCommand(frameworkFolder, arguments);
            }
            catch (Exception ex)
            {
                try
                {
                    // Clean the solutionCodeGroupName.
                    RemoveSecurityPolicy(machinePolicyLevel, solutionCodeGroupName);
                }
                catch {}

                string error = String.Format("Cannot create the security code group '{0}'.", assemblyCodeGroupName);
                throw new Exception(error, ex);
            }
        }

        internal static void RemoveSecurityPolicy(
            bool   machinePolicyLevel,
            string solutionCodeGroupName)
        {
            string frameworkFolder = GetFrameworkFolder();
            string policyLevel;
            if (machinePolicyLevel)
                policyLevel = "-m"; // Use Machine-level policy.
            else
                policyLevel = "-u"; // Use User-level policy.

            string arguments = policyLevel + " -q -rg \"" + solutionCodeGroupName + "\"";
            RunCaspolCommand(frameworkFolder, arguments);
        }

        private static string GetFrameworkFolder()
        {
            // Get the targeted Framework folder.
            Version version = new Version(2, 0, 50727);
            return GetRuntimeInstallationDirectory(version, true);
        }

        private static void RunCaspolCommand(string frameworkFolder, string arguments)
        {
            ProcessStartInfo processStartInfo = new ProcessStartInfo(Path.Combine(frameworkFolder, "caspol.exe"));
            processStartInfo.CreateNoWindow = true;
            processStartInfo.WindowStyle = ProcessWindowStyle.Hidden;
            processStartInfo.WorkingDirectory = frameworkFolder;
            processStartInfo.Arguments = arguments;
            processStartInfo.RedirectStandardError = true;
            processStartInfo.RedirectStandardOutput = true;
            processStartInfo.UseShellExecute = false;

            Process process = Process.Start(processStartInfo);
            string caspolOutputMessage = process.StandardOutput.ReadToEnd();
            process.WaitForExit();

            int exitCode = 0;
            if (process.HasExited)
                exitCode = process.ExitCode;

            if (exitCode != 0)
            {
                string message = null;
                if (!String.IsNullOrEmpty(caspolOutputMessage))
                {
                    String[] outputMessageLines = caspolOutputMessage.Split('\n');
                    for (int i = 2; i < outputMessageLines.Length; i++)
                    {
                        string line = outputMessageLines[i].Trim();
                        if (!String.IsNullOrEmpty(line))
                        {
                            message = line;
                            break;
                        }
                    }
                }
                if (String.IsNullOrEmpty(message))
                    message = "Cannot run the Code Access Security Policy tool (caspol.exe).";

                throw new ApplicationException(message);
            }
        }

        [Flags]
        private enum RuntimeInfo : uint
        {
            UpgradeVersion = 0x1,           // RUNTIME_INFO_UPGRADE_VERSION
            RequestIA64 = 0x2,              // RUNTIME_INFO_REQUEST_IA64
            RequestAmd64 = 0x4,             // RUNTIME_INFO_REQUEST_AMD64
            RequestX86 = 0x8,               // RUNTIME_INFO_REQUEST_X86
            DoNotReturnDirectory = 0x10,    // RUNTIME_INFO_DONT_RETURN_DIRECTORY
            DoNotReturnVersion = 0x20,      // RUNTIME_INFO_DONT_RETURN_VERSION
            DoNotShowErrorDialog = 0x40     // RUNTIME_INFO_DONT_SHOW_ERROR_DIALOG
        }

        [DllImport("mscoree.dll", CharSet = CharSet.Unicode, ExactSpelling = true, PreserveSig = true, SetLastError = false)]
        private static extern int /* [HRESULT] */ GetRequestedRuntimeInfo(
            string /* [LPCWSTR] */ pExe,
            string /* [LPCWSTR] */ pwszVersion,
            string /* [LPCWSTR] */ pConfigurationFile,
            uint /* [DWORD] */ startupFlags,
            RuntimeInfo /* [DWORD] */ runtimeInfoFlags,
            StringBuilder /* [LPWSTR] */ pDirectory,
            uint /* [DWORD] */ dwDirectory,
            out uint /* [DWORD *] */ dwDirectoryLength,
            StringBuilder /* [LPWSTR] */ pVersion,
            uint /* [DWORD] */ cchBuffer,
            out uint /* [DWORD *] */ dwLength
            );

        /// <summary>
        /// Returns the installation directory of the specified .NET runtime.
        /// </summary>
        /// <param name="version">
        /// The version of the runtime.
        /// </param>
        /// <param name="upgradeVersion">
        /// True to return the installation directory of the nearest compatible runtime version, or false for an exact match.
        /// </param>
        /// <returns>
        /// The .NET runtime installation directory.
        /// </returns>
        private static string GetRuntimeInstallationDirectory(Version version, bool upgradeVersion)
        {
            string versionString = "v" + version.ToString(3);
            RuntimeInfo runtimeInfo = RuntimeInfo.DoNotShowErrorDialog;
            if (upgradeVersion)
                runtimeInfo |= RuntimeInfo.UpgradeVersion;

            StringBuilder runtimeDirectory = new StringBuilder(270);
            StringBuilder runtimeVersion = new StringBuilder("v65535.65535.65535".Length);
            uint runtimeDirectoryLength;
            uint runtimeVersionLength;
            int errorCode = GetRequestedRuntimeInfo(null, versionString, null, 0, runtimeInfo, runtimeDirectory, (uint)runtimeDirectory.Capacity, out runtimeDirectoryLength, runtimeVersion, (uint)runtimeVersion.Capacity, out runtimeVersionLength);
            Marshal.ThrowExceptionForHR(errorCode);
            return Path.Combine(runtimeDirectory.ToString(), runtimeVersion.ToString());
        }
    }
}