/************************************************************************************************
* All code in this file is under the MS-RL License (https://opensource.org/licenses/MS-RL)      *
* By using the code in this file in any way, you agree to the above license terms.              *
* Copyright (C) LIGHTS IN LINE AB (https://www.lightsinline.se)                                 *
* Repository, Wiki, Issue tracker and more at https://git.lightsinline.se/products/VSTT-Plugins *
*                                                                                               *
* Contributors                                                                                  *
* LIGHTS IN LINE AB                                                                             *
* SWEDBANK AB                                                                                   *
* SKATTEVERKET                                                                                  *
************************************************************************************************/
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.VisualStudio.TestTools.WebTesting;
using System.ComponentModel;
using Microsoft.VisualStudio.TestTools.WebTesting.Rules;
using System.Text.RegularExpressions;
using System.IO;
using Microsoft.VisualStudio.TestTools.LoadTesting;

namespace LIL_VSTT_Plugins
{
    [DisplayName("Set Request Think Time"), Description("Changes the thinktime on requests with a set thinktime over 0 to the value of the ThinkTime context parameter")]
    public class SetRequestThinkTime : WebTestPlugin
    {
        [DisplayName("Debug"), DefaultValue(false), Description("Debug logging when thinktime is set on requests")]
        public bool DebugMode { get; set; }
        public override void PreRequest(object sender, PreRequestEventArgs e)
        {
            WebTestContext ctx = e.WebTest.Context;

            if (ctx.ContainsKey("ThinkTime"))
            {
                int tt = Int32.Parse(ctx["ThinkTime"].ToString());
                if (e.Request.ThinkTime > 0)
                {
                    if (DebugMode) e.WebTest.AddCommentToResult("Setting Think Time to " + tt);
                    e.Request.ThinkTime = tt;
                }
            }
        }
    }

    [DisplayName("Clear Cookies"), Description("Clears all cookies from previous iterations of this webtest")]
    public class ClearCookies : WebTestPlugin
    {
        public override void PreWebTest(object sender, PreWebTestEventArgs e)
        {
            e.WebTest.Context.CookieContainer = new System.Net.CookieContainer();
        }
    }

    [DisplayName("Think Time Emulator 10/190"), Description("Sets a context parameter named ThinkTime in each starting test to a random value between 10%-190% of the specified value.")]
    public class ThinkTimeEmulator10190 : ILoadTestPlugin
    {
        [DisplayName("Think Time"), DefaultValue(0), Description("The Think Time to be used seconds. Default is 0.")]
        public int ThinkTime { get; set; }

        //store the load test object.  
        LoadTest mLoadTest;
        Random rnd = new Random();

        public void Initialize(LoadTest loadTest)
        {
            mLoadTest = loadTest;

            //connect to the TestStarting event.
            mLoadTest.TestStarting += new EventHandler<TestStartingEventArgs>(mLoadTest_TestStarting);
        }


        void mLoadTest_TestStarting(object sender, TestStartingEventArgs e)
        {
            // Set the think time parameter in the tests context to a new value
            int min = ThinkTime/10;
            int max = ThinkTime*2 - min;
            e.TestContextProperties.Add("ThinkTime", rnd.Next(min, max));
        }
    }

    [DisplayName("Set Cookie from Query String"), Description("Sets a cookie with given name from the value of a given Query String parameter if found in the current request")]
    public class SetCookieFromQueryString : WebTestPlugin
    {
        [DisplayName("Cookie name"), Description("Name of the cookie to set")]
        public String CookieName { get; set; }

        [DisplayName("Query String parameter name"), Description("Name of the query string parameter to look for")]
        public String ParamName { get; set; }

        public override void PreRequest(object sender, PreRequestEventArgs e)
        {
            base.PreRequest(sender, e);
            QueryStringParameterCollection col = e.Request.QueryStringParameters;
            for(int x=0; x < col.Count; x++)
            {
                if (col[x].Name == ParamName)
                {
                    e.Request.Cookies.Add(new System.Net.Cookie(CookieName, col[x].Value));
                    return;
                }
            }
        }

        public override void PostRequest(object sender, PostRequestEventArgs e)
        {
            base.PostRequest(sender, e);
            if (e.Request.HasDependentRequests)
            {
                foreach (WebTestRequest item in e.Request.DependentRequests)
                {
                    QueryStringParameterCollection col = item.QueryStringParameters;
                    for (int x = 0; x < col.Count; x++)
                    {
                        if (col[x].Name == ParamName)
                        {
                            item.Cookies.Add(new System.Net.Cookie(CookieName, col[x].Value));
                        }
                    }
                }
            }
        }
    }

    [DisplayName("Add Header"), Description("Adds the specified header to all requests matching a given URL regex")]
    public class AddHeader : WebTestPlugin
    {
        [DisplayName("Header Name"), Description("Name of the header to set")]
        public String HeaderName { get; set; }

        [DisplayName("Header Value"), Description("Value of the header to set")]
        public String HeaderValue { get; set; }

        [DisplayName("URL RegEx"), Description("Regular Expression to match on the URL. Empty matches all requests.")]
        public String RegEx { get; set; }

        public override void PreRequest(object sender, PreRequestEventArgs e)
        {
            base.PreRequest(sender, e);
            if (Regex.Match(e.Request.Url, RegEx).Success || RegEx.Length == 0)
            {
                e.Request.Headers.Add(HeaderName, HeaderValue);
            }
        }

        public override void PostRequest(object sender, PostRequestEventArgs e)
        {
            base.PostRequest(sender, e);
            if (e.Request.HasDependentRequests)
            {
                foreach (WebTestRequest item in e.Request.DependentRequests)
                {
                    if (Regex.Match(item.Url, RegEx).Success || RegEx.Length == 0)
                    {
                        item.Headers.Add(HeaderName, HeaderValue);
                    }
                }
            }
        }
    }

    [DisplayName("Log Context Arrays to File"), Description("Logs the specified context parameter arrays to file after each test iteration, one row for each array value")]
    public class LogContextArrayToFile : WebTestPlugin
    {
        [DisplayName("Parameter arrays"), Description("Comma separated list of array parameters to log")]
        public String Params { get; set; }

        [DisplayName("File Name"), Description("The file name to use for logging")]
        public String Filename { get; set; }

        private bool header = true;
        private string[] paramlist = null;

        public override void PostWebTest(object sender, PostWebTestEventArgs e)
        {
            if (Params.Length == 0 || Filename.Length == 0) return;
            if (paramlist == null) paramlist = Params.Split(',');
            if (header)
            {
                File.AppendAllText(Filename, Params + "\r\n");
                header = false;
            }
            // Check that the first param array has a count
            int count = 0;
            WebTestContext ctx = e.WebTest.Context;
            if (ctx.ContainsKey(paramlist[0] + "_count")) count = Int32.Parse(ctx[paramlist[0] + "_count"].ToString());
            for (int i = 1; i <= count; i++)
            {
                string row = "";
                foreach (string param in paramlist)
                {
                    if (row.Length > 0) row += ",";
                    if (ctx.ContainsKey(param + "_" + i)) row += ctx[param + "_" + i].ToString();
                }
                File.AppendAllText(Filename, row + "\r\n");
            }
            base.PostWebTest(sender, e);
        }
    }

    [DisplayName("Log Context to File"), Description("Logs the specified context parameters to file after each test iteration")]
    public class LogContextToFile : WebTestPlugin
    {
        [DisplayName("Regular Expression"), Description("The RegEx to use to match context parameter names")]
        public String myRegex { get; set; }

        [DisplayName("File Name"), Description("The file name to use for logging")]
        public String Filename { get; set; }

        [DisplayName("Write header"), DefaultValue(false), Description("Writes the parameter names as a header. Will write a new header for each user (dont use in loadtest)")]
        public bool useHeader { get; set; }

        [DisplayName("Log if passed"), DefaultValue(true), Description("Logs the parameters if the webtest passed")]
        public bool logPassed { get; set; }

        [DisplayName("Log if failed"), DefaultValue(true), Description("Logs the parameters if the webtest failed")]
        public bool logFailed { get; set; }

        private bool header = true;

        public override void PostWebTest(object sender, PostWebTestEventArgs e)
        {
            if ((e.WebTest.Outcome == Outcome.Pass && logPassed) || (e.WebTest.Outcome == Outcome.Fail && logFailed)) {
                string row = "";
                string hrow = "";

                foreach (KeyValuePair<string, object> pair in e.WebTest.Context)
                {
                    if (Regex.Match(pair.Key, myRegex).Success)
                    {
                        if (header && useHeader)
                        {
                            if (hrow.Length == 0) hrow = pair.Key.ToString();
                            else hrow += "," + pair.Key.ToString();
                        }
                        if (row.Length == 0) row = pair.Value.ToString();
                        else row += "," + pair.Value.ToString();
                    }
                }
                if (header && useHeader)
                {
                    File.AppendAllText(Filename + e.WebTest.Context.WebTestUserId + ".log", hrow + "\r\n");
                    header = false;
                }
                File.AppendAllText(Filename + e.WebTest.Context.WebTestUserId + ".log", row + "\r\n");
                base.PostWebTest(sender, e);
            }
        }
    }

    [DisplayName("Multi Regular Expression"), Description("Saves all matches for a regular expression including a count")]
    public class MultiRegExExtract : ExtractionRule
    {
        [DisplayName("Regular Expression"), Description("The RegEx to use, supports grouping, for example (.+?)")]
        public String Regex { get; set; }

        [DisplayName("Group Name"), DefaultValue("1"), Description("The group to use as the value for the parameter, default is 1")]
        public String GroupName { get; set; }

        [DisplayName("Pass if not found"), DefaultValue(false), Description("Wherever this extration rule should pass even if no matches are found. Default false.")]
        public bool PassIfNotFound { get; set; }

        public override void Extract(object sender, ExtractionEventArgs e)
        {
            MatchCollection matches = System.Text.RegularExpressions.Regex.Matches(
                e.Response.BodyString,
                Regex,
                RegexOptions.IgnoreCase);
            int index = 0;
            foreach (Match match in matches)
            {
                index++;
                e.WebTest.Context[ContextParameterName + "_" + index] = match.Groups[GroupName].Value;
            }
            if (index > 0)
            {
                e.WebTest.Context[ContextParameterName + "_count"] = index.ToString();
                e.Success = true;
                e.Message = "Found " + index + " matches";
            }
            else
            {
                e.WebTest.Context[ContextParameterName + "_count"] = index.ToString();
                e.Success = PassIfNotFound;
                e.Message = "No matches found. Body size is " + e.Response.BodyString.Length;
            }
        }
    }

    [DisplayName("Extract Count"), Description("Counts all matches for a given RegEx and saves the count to the specified parameter")]
    public class ExtractCount : ExtractionRule
    {
        [DisplayName("Regular Expression"), Description("The RegEx to use")]
        public String Regex { get; set; }

        [DisplayName("Fail on zero count"), DefaultValue(false), Description("Wherever this extration rule should fail if no matches are found")]
        public bool FailOnZero { get; set; }

        public override void Extract(object sender, ExtractionEventArgs e)
        {
            MatchCollection matches = System.Text.RegularExpressions.Regex.Matches(
                e.Response.BodyString,
                Regex,
                RegexOptions.IgnoreCase);
            
            e.WebTest.Context[ContextParameterName] = matches.Count.ToString();
            e.Message = "Found " + matches.Count + " matches";
            if (FailOnZero && matches.Count == 0)
            {
                e.Success = false;
            }
        }
    }

    [DisplayName("Regular Expression Loop"), Description("Loop Condition that matches once on each regexp in previous response body")]
    public class RegExpLoop : ConditionalRule
    {
        private Int32 CurrentMatch { get; set; }
        private String LastUrl { get; set; }
        private MatchCollection Matches { get; set; }

        [DisplayName("Context parameter"), IsContextParameterName(true),
        Description("Name of context parameter where the current value should be set in each loop")]
        public string ContextParameterName { get; set; }

        [DisplayName("Regex"), Description("The RegEx to use, supports grouping, for example (.+?)")]
        public String Regex { get; set; }

        [DisplayName("GroupName"), Description("The group to use as the value for the parameter if several are specified, default is 1")]
        public String GroupName { get; set; }

        // Methods
        public override void CheckCondition(object sender, ConditionalEventArgs e)
        {
            if (CurrentMatch < Matches.Count)
            {
                e.WebTest.Context[ContextParameterName] =
                Matches[CurrentMatch].Groups[GroupName].Value;
                e.IsMet = true;
                CurrentMatch++;
                return;
            }

            e.IsMet = false;
        }

        public override void Initialize(object sender, ConditionalEventArgs e)
        {
            CurrentMatch = 0;
            Matches = System.Text.RegularExpressions.Regex.Matches(
            e.WebTest.LastResponse.BodyString,
            Regex,
            RegexOptions.IgnoreCase);
        }

        public override string StringRepresentation()
        {
            return "Regex condition " + Regex;
        }
    }
}