using System;
using System.Collections;
using System.Collections.Generic;
using System.IO;
using System.Text.RegularExpressions;
using Microsoft.Win32;

namespace SharpDPAPI.Commands
{
    public class Search : ICommand
    {
        public static string CommandName => "search";
        delegate void ProcessFileCallback(string path, bool showErrors);


        private static readonly byte[] dpapiBlobHeader =
        {
            // Version(4 bytes) | DPAPI Proivder Guid(16-bytes - df9d8cd0-1501-11d1-8c7a-00c04fc297eb)
            0x01, 0x00, 0x00, 0x00, 0xD0, 0x8C, 0x9D, 0xDF, 0x01, 0x15, 0xD1, 0x11, 0x8C, 0x7A, 0x00, 0xC0, 0x4F, 0xC2, 0x97, 0xEB
        };

        private readonly byte[][] dpapiBlobSearches =
        {
            dpapiBlobHeader,

            // The following are potential base64 representations of the DPAPI provider GUID
            // Generated by putting dpapiProviderGuid into the script here: https://www.leeholmes.com/blog/2017/09/21/searching-for-content-in-base-64-strings/
            System.Text.Encoding.ASCII.GetBytes("AAAA0Iyd3wEV0RGMegDAT8KX6"),
            System.Text.Encoding.ASCII.GetBytes("AQAAANCMnd8BFdERjHoAwE/Cl+"),
            System.Text.Encoding.ASCII.GetBytes("EAAADQjJ3fARXREYx6AMBPwpfr"),

            // Hex string representation
            System.Text.Encoding.ASCII.GetBytes("01000000D08C9DDF0115D1118C7A00C04FC297EB")
        };

        public void Execute(Dictionary<string, string> arguments)
        {
            Console.WriteLine("\r\n[*] Action: Searching for DPAPI blobs");

            if (!arguments.ContainsKey("/type"))
                arguments["/type"] = "registry";

            Console.WriteLine($"[*] Search type: {arguments["/type"]}");

            switch (arguments["/type"])
            {
                case "registry":
                    SearchRegistry(arguments);
                    break;

                case "folder":
                    SearchFolder(arguments);
                    break;

                case "file":
                    SearchFile(arguments);
                    break;

                case "base64":
                    SearchBase64(arguments);
                    break;

                default:
                    throw new ArgumentException($"Unknown /type parameter '{arguments["/type"]}'");
            }
        }

        private void SearchFolder(Dictionary<string, string> arguments)
        {
            if (!arguments.ContainsKey("/path"))
                throw new ArgumentException("/path argument not specified");

            var path = arguments["/path"];
            var showErrors = arguments.ContainsKey("/showErrors");

            if (!Directory.Exists(path))
                throw new ArgumentException($"The folder '{path}' does not exist");

            uint maxBytes = 1024;
            if (arguments.ContainsKey("/maxBytes") && !uint.TryParse(arguments["/maxBytes"], out maxBytes))
                throw new ArgumentException($"Invalid uint value '{arguments["/maxBytes"]}' in the /maxBytes argument");

            Console.WriteLine($"[*] Searching for the folder {path} for files potentially containing DPAPI blobs in the first {maxBytes} bytes\n");

            FindFiles(path, maxBytes, showErrors, (filePath, displayErrors) =>
            {
                try
                {
                    if (FileContainsDpapiBlob(filePath, maxBytes))
                        Console.WriteLine(filePath);
                }
                catch (Exception e)
                {
                    if(displayErrors)
                        Console.WriteLine($"File Processing ERROR: {filePath} - {e.Message}");
                }
            });
        }


        private void FindFiles(string path, uint maxBytes, bool showErrors, ProcessFileCallback processFile)
        {
            // Modified largely from https://developerslogblog.wordpress.com/2020/02/25/c-how-to-find-all-files-recursively-in-a-folder/
            var paths = new List<string>();
            var directoriesQueue = new Queue();
            directoriesQueue.Enqueue(path);

            while (directoriesQueue.Count > 0)
            {
                var currentPath = (string)directoriesQueue.Dequeue();

                try
                {
                    //Console.WriteLine("Processing folder " + currentPath);
                    var directories = Directory.GetDirectories(currentPath);
                    foreach (var directory in directories)
                        directoriesQueue.Enqueue(directory);

                    foreach (var file in Directory.GetFiles(currentPath))
                        processFile(file, showErrors);
                }
                catch(Exception e)
                {
                    if(showErrors)
                        Console.WriteLine($"ERROR: {e}");
                }
            }
        }


        private void SearchFile(Dictionary<string, string> arguments)
        {
            if(!arguments.ContainsKey("/path"))
                throw new ArgumentException("/path argument not specified");

            var path = arguments["/path"];

            if(!File.Exists(path))
                throw new ArgumentException($"The file '{path}' does not exist");

            uint maxBytes = 1024;
            if (arguments.ContainsKey("/maxBytes") && !uint.TryParse(arguments["/maxBytes"], out maxBytes))
                throw new ArgumentException($"Invalid uint value '{arguments["/maxBytes"]}' in the /maxBytes argument");

            Console.WriteLine($"[*] Searching for DPAPI blobs in the file {path}\n");
            if (FileContainsDpapiBlob(path, maxBytes))
                Console.WriteLine($"Found potential DPAPI blob at {path}");
            else
                Console.WriteLine("No DPAPI blob found");
        }

        private bool FileContainsDpapiBlob(string path, uint bytesToSearch)
        {
            var fileContents = new byte[bytesToSearch];
            using (var file = new FileStream(path, FileMode.Open))
            {
                file.Read(fileContents, 0, (int) bytesToSearch);
            }

            return ContainsDpapiBlob(fileContents);
        }

        private void SearchBase64(Dictionary<string, string> arguments)
        {
            if (!arguments.ContainsKey("/base64"))
                throw new ArgumentException("/base64 argument not found");

            ContainsDpapiBlob(Convert.FromBase64String(arguments["/base64"]));
        }

        private void SearchRegistry(Dictionary<string, string> arguments)
        {
            var showErrors = arguments.ContainsKey("/showErrors");


            if (arguments.ContainsKey("/path"))
            {
                string keyPath;
                RegistryHive hive;
                // Based on Seatbelt's code here: https://github.com/GhostPack/Seatbelt/blob/master/Seatbelt/Commands/Windows/RegistryValueCommand.cs#L43-L60
                var path = arguments["/path"];
                var separatorPos = path.IndexOf("\\");

                if (separatorPos == -1)                   // e.g. HKLM
                {
                    hive = GetHive(path);
                    keyPath = "\\";
                }
                else if (separatorPos == path.Length)   // e.g. HKLM\
                {
                    var hiveStr = path.Substring(0, separatorPos);
                    hive = GetHive(hiveStr);
                    keyPath = "\\";
                }
                else                                       // e.g. HKLM\Software
                {
                    var hiveStr = path.Substring(0, separatorPos);
                    hive = GetHive(hiveStr);
                    keyPath = path.Substring(separatorPos + 1);
                }

                Console.WriteLine($"[*] Searching the key '{path}' for DPAPI blobs. Hive: {hive} Path: {keyPath}{(showErrors ? " (Displaying all errors)" : "")}\n");

                var root = RegistryKey.OpenRemoteBaseKey(hive, "").OpenSubKey(keyPath, RegistryKeyPermissionCheck.ReadSubTree);
                foreach (var match in FindRegistryBlobs(root, showErrors))
                {
                    Console.WriteLine(match);
                }
            }
            else
            {
                Console.WriteLine($"[*] Searching the key HLKM and HKCU hives for DPAPI blobs\n");

                var matchingKeys = new List<string>();
                Console.WriteLine("[*] Searching USERS hive:\n");
                foreach (var match in FindRegistryBlobs(Registry.Users.OpenSubKey("\\", RegistryKeyPermissionCheck.ReadSubTree), showErrors))
                {
                    Console.WriteLine(match);
                }

                matchingKeys.Clear();
                Console.WriteLine("\n\n[*] Searching HKLM hive:\n");
                
                foreach (var match in FindRegistryBlobs(Registry.LocalMachine.OpenSubKey("\\", RegistryKeyPermissionCheck.ReadSubTree), showErrors))
                {
                    Console.WriteLine(match);
                }
            }
        }

        // From Seatbelt: https://github.com/GhostPack/Seatbelt/blob/master/Seatbelt/Util/RegistryUtil.cs#L440
        public static RegistryHive GetHive(string name)
        {
            switch (name.ToUpper())
            {
                case "HKCR":
                case "HKEY_CLASSES_ROOT":
                    return RegistryHive.ClassesRoot;

                case "HKEY_CURRENT_CONFIG":
                    return RegistryHive.CurrentConfig;

                case "HKCU":
                case "HKEY_CURRENT_USER":
                    return RegistryHive.CurrentUser;

                case "HKLM":
                case "HKEY_LOCAL_MACHINE":
                    return RegistryHive.LocalMachine;

                case "HKEY_PERFORMANCE_DATA":
                    return RegistryHive.PerformanceData;

                case "HKU":
                case "HKEY_USERS":
                    return RegistryHive.Users;

                default:
                    throw new Exception("UnknownRegistryHive");
            }
        }

        private IEnumerable<string> FindRegistryBlobs(RegistryKey root, bool reportErrors)
        {
            var toCheck = new LinkedList<RegistryKey>();
            toCheck.AddLast(root);
            Console.WriteLine("Root: " + root);

            while (toCheck.Count > 0)
            {
                root = toCheck.First.Value;
                toCheck.RemoveFirst();

                if (root == null)
                    continue;

                var valueNames = new string[] { };

                try
                {
                    valueNames = root.GetValueNames();
                }
                catch
                {
                    if (reportErrors)
                        Console.WriteLine($"ERROR: Could not list values for {root}");
                }

                foreach (var name in valueNames)
                {
                    switch (root.GetValueKind(name))
                    {
                        case RegistryValueKind.String:
                            string matchingPath = null;
                            try
                            {
                                var str = (string)root.GetValue(name);

                                // Regex generated by putting dpapiProviderGuid into the script here: https://www.leeholmes.com/blog/2017/09/21/searching-for-content-in-base-64-strings/
                                if (Regex.IsMatch(str, "(AAAA0Iyd3wEV0RGMegDAT8KX6|AQAAANCMnd8BFdERjHoAwE/Cl+|EAAADQjJ3fARXREYx6AMBPwpfr|01000000D08C9DDF0115D1118C7A00C04FC297EB)"))
                                {
                                    var bytes = Convert.FromBase64String(str);

                                    if (ContainsDpapiBlob(bytes))
                                        matchingPath = $"{root.Name}{name}";
                                }
                            }
                            catch
                            {
                                // Not base64 content
                            }
                            if(matchingPath != null)
                                yield return matchingPath;
                            break;

                        case RegistryValueKind.Binary:
                            if(ContainsDpapiBlob((byte[])root.GetValue(name)))
                            {
                                yield return $"{root.Name}{name}";
                            }
                            break;
                    }
                }


                var subkeyNames = new string[] { };

                try
                {
                    subkeyNames = root.GetSubKeyNames();
                }
                catch
                {
                    if (reportErrors)
                        Console.WriteLine($"ERROR: Could not list subkeys of {root}");
                }
                foreach (var sub in subkeyNames)
                {
                    try
                    {
                        var subkey = root.OpenSubKey(sub);
                        toCheck.AddLast(subkey);
                    }
                    catch(Exception e)
                    {
                        if (reportErrors)
                            Console.WriteLine($"ERROR: Failed to open {root}\\{sub}. Message: {e.Message}");
                    }
                }
            }
        }

        private bool ContainsDpapiBlob(byte[] bytes)
        {
            //return bytes.Contains(dpapiProviderGuid);
            foreach (var searchBytes in dpapiBlobSearches)
            {
                if (bytes.Contains(searchBytes))
                    return true;
            }

            return false;
        }
    }
}