using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;
using System.Management;
using System.Net;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Text;
using System.Text.RegularExpressions;
using Microsoft.Win32;

namespace ISphereWinHelper
{
    internal static class LiveProbeInventory
    {
        private static readonly string[] InterestingTerms = new[]
        {
            "isphere", "implatform", "impp", "importal", "smark", "smack", "xmpp", "msg", "loginlog", "packetreader", "sendreceive"
        };

        private static readonly string[] InterestingExtensions = new[]
        {
            ".exe", ".dll", ".config", ".xml", ".ini", ".json", ".log", ".db", ".sqlite", ".dat", ".idx", ".ldb", ".txt"
        };

        public static Dictionary<string, object> CollectProcessDetails(Dictionary<string, object> runtime)
        {
            var pids = TargetPids(runtime);
            var details = new List<Dictionary<string, object>>();
            foreach (int pid in pids)
            {
                string query = "SELECT ProcessId,ParentProcessId,Name,ExecutablePath,CommandLine,CreationDate FROM Win32_Process WHERE ProcessId=" + pid.ToString(CultureInfo.InvariantCulture);
                foreach (ManagementObject obj in QueryWmi(query))
                {
                    using (obj)
                    {
                        details.Add(new Dictionary<string, object>
                        {
                            { "pid", SafeWmi(obj, "ProcessId") },
                            { "parent_pid", SafeWmi(obj, "ParentProcessId") },
                            { "name", SafeWmi(obj, "Name") },
                            { "executable_path", SafeWmi(obj, "ExecutablePath") },
                            { "command_line_redacted", RedactSecrets(SafeWmi(obj, "CommandLine")) },
                            { "creation_date", SafeWmi(obj, "CreationDate") }
                        });
                    }
                }
            }
            return new Dictionary<string, object>
            {
                { "target_pid_count", pids.Count },
                { "details", details }
            };
        }

        public static Dictionary<string, object> CollectNetwork(Dictionary<string, object> runtime)
        {
            var targetPids = TargetPids(runtime);
            return new Dictionary<string, object>
            {
                { "tcp", CaptureTargetTcpConnections(targetPids) },
                { "named_pipes", CaptureNamedPipes() }
            };
        }

        public static Dictionary<string, object> CollectServicesRegistryShortcuts()
        {
            return new Dictionary<string, object>
            {
                { "services", CaptureServices() },
                { "registry_uninstall", CaptureUninstallEntries() },
                { "shortcuts", CaptureShortcuts() }
            };
        }

        public static Dictionary<string, object> CollectFilesystem(Dictionary<string, object> runtime)
        {
            var roots = CandidateRoots(runtime);
            var entries = new List<Dictionary<string, object>>();
            foreach (string root in roots)
            {
                entries.Add(InventoryRoot(root));
            }
            return new Dictionary<string, object>
            {
                { "root_count", entries.Count },
                { "roots", entries }
            };
        }

        public static Dictionary<string, object> SummarizeFeasibility(Dictionary<string, object> runtime, Dictionary<string, object> filesystem, Dictionary<string, object> network)
        {
            var targetPids = TargetPids(runtime);
            bool hasTarget = targetPids.Count > 0;
            bool hasSmack = AnyAssemblyPresent(runtime, "smack.dll");
            bool hasInterface = AnyAssemblyPresent(runtime, "IMPP.Interface.dll");
            bool hasService = AnyAssemblyPresent(runtime, "IMPP.Service.dll") || AnyAssemblyPresent(runtime, "IMPP.ServiceBase.dll");
            bool hasTcpFileTransfer = AnyAssemblyPresent(runtime, "TcpFileTransfer.dll");
            string route = hasTarget && hasSmack && hasInterface ? "B_ROUTE_PROMISING" : "NEEDS_REVIEW_OR_A_ROUTE";

            return new Dictionary<string, object>
            {
                { "target_process_seen", hasTarget },
                { "target_pid_count", targetPids.Count },
                { "has_smack_loaded", hasSmack },
                { "has_impp_interface_loaded", hasInterface },
                { "has_impp_service_loaded", hasService },
                { "has_tcp_file_transfer_loaded", hasTcpFileTransfer },
                { "initial_route_hint", route },
                { "next_decision", hasTarget ? "review loaded modules, bitness, install files, and UIA target windows" : "run again after manual login in the working environment" }
            };
        }

        private static Dictionary<string, object> InventoryRoot(string root)
        {
            var files = new List<Dictionary<string, object>>();
            var directories = new List<string>();
            if (!Directory.Exists(root))
            {
                return new Dictionary<string, object>
                {
                    { "root", root },
                    { "exists", false },
                    { "directories_sample", directories },
                    { "files", files }
                };
            }

            foreach (string dir in SafeDirectories(root, 3, 120))
            {
                directories.Add(dir);
            }

            foreach (string file in SafeFiles(root, 4, 1000))
            {
                if (IsInterestingFile(file))
                {
                    files.Add(DescribeFile(file));
                }
            }

            files.Sort(delegate(Dictionary<string, object> left, Dictionary<string, object> right)
            {
                return string.Compare(Convert.ToString(left["path"]), Convert.ToString(right["path"]), StringComparison.OrdinalIgnoreCase);
            });

            return new Dictionary<string, object>
            {
                { "root", root },
                { "exists", true },
                { "directories_sample", directories },
                { "interesting_file_count", files.Count },
                { "files", files }
            };
        }

        private static Dictionary<string, object> DescribeFile(string path)
        {
            var info = new FileInfo(path);
            string version = "";
            string productVersion = "";
            try
            {
                FileVersionInfo vi = FileVersionInfo.GetVersionInfo(path);
                version = vi.FileVersion ?? "";
                productVersion = vi.ProductVersion ?? "";
            }
            catch
            {
            }

            bool keyBinary = IsKeyBinary(info.Name);
            bool smallText = IsSmallTextFile(info);
            return new Dictionary<string, object>
            {
                { "path", info.FullName },
                { "name", info.Name },
                { "extension", info.Extension },
                { "length", info.Length },
                { "creation_time_utc", info.CreationTimeUtc.ToString("o") },
                { "last_write_time_utc", info.LastWriteTimeUtc.ToString("o") },
                { "file_version", version },
                { "product_version", productVersion },
                { "sha256", keyBinary || info.Length <= 8 * 1024 * 1024 ? Sha256(info.FullName) : "" },
                { "text_sample_redacted", smallText ? RedactedTextSample(info.FullName, 4096) : "" }
            };
        }

        private static SortedSet<string> CandidateRoots(Dictionary<string, object> runtime)
        {
            var roots = new SortedSet<string>(StringComparer.OrdinalIgnoreCase);
            foreach (Dictionary<string, object> target in GetDictionaryList(runtime, "targets"))
            {
                string exe = GetString(target, "executable_path");
                if (!string.IsNullOrWhiteSpace(exe))
                {
                    AddRoot(roots, Path.GetDirectoryName(exe));
                }
            }

            AddRoot(roots, Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), "Temp", "importal"));
            AddRoot(roots, Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), "importal"));
            AddRoot(roots, Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.ApplicationData), "importal"));
            AddRoot(roots, Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.CommonApplicationData), "importal"));
            AddRoot(roots, Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.MyDocuments), "iSphere"));
            AddRoot(roots, Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.MyDocuments), "IMPP"));
            return roots;
        }

        private static void AddRoot(SortedSet<string> roots, string path)
        {
            if (string.IsNullOrWhiteSpace(path))
            {
                return;
            }
            try
            {
                if (Directory.Exists(path))
                {
                    roots.Add(Path.GetFullPath(path));
                }
            }
            catch
            {
            }
        }

        private static List<Dictionary<string, object>> CaptureServices()
        {
            var services = new List<Dictionary<string, object>>();
            foreach (ManagementObject obj in QueryWmi("SELECT Name,DisplayName,State,StartMode,PathName,ProcessId FROM Win32_Service"))
            {
                using (obj)
                {
                    string combined = (SafeWmi(obj, "Name") + " " + SafeWmi(obj, "DisplayName") + " " + SafeWmi(obj, "PathName")).ToLowerInvariant();
                    if (!ContainsInterestingTerm(combined))
                    {
                        continue;
                    }
                    services.Add(new Dictionary<string, object>
                    {
                        { "name", SafeWmi(obj, "Name") },
                        { "display_name", SafeWmi(obj, "DisplayName") },
                        { "state", SafeWmi(obj, "State") },
                        { "start_mode", SafeWmi(obj, "StartMode") },
                        { "path_name_redacted", RedactSecrets(SafeWmi(obj, "PathName")) },
                        { "process_id", SafeWmi(obj, "ProcessId") }
                    });
                }
            }
            return services;
        }

        private static List<Dictionary<string, object>> CaptureUninstallEntries()
        {
            var entries = new List<Dictionary<string, object>>();
            CaptureUninstallHive(entries, RegistryHive.LocalMachine, RegistryView.Registry64);
            CaptureUninstallHive(entries, RegistryHive.LocalMachine, RegistryView.Registry32);
            CaptureUninstallHive(entries, RegistryHive.CurrentUser, RegistryView.Registry64);
            CaptureUninstallHive(entries, RegistryHive.CurrentUser, RegistryView.Registry32);
            return entries;
        }

        private static void CaptureUninstallHive(List<Dictionary<string, object>> entries, RegistryHive hive, RegistryView view)
        {
            try
            {
                using (RegistryKey baseKey = RegistryKey.OpenBaseKey(hive, view))
                using (RegistryKey uninstall = baseKey.OpenSubKey(@"SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall"))
                {
                    if (uninstall == null)
                    {
                        return;
                    }
                    foreach (string subName in uninstall.GetSubKeyNames())
                    {
                        using (RegistryKey sub = uninstall.OpenSubKey(subName))
                        {
                            if (sub == null)
                            {
                                continue;
                            }
                            string displayName = Convert.ToString(sub.GetValue("DisplayName", ""));
                            string installLocation = Convert.ToString(sub.GetValue("InstallLocation", ""));
                            string publisher = Convert.ToString(sub.GetValue("Publisher", ""));
                            string combined = (subName + " " + displayName + " " + installLocation + " " + publisher).ToLowerInvariant();
                            if (!ContainsInterestingTerm(combined))
                            {
                                continue;
                            }
                            entries.Add(new Dictionary<string, object>
                            {
                                { "hive", hive.ToString() },
                                { "view", view.ToString() },
                                { "key", subName },
                                { "display_name", displayName },
                                { "display_version", Convert.ToString(sub.GetValue("DisplayVersion", "")) },
                                { "publisher", publisher },
                                { "install_location", installLocation },
                                { "uninstall_string_redacted", RedactSecrets(Convert.ToString(sub.GetValue("UninstallString", ""))) }
                            });
                        }
                    }
                }
            }
            catch
            {
            }
        }

        private static List<Dictionary<string, object>> CaptureShortcuts()
        {
            var shortcuts = new List<Dictionary<string, object>>();
            string[] roots = new[]
            {
                Environment.GetFolderPath(Environment.SpecialFolder.DesktopDirectory),
                Environment.GetFolderPath(Environment.SpecialFolder.CommonDesktopDirectory),
                Environment.GetFolderPath(Environment.SpecialFolder.StartMenu),
                Environment.GetFolderPath(Environment.SpecialFolder.CommonStartMenu)
            };
            foreach (string root in roots)
            {
                if (string.IsNullOrWhiteSpace(root) || !Directory.Exists(root))
                {
                    continue;
                }
                foreach (string lnk in SafeFiles(root, 5, 500))
                {
                    if (!lnk.EndsWith(".lnk", StringComparison.OrdinalIgnoreCase))
                    {
                        continue;
                    }
                    string lower = lnk.ToLowerInvariant();
                    if (!ContainsInterestingTerm(lower))
                    {
                        continue;
                    }
                    FileInfo info = new FileInfo(lnk);
                    shortcuts.Add(new Dictionary<string, object>
                    {
                        { "path", info.FullName },
                        { "length", info.Length },
                        { "last_write_time_utc", info.LastWriteTimeUtc.ToString("o") },
                        { "sha256", info.Length <= 1024 * 1024 ? Sha256(info.FullName) : "" }
                    });
                }
            }
            return shortcuts;
        }

        private static List<Dictionary<string, object>> CaptureNamedPipes()
        {
            var pipes = new List<Dictionary<string, object>>();
            try
            {
                foreach (string pipe in Directory.GetFiles(@"\\.\pipe\"))
                {
                    string lower = pipe.ToLowerInvariant();
                    if (!ContainsInterestingTerm(lower))
                    {
                        continue;
                    }
                    pipes.Add(new Dictionary<string, object> { { "path", pipe } });
                }
            }
            catch (Exception ex)
            {
                pipes.Add(new Dictionary<string, object> { { "error", ex.GetType().Name + ": " + ex.Message } });
            }
            return pipes;
        }

        private const int AF_INET = 2;
        private const int TCP_TABLE_OWNER_PID_ALL = 5;

        [DllImport("iphlpapi.dll", SetLastError = true)]
        private static extern uint GetExtendedTcpTable(IntPtr pTcpTable, ref int dwOutBufLen, bool sort, int ipVersion, int tblClass, uint reserved);

        [StructLayout(LayoutKind.Sequential)]
        private struct MIB_TCPROW_OWNER_PID
        {
            public uint state;
            public uint localAddr;
            public uint localPort;
            public uint remoteAddr;
            public uint remotePort;
            public uint owningPid;
        }

        private static Dictionary<string, object> CaptureTargetTcpConnections(HashSet<int> targetPids)
        {
            var rows = new List<Dictionary<string, object>>();
            string error = "";
            try
            {
                int bufferSize = 0;
                GetExtendedTcpTable(IntPtr.Zero, ref bufferSize, true, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0);
                IntPtr buffer = Marshal.AllocHGlobal(bufferSize);
                try
                {
                    uint result = GetExtendedTcpTable(buffer, ref bufferSize, true, AF_INET, TCP_TABLE_OWNER_PID_ALL, 0);
                    if (result != 0)
                    {
                        error = "GetExtendedTcpTable result=" + result.ToString(CultureInfo.InvariantCulture);
                    }
                    else
                    {
                        int count = Marshal.ReadInt32(buffer);
                        IntPtr rowPtr = new IntPtr(buffer.ToInt64() + 4);
                        int rowSize = Marshal.SizeOf(typeof(MIB_TCPROW_OWNER_PID));
                        for (int i = 0; i < count; i++)
                        {
                            var row = (MIB_TCPROW_OWNER_PID)Marshal.PtrToStructure(rowPtr, typeof(MIB_TCPROW_OWNER_PID));
                            int pid = unchecked((int)row.owningPid);
                            if (targetPids.Contains(pid))
                            {
                                rows.Add(new Dictionary<string, object>
                                {
                                    { "pid", pid },
                                    { "state", TcpStateName(row.state) },
                                    { "local_address", UIntToIp(row.localAddr) },
                                    { "local_port", NetworkToHostPort(row.localPort) },
                                    { "remote_address", UIntToIp(row.remoteAddr) },
                                    { "remote_port", NetworkToHostPort(row.remotePort) }
                                });
                            }
                            rowPtr = new IntPtr(rowPtr.ToInt64() + rowSize);
                        }
                    }
                }
                finally
                {
                    Marshal.FreeHGlobal(buffer);
                }
            }
            catch (Exception ex)
            {
                error = ex.GetType().Name + ": " + ex.Message;
            }

            return new Dictionary<string, object>
            {
                { "target_pid_count", targetPids.Count },
                { "connection_count", rows.Count },
                { "connections", rows },
                { "error", error }
            };
        }

        private static string UIntToIp(uint value)
        {
            try
            {
                return new IPAddress(value).ToString();
            }
            catch
            {
                return value.ToString(CultureInfo.InvariantCulture);
            }
        }

        private static int NetworkToHostPort(uint port)
        {
            return (int)(((port & 0xFF) << 8) | ((port >> 8) & 0xFF));
        }

        private static string TcpStateName(uint state)
        {
            switch (state)
            {
                case 1: return "CLOSED";
                case 2: return "LISTEN";
                case 3: return "SYN_SENT";
                case 4: return "SYN_RCVD";
                case 5: return "ESTABLISHED";
                case 6: return "FIN_WAIT1";
                case 7: return "FIN_WAIT2";
                case 8: return "CLOSE_WAIT";
                case 9: return "CLOSING";
                case 10: return "LAST_ACK";
                case 11: return "TIME_WAIT";
                case 12: return "DELETE_TCB";
                default: return state.ToString(CultureInfo.InvariantCulture);
            }
        }

        private static IEnumerable<ManagementObject> QueryWmi(string query)
        {
            var list = new List<ManagementObject>();
            try
            {
                using (var searcher = new ManagementObjectSearcher(query))
                using (ManagementObjectCollection results = searcher.Get())
                {
                    foreach (ManagementObject obj in results)
                    {
                        list.Add(obj);
                    }
                }
            }
            catch
            {
            }
            return list;
        }

        private static string SafeWmi(ManagementObject obj, string property)
        {
            try
            {
                object value = obj[property];
                return value == null ? "" : Convert.ToString(value);
            }
            catch
            {
                return "";
            }
        }

        private static bool AnyAssemblyPresent(Dictionary<string, object> runtime, string name)
        {
            foreach (Dictionary<string, object> target in GetDictionaryList(runtime, "targets"))
            {
                object rawPresence;
                if (!target.TryGetValue("assembly_presence", out rawPresence))
                {
                    continue;
                }
                var presence = rawPresence as Dictionary<string, object>;
                if (presence == null)
                {
                    continue;
                }
                object value;
                if (presence.TryGetValue(name, out value) && value is bool && (bool)value)
                {
                    return true;
                }
            }
            return false;
        }

        private static HashSet<int> TargetPids(Dictionary<string, object> runtime)
        {
            var pids = new HashSet<int>();
            foreach (Dictionary<string, object> target in GetDictionaryList(runtime, "targets"))
            {
                int pid = GetInt(target, "pid");
                if (pid > 0)
                {
                    pids.Add(pid);
                }
            }
            return pids;
        }

        private static List<Dictionary<string, object>> GetDictionaryList(Dictionary<string, object> source, string key)
        {
            var list = new List<Dictionary<string, object>>();
            object value;
            if (source == null || !source.TryGetValue(key, out value) || value == null)
            {
                return list;
            }
            var direct = value as List<Dictionary<string, object>>;
            if (direct != null)
            {
                return direct;
            }
            return list;
        }

        private static string GetString(Dictionary<string, object> source, string key)
        {
            object value;
            if (source != null && source.TryGetValue(key, out value) && value != null)
            {
                return Convert.ToString(value);
            }
            return "";
        }

        private static int GetInt(Dictionary<string, object> source, string key)
        {
            object value;
            if (source != null && source.TryGetValue(key, out value) && value != null)
            {
                try
                {
                    return Convert.ToInt32(value);
                }
                catch
                {
                }
            }
            return 0;
        }

        private static List<string> SafeDirectories(string root, int maxDepth, int maxItems)
        {
            var result = new List<string>();
            SafeDirectories(root, 0, maxDepth, maxItems, result);
            return result;
        }

        private static void SafeDirectories(string root, int depth, int maxDepth, int maxItems, List<string> result)
        {
            if (depth > maxDepth || result.Count >= maxItems || !Directory.Exists(root))
            {
                return;
            }
            string[] dirs;
            try
            {
                dirs = Directory.GetDirectories(root);
            }
            catch
            {
                return;
            }
            foreach (string dir in dirs)
            {
                if (result.Count >= maxItems)
                {
                    return;
                }
                result.Add(dir);
                SafeDirectories(dir, depth + 1, maxDepth, maxItems, result);
            }
        }

        private static List<string> SafeFiles(string root, int maxDepth, int maxItems)
        {
            var result = new List<string>();
            SafeFiles(root, 0, maxDepth, maxItems, result);
            return result;
        }

        private static void SafeFiles(string root, int depth, int maxDepth, int maxItems, List<string> result)
        {
            if (depth > maxDepth || result.Count >= maxItems || !Directory.Exists(root))
            {
                return;
            }
            string[] files;
            try
            {
                files = Directory.GetFiles(root);
            }
            catch
            {
                return;
            }
            foreach (string file in files)
            {
                if (result.Count >= maxItems)
                {
                    return;
                }
                result.Add(file);
            }
            string[] dirs;
            try
            {
                dirs = Directory.GetDirectories(root);
            }
            catch
            {
                return;
            }
            foreach (string dir in dirs)
            {
                if (result.Count >= maxItems)
                {
                    return;
                }
                SafeFiles(dir, depth + 1, maxDepth, maxItems, result);
            }
        }

        private static bool IsInterestingFile(string path)
        {
            string name = Path.GetFileName(path).ToLowerInvariant();
            string ext = Path.GetExtension(path).ToLowerInvariant();
            if (ContainsInterestingTerm(name))
            {
                return true;
            }
            foreach (string interestingExt in InterestingExtensions)
            {
                if (ext == interestingExt)
                {
                    return true;
                }
            }
            return false;
        }

        private static bool IsKeyBinary(string name)
        {
            string lower = name.ToLowerInvariant();
            return lower == "implatformclient.exe" ||
                   lower == "smack.dll" ||
                   lower == "impp.interface.dll" ||
                   lower == "impp.service.dll" ||
                   lower == "impp.servicebase.dll" ||
                   lower == "tcpfiletransfer.dll" ||
                   lower == "hyhc.impp.dal.dll" ||
                   lower == "inetwork.dll" ||
                   lower == "ioclientnetwork.dll";
        }

        private static bool IsSmallTextFile(FileInfo info)
        {
            if (info.Length > 256 * 1024)
            {
                return false;
            }
            string ext = info.Extension.ToLowerInvariant();
            return ext == ".config" || ext == ".xml" || ext == ".ini" || ext == ".json" || ext == ".log" || ext == ".txt";
        }

        private static string RedactedTextSample(string path, int maxChars)
        {
            try
            {
                string text = File.ReadAllText(path, Encoding.UTF8);
                if (text.Length > maxChars)
                {
                    text = text.Substring(0, maxChars);
                }
                return RedactSecrets(text);
            }
            catch
            {
                try
                {
                    byte[] bytes = File.ReadAllBytes(path);
                    int len = Math.Min(bytes.Length, maxChars);
                    string text = Encoding.Default.GetString(bytes, 0, len);
                    return RedactSecrets(text);
                }
                catch
                {
                    return "";
                }
            }
        }

        private static string RedactSecrets(string text)
        {
            if (string.IsNullOrEmpty(text))
            {
                return "";
            }
            string redacted = Regex.Replace(text, "(password|passwd|pwd|token|secret|session|cookie|auth|authorization)\\s*[:=]\\s*[^\\s;&<>\\\"]+", "$1=<redacted>", RegexOptions.IgnoreCase);
            redacted = Regex.Replace(redacted, "(UserID|UserName|LoginName)\\s*[:=]\\s*[^\\s;&<>\\\"]+", "$1=<redacted>", RegexOptions.IgnoreCase);
            return redacted;
        }

        private static bool ContainsInterestingTerm(string text)
        {
            if (string.IsNullOrEmpty(text))
            {
                return false;
            }
            string lower = text.ToLowerInvariant();
            foreach (string term in InterestingTerms)
            {
                if (lower.Contains(term))
                {
                    return true;
                }
            }
            return false;
        }

        private static string Sha256(string path)
        {
            try
            {
                using (SHA256 sha = SHA256.Create())
                using (FileStream stream = File.OpenRead(path))
                {
                    byte[] hash = sha.ComputeHash(stream);
                    var sb = new StringBuilder(hash.Length * 2);
                    foreach (byte b in hash)
                    {
                        sb.Append(b.ToString("x2", CultureInfo.InvariantCulture));
                    }
                    return sb.ToString();
                }
            }
            catch
            {
                return "";
            }
        }
    }
}