/**
 * Ashita - Copyright (c) 2014 - 2016 atom0s [atom0s@live.com]
 *
 * This work is licensed under the Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International License.
 * To view a copy of this license, visit http://creativecommons.org/licenses/by-nc-nd/4.0/ or send a letter to
 * Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
 *
 * By using Ashita, you agree to the above license and its terms.
 *
 *      Attribution - You must give appropriate credit, provide a link to the license and indicate if changes were
 *                    made. You must do so in any reasonable manner, but not in any way that suggests the licensor
 *                    endorses you or your use.
 *
 *   Non-Commercial - You may not use the material (Ashita) for commercial purposes.
 *
 *   No-Derivatives - If you remix, transform, or build upon the material (Ashita), you may not distribute the
 *                    modified material. You are, however, allowed to submit the modified works back to the original
 *                    Ashita project in attempt to have it added to the original project.
 *
 * You may not apply legal terms or technological measures that legally restrict others
 * from doing anything the license permits.
 *
 * No warranties are given.
 */

#ifndef __ASHITA_AS_MEMORY_H_INCLUDED__
#define __ASHITA_AS_MEMORY_H_INCLUDED__

#if defined (_MSC_VER) && (_MSC_VER >= 1020)
#pragma once
#endif

#include <Windows.h>
#include <sstream>
#include <vector>

namespace Ashita
{
    class Memory
    {
    public:
        /**
         * Determines the length of data from the given offset valid to scan with.
         *
         * @param {std::vector} pattern - The pattern to find the length within.
         * @param {uintptr_t} offset - The offset to begin iterating at.
         * @returns {uintptr_t} The length of data from the given offset.
         */
        static uintptr_t FindLength(const std::vector<std::pair<uint8_t, bool>>& pattern, uintptr_t offset)
        {
            // Ensure the offset is within the pattern range..
            if (offset >= pattern.size())
                return 0;

            auto len = 0;
            for (size_t x = 0; x < pattern.size(); x++)
            {
                // Compare the data for wildcards or null characters..
                if (pattern[x].second == false ||
                    pattern[x].first == 0x00)
                    break;

                len++;
            }

            return len;
        }

        /**
         * Finds the largest chunk of valid data within the pattern to scan for.
         *
         * @param {std::vector} pattern - The pattern to find the chunk within.
         * @param {int32_t [2]} output - The output data to store the offset and length of data to use.
         */
        static void FindLargestChunk(const std::vector<std::pair<uint8_t, bool>>& pattern, uintptr_t output[2])
        {
            auto offset = (uintptr_t)0;
            auto length = (uintptr_t)FindLength(pattern, 0);
            auto maxlen = (uintptr_t)pattern.size();

            // Loop the pattern and find its largest chunk..
            for (auto x = length; x < maxlen; x++)
            {
                // Ensure the data is required..
                if (pattern[x].second != true)
                    continue;

                // Find the length at the given offset..
                auto count = FindLength(pattern, x);
                if (count > length)
                {
                    offset = x;
                    length = count;
                }

                // Step past the current data..
                if (count > 0)
                    x += count - 1;
            }

            // Return the largest offset and length..
            output[0] = offset;
            output[1] = length;
        }

        /**
         * Finds a pattern within the given range of data.
         *
         * @param {uintptr_t} baseAddress - The address of the data to begin searching within.
         * @param {uintptr_t} baseSize - The size of the data to search within.
         * @param {const char*} pattern - The pattern to search for.
         * @param {uintptr_t} offset - The offset from the found location to add to the pointer.
         * @param {uintptr_t} count - The result count to use if the pattern is found more than once.
         * @returns {uintptr_t} The address where the pattern was located, 0 otherwise.
         */
        static uintptr_t FindPattern(uintptr_t baseAddress, uintptr_t baseSize, const char* pattern, uintptr_t offset, uintptr_t count)
        {
            // Ensure the incoming pattern is properly aligned..
            if (strlen(pattern) % 2 > 0)
                return 0;

            // Convert the pattern to a vector of data..
            std::vector<std::pair<uint8_t, bool>> vpattern;
            for (size_t x = 0, y = strlen(pattern) / 2; x < y; x++)
            {
                // Obtain the current byte..
                std::stringstream stream(std::string(pattern + (x * 2), 2));

                // Check if this is a wildcard..
                if (stream.str() == "??")
                    vpattern.push_back(std::make_pair(00, false));
                else
                {
                    auto byte = strtol(stream.str().c_str(), nullptr, 16);
                    vpattern.push_back(std::make_pair((uint8_t)byte, true));
                }
            }

            // Find the largest chunk of valid data to scan with..
            uintptr_t p[2] = { 0 };
            FindLargestChunk(vpattern, p);

            // Prepare the required variables..
            const auto psize = vpattern.size();
            const auto pstart = p[0];
            const auto plen = (intptr_t)p[1];
            const auto pfirst = vpattern[0].first;

            // Build the wildcard table..
            uint8_t wildcard[UCHAR_MAX + 1] = { 0 };
            for (auto x = pstart; x < pstart + plen; x++)
                wildcard[(uint8_t)vpattern[x].first] = 1;

            const auto data = (uint8_t*)baseAddress;
            const auto size = (uint32_t)baseSize;

            std::vector<uintptr_t> matches;

            // Loop the data and look for the pattern..
            for (intptr_t x = size - psize; x >= 0; x--)
            {
                auto c = data[x];
                auto w = wildcard[c];
                auto k = 0;

                // Step over wildcard characters..
                while (w == 0 && x > (intptr_t)plen)
                {
                    x -= plen;
                    w = wildcard[data[x]];
                    k = 1;
                }

                if (k == 1)
                {
                    x++;
                    continue;
                }

                // Ensure the first byte matches..
                if (c != pfirst) continue;

                // Validate the size..
                if (x - pstart < 0 || x - pstart + psize > size)
                    return 0;

                // Walk the pattern and match its data..
                for (size_t y = 0; y < psize; y++)
                {
                    if (y == pstart || vpattern[y].second != true)
                        continue;
                    if (data[x - pstart + y] != vpattern[y].first)
                        break;
                    if (y + 1 == psize)
                        matches.insert(matches.begin(), (uintptr_t)(data + (x - pstart)) + offset);
                }
            }

            // Handle the count match..
            if (count < matches.size())
                return matches[count];
            return 0;
        }

        /**
         * Obtains the calling module handle for the given address.
         *
         * @param {uintptr_t} returnAddress - The address to locate the module owner of.
         * @returns {HMODULE} The module handle if found, nullptr otherwise.
         */
        static HMODULE __stdcall GetCallingModule(uintptr_t returnAddress)
        {
            if (returnAddress == 0)
                return nullptr;

            MEMORY_BASIC_INFORMATION mbi = { 0 };
            if (::VirtualQuery((LPCVOID)returnAddress, &mbi, sizeof(MEMORY_BASIC_INFORMATION)) == sizeof(MEMORY_BASIC_INFORMATION))
                return (HMODULE)mbi.AllocationBase;

            return nullptr;
        }
    };
}; // namespace Ashita

#endif // __ASHITA_AS_MEMORY_H_INCLUDED__