using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Threading;
using System.Runtime.InteropServices;
using System.ServiceProcess;
using System.Text;
using System.Threading.Tasks;
using Microsoft.Win32;

namespace ame_integrity_check
{
    public enum Type
    {
        File = 1,
        Directory = 2,
        Process = 3,
        Service = 4
    }
    internal class Scanner : IDisposable
    {
        private int result = 1;
        private bool allFound = true;
        private bool found = false;
        private bool errorOverride = false;
        public Task<bool> displayTask;
        
        public void Dispose() => GC.SuppressFinalize(this);

        public async Task<bool> DisplayQuery(string text, int time = 150)
        {
            Out.WriteCustomString(text, 3, 11);
            
            string maxSpaces = "      ";
            for (int i = 1; i < 6; i++) {
                Console.SetCursorPosition(59, Console.CursorTop);
                var spaces = maxSpaces.Remove(0, i);
                
                Console.Write($"[ {spaces.PadLeft(spaces.Length + i, '*')} ]");
                Thread.Sleep(time);
            }
            
            return true;
        }

        private static void SetQueryStatus(string status, ConsoleColor color)
        {
            Console.SetCursorPosition(64 - status.Length, Console.CursorTop);
            Console.Write(" [ ");
            Out.WriteCustomString(status, 3, 0, foregroundColor: color);
            Console.WriteLine(" ]");
        }
        
        public async void Query(Type type, string item, bool finalize = false, bool modifyResult = true)
        {
            item = Environment.ExpandEnvironmentVariables(item);
            bool foundItem = false;
            try {
                switch (type) {
                    case Type.File:
                        if (item.Contains("*"))
                        {
                            var lastToken = item.LastIndexOf("\\");
                            var parentPath = item.Remove(lastToken).TrimEnd('\\');

                            if (parentPath.Contains("*")) throw new ArgumentException("Parent directories to a given file filter cannot contain wildcards.");
                            var filter = item.Substring(lastToken + 1);

                            foundItem = Directory.GetFiles(parentPath, filter).Any();
                            break;
                        }
                        
                        foundItem = File.Exists(item);
                        break;
                    case Type.Directory:
                        if (item.Contains("*"))
                        {
                            var lastToken = item.LastIndexOf("\\");
                            var parentPath = item.Remove(lastToken).TrimEnd('\\');
                            
                            if (parentPath.Contains("*")) throw new ArgumentException("Parent directories to a given file filter cannot contain wildcards.");
                            var filter = item.Substring(lastToken + 1);

                            var foundDirs = Directory.GetDirectories(parentPath, filter);
                            
                            foreach (var foundDir in foundDirs)
                            {
                                foreach (var file in Directory.GetFiles(foundDir, "*", SearchOption.AllDirectories))
                                {
                                    if (!file.ToLower().EndsWith(".mui") && !file.ToLower().EndsWith(".pri") && !file.ToLower().EndsWith(".res"))
                                    {
                                        foundItem = true;
                                    }
                                }
                            }
                        }
                        else
                        {

                            if (Directory.Exists(item))
                            {
                                foreach (var file in Directory.GetFiles(item, "*", SearchOption.AllDirectories))
                                {
                                    if (!file.ToLower().EndsWith(".mui") && !file.ToLower().EndsWith(".pri") && !file.ToLower().EndsWith(".res"))
                                    {
                                        foundItem = true;
                                    }
                                }
                            }
                        }

                        break;
                    case Type.Process:
                        foundItem = Process.GetProcessesByName(item).Any();
                        break;
                    case Type.Service:
                        foundItem = ServiceController.GetServices().Any(x => x.ServiceName.Equals("wuauserv", StringComparison.CurrentCultureIgnoreCase));
                        break;
                    default:
                        foundItem = false;
                        break;
                }
            } catch (Exception e) {
                if (e.GetType().ToString() == "System.UnauthorizedAccessException" || e.GetType().ToString() == "System.Security.SecurityException")
                {
                    foundItem = true;
                }
                else
                {
                    errorOverride = true;
                }
            }

            if (foundItem) found = true;

            if (!finalize) return;

            await displayTask;
            
            if (errorOverride) {
                errorOverride = false;
                SetQueryStatus("ERROR", ConsoleColor.DarkRed);
                found = false;
                return;
            }
            
            if (!found) {
                if (modifyResult) allFound = false;
                SetQueryStatus("Absent", ConsoleColor.Green);
            } else {
                result = 2;
                if (allFound) result = 3;

                SetQueryStatus("Present", ConsoleColor.DarkRed);
            }
            found = false;
        }

        public void DisplayResult()
        {
            switch (result) {
                case 1:
                    Out.WriteCustomString("\n\nAME Integrity validated", 1, foregroundColor: ConsoleColor.Green);
                    break;
                case 2:
                    Out.WriteCustomString("\n\nAME integrity compromised, contact the team for help.", 1, foregroundColor: ConsoleColor.Red);
                    break;
                case 3:
                    Out.WriteCustomString("\n\nYour system is not ameliorated.", 1, foregroundColor: ConsoleColor.Red);
                    break;
            }
        }
    }
    
    internal static class Program
    {
        public static string PreviousTitle;
        public static int PreviousBufferHeight;
        public static int PreviousBufferWidth;
        public static int PreviousSizeHeight;
        public static int PreviousSizeWidth;

        public const string Ver = "1.0";
        
        public static void Main(string[] args)
        {
            PreviousSizeHeight = Console.WindowHeight;
            PreviousSizeWidth = Console.WindowWidth;
            PreviousBufferHeight = Console.BufferHeight;
            PreviousBufferWidth = Console.BufferWidth;

            try
            {
                Console.SetWindowSize(80, 26);
                Console.SetBufferSize(80, 26);
                Console.SetWindowSize(80, 26);
            } catch (Exception e)
            {

            }
            
            Console.Clear();

            Console.CursorVisible = false;
            PreviousTitle = Console.Title;
            Console.Title = "AME Integrity Check";
            
            try { Out.DisableResize(); } catch (Exception) { }
            try { Out.DisableQuickEdit(); } catch (Exception) { }

            Console.CancelKeyPress += Exit;

            Out.WriteCustomString("\n__________________________________________________________" +
                                  $"\n\n| AME Integrity Check v{Ver} |\n\n", 1);
            
            displayStart:
            
            Menu mainMenu = new Menu() {
                Choices = {"Check AME Integrity", "Get Support", "", "Exit"},
                EndString = "\n__________________________________________________________\n",
                Statement = "Use the arrows keys to navigate"
            };
            
            var choice = mainMenu.Load();
            switch (choice) {
                case 0:
                    RunCheck();
                    Out.ResetPane();
                    goto displayStart;
                case 1:
                    Process.Start("https://t.me/joinchat/CR-xFBGQKVt7HPZKgZfbxg");
                    Out.ResetPane();
                    goto displayStart;
                case 2:
                    PrepProcessExit();
                    Environment.Exit(0);
                    break;
                default:
                    PrepProcessExit();
                    Environment.Exit(1);
                    break;
            }
        }

        private static bool detected = false;
        private static int result = 1;
        private static void RunCheck()
        {
            Out.ResetPane();

            bool legacy = false;
            
            var registryKey = Registry.LocalMachine.OpenSubKey(@"SOFTWARE\Microsoft\Windows NT\CurrentVersion");
            int winVer = 1;
            if (registryKey != null) winVer = Int32.Parse(registryKey.GetValue("CurrentBuildNumber").ToString());
            if (winVer < 19044) legacy = true;

            using (var scanner = new Scanner())
            {
                scanner.displayTask = scanner.DisplayQuery("Checking for Windows Defender activity...", 250);
                scanner.Query(Type.Process, "MsMpEng", true);
                
                scanner.displayTask = scanner.DisplayQuery("Checking Windows Defender files...", 200);
                if (!legacy) {
                    scanner.Query(Type.Directory, "%ProgramFiles%\\Windows Defender");
                    scanner.Query(Type.Directory, "%ProgramData%\\Microsoft\\Windows Defender", true);
                }
                else scanner.Query(Type.Directory, "%ProgramFiles%\\Windows Defender", true);

                if (!legacy) {
                    scanner.displayTask = scanner.DisplayQuery("Checking Windows Update service...", 350);
                    scanner.Query(Type.Service, "wuauserv", true);
                }

                scanner.displayTask = scanner.DisplayQuery("Checking Windows Update files...", 220);
                scanner.Query(Type.File, "%WINDIR%\\System32\\wuaueng.dll");
                scanner.Query(Type.File, "%WINDIR%\\System32\\wuapi.dll", true);
                
                scanner.displayTask = scanner.DisplayQuery("Checking Microsoft Edge...", 200);
                scanner.Query(Type.Directory, "%ProgramFiles(x86)%\\Microsoft\\Edge");
                scanner.Query(Type.Directory, "%WINDIR%\\SystemApps\\*MicrosoftEdge*", true);
                
                scanner.displayTask = scanner.DisplayQuery("Checking for Microsoft Store activity...", 200);
                scanner.Query(Type.Process, "WinStore.App", true, false);
                
                scanner.displayTask = scanner.DisplayQuery("Checking Windows SmartScreen...");
                scanner.Query(Type.Process, "smartscreen");
                scanner.Query(Type.File, "%WINDIR%\\System32\\smartscreen.exe", true);
                
                scanner.displayTask = scanner.DisplayQuery("Checking SIH Client...");
                scanner.Query(Type.File, "%WINDIR%\\System32\\SIHClient.exe", true);
                
                scanner.displayTask = scanner.DisplayQuery("Checking Storage Sense...", 300);
                scanner.Query(Type.File, "%WINDIR%\\System32\\StorSvc.dll", true);
                
                scanner.DisplayResult();
            }

            Out.WriteCustomString("__________________________________________________________\n\n\nPress any key to return to the Menu: ", 3, 11);
            
            ClearBuffer();
            
            Console.CursorVisible = true;
            Console.ReadKey(false);
            Console.CursorVisible = false;
        }

        private static bool anyNotPresent;

        private static void ClearBuffer()
        {
            var posCacheTop = Console.CursorTop;
            var posCacheLeft = Console.CursorLeft;

            while (Console.KeyAvailable) {
                Console.ReadKey(true);
            }
            //Console.SetCursorPosition(posCacheLeft, posCacheTop);
            //Console.Write("".PadLeft(Console.WindowHeight, ' '));
            //Console.SetCursorPosition(posCacheLeft, posCacheTop);
        }
        private static void Exit(object sender, ConsoleCancelEventArgs args)
        {
            PrepProcessExit();
            Environment.Exit(0);
        }
        
        private static void PrepProcessExit()
        {
            var parent = ParentProcess.ProcessName;
            if (parent.Equals("Explorer", StringComparison.CurrentCultureIgnoreCase)) return;

            try { Out.EnableResize(); } catch (Exception) { }
            try { Out.EnableQuickEdit(); } catch (Exception) { }
            
            Console.CursorVisible = true;
            Console.Clear();
            Console.Title = PreviousTitle;

            try
            {
                Console.SetWindowSize(PreviousSizeWidth, PreviousSizeHeight);
                Console.SetBufferSize(PreviousBufferWidth, PreviousBufferHeight);
            } catch (Exception e)
            {

            }
        }
    }


    internal static class Out
    {
        public static void ResetPane(int fromTop = 6)
        {
            Console.SetCursorPosition(Console.CursorLeft, fromTop - 1);
            var length = Console.WindowHeight - fromTop - 1;

            for (int i = 0; i < length; i++) {
                Console.Write("".PadRight(Console.WindowWidth, ' '));
            }
            Console.SetCursorPosition(0, fromTop);
        }
        public static void WriteCustomString(string text, int type, int offset = 11, ConsoleColor foregroundColor = ConsoleColor.DarkYellow, ConsoleColor backgroundColor = ConsoleColor.DarkYellow)
        {
            bool resetColor = false;
            
            ConsoleColor foregroundCache = Console.ForegroundColor;
            ConsoleColor backgroundCache = Console.BackgroundColor;
            
            if (foregroundColor == ConsoleColor.DarkYellow) {
                foregroundColor = Console.ForegroundColor;
            }
            if (backgroundColor == ConsoleColor.DarkYellow) {
                backgroundColor = Console.BackgroundColor;
            }
            
            if (Console.ForegroundColor != foregroundColor) {
                resetColor = true;

                Console.ForegroundColor = foregroundColor;
            }
            if (Console.BackgroundColor != backgroundColor) {
                resetColor = true;

                Console.BackgroundColor = backgroundColor;
            }

            switch (type) {
                case 1:
                    // Indented and centered
                    Console.WriteLine(CenterString(text, offset));
                    break;
                case 2:
                    Console.WriteLine(text.Insert(0, new string(' ', offset)));
                    break;
                case 3:
                    var loopOnce = false;
                    foreach (string line in text.Split('\n')) {
                        if (line == "") {
                            Console.WriteLine();
                            continue;
                        }
                        Console.Write(line.Insert(0, new string(' ', offset)));
                    }

                    break;
                default:
                    break;
            }

            if (resetColor) {
                Console.ForegroundColor = foregroundCache;
                Console.BackgroundColor = backgroundCache;
            }
        }

        private static string CenterString(string text, int offset = 11, int width = 58)
        {
            StringBuilder subLines = new StringBuilder();
            string newLine = "";
            
            foreach (string line in text.Split('\n')) {
                if (!subLines.ToString().Equals("")) newLine = "\n";
                
                if (line == "" && !subLines.ToString().Equals("")) {
                    subLines.Append("\n");
                    continue;
                }
                
                var space = "";
                if (!(line.Length % 2).Equals(0) && line.Length != width) {
                    space = " ";
                }
                
                if (line.Length > width) {
                    for (int index = 0; index < line.Length; index += (width - 10)) {
                        if (!subLines.ToString().Equals("")) newLine = "\n";
                        
                        var subLine = line.Substring(index, Math.Min(width - 10, line.Length - index));
                        subLine = subLine.Trim(' ');

                        var subCentered = CenterString(subLine, offset, width);
                        subLines.Append(newLine + subCentered);
                    }
                    continue;
                }

                var leadingLength = (width - line.Length) / 2;
                subLines.Append(newLine + space + line.PadLeft(line.Length + leadingLength, ' ').Insert(0, new string(' ', offset)));
            }

            return subLines.ToString();
        }
        
        private const int MF_BYCOMMAND = 0x00000000;
        public const int SC_CLOSE = 0xF060;
        public const int SC_MINIMIZE = 0xF020;
        public const int SC_MAXIMIZE = 0xF030;
        public const int SC_SIZE = 0xF000;//resize

        [DllImport("user32.dll")]
        public static extern int DeleteMenu(IntPtr hMenu, int nPosition, int wFlags);

        [DllImport("user32.dll")]
        private static extern IntPtr GetSystemMenu(IntPtr hWnd, bool bRevert);

        [DllImport("kernel32.dll", ExactSpelling = true)]
        private static extern IntPtr GetConsoleWindow();

        public static void DisableResize()
        {
            IntPtr handle = GetConsoleWindow();
            IntPtr sysMenu = GetSystemMenu(handle, false);

            if (handle != IntPtr.Zero)
            {
                //DeleteMenu(sysMenu, SC_CLOSE, MF_BYCOMMAND);
                //DeleteMenu(sysMenu, SC_MINIMIZE, MF_BYCOMMAND);
                DeleteMenu(sysMenu, SC_MAXIMIZE, MF_BYCOMMAND);
                DeleteMenu(sysMenu, SC_SIZE, MF_BYCOMMAND);//resize
            }
        }
        public static void EnableResize()
        {
            IntPtr handle = GetConsoleWindow();
            GetSystemMenu(handle, true);
        }
        
        const uint CHECK_QUICK_EDIT = 0x0040;
        const int ENABLE_QUICK_EDIT = 0x40 | 0x80;
        
        // STD_INPUT_HANDLE (DWORD): -10 is the standard input device.
        const int STD_INPUT_HANDLE = -10;

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr GetStdHandle(int nStdHandle);

        [DllImport("kernel32.dll")]
        static extern bool GetConsoleMode(IntPtr hConsoleHandle, out uint lpMode);

        [DllImport("kernel32.dll")]
        static extern bool SetConsoleMode(IntPtr hConsoleHandle, uint dwMode);

        internal static void DisableQuickEdit() {

            IntPtr consoleHandle = GetStdHandle(STD_INPUT_HANDLE);

            // get current console mode
            uint consoleMode;
            GetConsoleMode(consoleHandle, out consoleMode);
            
            // set the new mode
            SetConsoleMode(consoleHandle, consoleMode &= ~CHECK_QUICK_EDIT);
        }
        internal static void EnableQuickEdit() {
            IntPtr consoleHandle = GetStdHandle(STD_INPUT_HANDLE);

            // get current console mode
            uint consoleMode;
            GetConsoleMode(consoleHandle, out consoleMode);
            
            // set the new mode
            SetConsoleMode(consoleHandle, consoleMode|(ENABLE_QUICK_EDIT));
        }
    }
    internal class Menu
    {
        public List<string> Choices;
        public string EndString;
        public string Statement;
        
        public int Offset;


        public Menu(int offset = 16, string endString = null)
        {
            Choices = new List<string>();

            EndString = endString;
            Offset = offset;
        }

        public int Load()
        {
            int max = -1;
            var posCache = Console.CursorTop;

            bool loopedOnce = false;
            foreach (string choice in Choices) {
                Console.SetCursorPosition(Offset - 2, Console.CursorTop);
                if (!loopedOnce) {
                    Out.WriteCustomString("> " + choice, 2, 0, foregroundColor: ConsoleColor.Green);
                    loopedOnce = true;
                } else {
                    Out.WriteCustomString(choice, 2, 2);
                }
                
                max += 1;
            }
            
            if (EndString != null) Out.WriteCustomString(EndString, 1);
            if (Statement != null) Out.WriteCustomString(Statement, 2);
            
            int index = 0;
            ConsoleKey keyPressed;
            
            Console.SetCursorPosition(Offset, posCache);

            while ((keyPressed = Console.ReadKey(true).Key) != ConsoleKey.Enter) {
                if (keyPressed == ConsoleKey.DownArrow) {
                    if (index >= max) continue;
                    Console.SetCursorPosition(Offset - 2, Console.CursorTop);
                    Out.WriteCustomString(Choices[index], 3, 2);

                    if (!String.IsNullOrEmpty(Choices[index + 1])) {
                        index += 1;
                        Console.SetCursorPosition(Offset - 2, Console.CursorTop + 1);
                        Out.WriteCustomString("> " + Choices[index], 3, 0, foregroundColor: ConsoleColor.Green);
                    } else {
                        index += 2;
                        Console.SetCursorPosition(Offset - 2, Console.CursorTop + 2);
                        Out.WriteCustomString("> " + Choices[index], 3, 0, foregroundColor: ConsoleColor.Green);
                    }
                }
                if (keyPressed == ConsoleKey.UpArrow) {
                    if (!(index > 0)) continue;
                    Console.SetCursorPosition(Offset - 2, Console.CursorTop);
                    Out.WriteCustomString(Choices[index], 3, 2);
                    
                    if (!String.IsNullOrEmpty(Choices[index - 1])) {
                        index -= 1;
                        Console.SetCursorPosition(Offset - 2, Console.CursorTop - 1);
                        Out.WriteCustomString("> " + Choices[index], 3, 0, foregroundColor: ConsoleColor.Green);
                    } else {
                        index -= 2;
                        Console.SetCursorPosition(Offset - 2, Console.CursorTop - 2);
                        Out.WriteCustomString("> " + Choices[index], 3, 0, foregroundColor: ConsoleColor.Green);
                    }
                }
            }
            return index;
        }
    }
    public static class ParentProcess
    {
        public static string ProcessName
        {
            get { return GetParentProcess().ProcessName; }
        }

        private static Process GetParentProcess()
        {
            int iParentPid = 0;
            int iCurrentPid = Process.GetCurrentProcess().Id;
 
            IntPtr oHnd = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
 
            if (oHnd == IntPtr.Zero)
                return null;
 
            PROCESSENTRY32 oProcInfo = new PROCESSENTRY32();
 
            oProcInfo.dwSize =
            (uint)System.Runtime.InteropServices.Marshal.SizeOf(typeof(PROCESSENTRY32));
 
            if (Process32First(oHnd, ref oProcInfo) == false)
                return null;
 
            do
            {
                if (iCurrentPid == oProcInfo.th32ProcessID)
                    iParentPid = (int)oProcInfo.th32ParentProcessID;
            }
            while (iParentPid == 0 && Process32Next(oHnd, ref oProcInfo));
 
            if (iParentPid > 0)
                return Process.GetProcessById(iParentPid);
            else
                return null;
        }
 
        static uint TH32CS_SNAPPROCESS = 2;
 
        [StructLayout(LayoutKind.Sequential)]
        public struct PROCESSENTRY32
        {
            public uint dwSize;
            public uint cntUsage;
            public uint th32ProcessID;
            public IntPtr th32DefaultHeapID;
            public uint th32ModuleID;
            public uint cntThreads;
            public uint th32ParentProcessID;
            public int pcPriClassBase;
            public uint dwFlags;
            [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
            public string szExeFile;
        };
 
        [DllImport("kernel32.dll", SetLastError = true)]
        static extern IntPtr CreateToolhelp32Snapshot(uint dwFlags, uint th32ProcessID);
 
        [DllImport("kernel32.dll")]
        static extern bool Process32First(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);
 
        [DllImport("kernel32.dll")]
        static extern bool Process32Next(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);
    }
}