// InstallDrv.cpp : Этот файл содержит функцию "main". Здесь начинается и заканчивается выполнение программы.
//

#define _CRT_SECURE_NO_WARNINGS

#include <windows.h>
#include <stdio.h> 

#define bzero(b) (memset((b), '\0', sizeof(b)), (void) 0)  

char volumeName[_MAX_PATH];

BOOL SetPrivilege(
    HANDLE hToken,
    LPCTSTR lpszPrivilege,  // name of privilege to enable/disable
    BOOL bEnablePrivilege   // to enable or disable privilege
)
{
    TOKEN_PRIVILEGES tp;
    LUID luid;

    if (!LookupPrivilegeValue(
        NULL,            // lookup privilege on local system
        lpszPrivilege,   // privilege to lookup 
        &luid))        // receives LUID of privilege
    {
        //           printf("LookupPrivilegeValue error: %u\n", GetLastError() ); 
        return FALSE;
    }

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    if (bEnablePrivilege) {
        tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    }
    else {
        tp.Privileges[0].Attributes = 0;
    }

    // Enable the privilege or disable all privileges.

    if (!AdjustTokenPrivileges(
        hToken,
        FALSE,
        &tp,
        sizeof(TOKEN_PRIVILEGES),
        (PTOKEN_PRIVILEGES)NULL,
        (PDWORD)NULL)) {
        //            printf("AdjustTokenPrivileges error: %u\n", GetLastError() ); 
        return FALSE;
    }

    if (GetLastError() == ERROR_NOT_ALL_ASSIGNED) {
        //            printf("The token does not have the specified privilege. \n");
        return FALSE;
    }

    return TRUE;
}

int setSeDebug(HANDLE handleThreadId)
{
    HANDLE hToken;
    if (!OpenThreadToken(handleThreadId,
        TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
        FALSE,
        &hToken)
        ) {
        if (GetLastError() == ERROR_NO_TOKEN) {
            if (!ImpersonateSelf(SecurityImpersonation)) {
                //Log2File("Error setting impersonation! [UnsetSeDebug()]", L_DEBUG);
                return 0;
            }

            if (!OpenThreadToken(handleThreadId,
                TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
                FALSE,
                &hToken)
                ) {
                //Log2File("Error Opening Thread Token! [UnsetSeDebug()]", L_DEBUG);
                return 0;
            }
        }
    }

    //now disable SeDebug
    if (!SetPrivilege(hToken, SE_DEBUG_NAME, TRUE)) {
        //Log2File("Error unsetting SeDebug Privilege [SetPrivilege()]", L_WARN);
        return 0;
    }

    CloseHandle(hToken);
    return 1;
}

int runProcess(char* strDevcon, BOOL isVolume) {
//    char strDevcon[_MAX_PATH];
//    strcpy_s(strDevcon, _MAX_PATH, "1.bat ");

//    for (int i = 1; i < argc; ++i) {
//        strcat_s(strDevcon, _MAX_PATH, argv[i]);
//        strcat_s(strDevcon, _MAX_PATH, " ");
//    }

    char buf[1024 * 10];           //буфер ввода/вывода

    PROCESS_INFORMATION piProcInfo;
    STARTUPINFOA siStartInfo;
    SECURITY_ATTRIBUTES sa;
    SECURITY_DESCRIPTOR sd;        //структура security для пайпов
    HANDLE newstdin, newstdout, read_stdout, write_stdin;  //дескрипторы пайпов

    //    if (IsWinNT())        //инициализация security для Windows NT
    //    {
    InitializeSecurityDescriptor(&sd, SECURITY_DESCRIPTOR_REVISION);
    SetSecurityDescriptorDacl(&sd, TRUE, NULL, FALSE);
    sa.lpSecurityDescriptor = &sd;
    //    }

    //    else sa.lpSecurityDescriptor = NULL;

    sa.nLength = sizeof(SECURITY_ATTRIBUTES);
    sa.bInheritHandle = TRUE;       //разрешаем наследование дескрипторов

    if (!CreatePipe(&newstdin, &write_stdin, &sa, 0))   //создаем пайп
    {
        printf("Error creating input channel");
        return 1;
    }

    if (!CreatePipe(&read_stdout, &newstdout, &sa, 0)) //создаем пайп
        // для stdout
    {
        printf("Error creating output channel");

        CloseHandle(newstdin);
        CloseHandle(write_stdin);
        return 1;
    }
    ZeroMemory(&siStartInfo, sizeof(siStartInfo));
    ZeroMemory(&piProcInfo, sizeof(piProcInfo));
    GetStartupInfoA(&siStartInfo);      //создаем startupinfo для дочернего процесса
    siStartInfo.cb = sizeof(STARTUPINFO);
    siStartInfo.wShowWindow = SW_HIDE;     //прячем окно
    siStartInfo.dwFlags = STARTF_USESTDHANDLES | STARTF_USESHOWWINDOW;
    //   siStartInfo.lpDesktop   = NULL;
    siStartInfo.hStdOutput = newstdout;
    siStartInfo.hStdError = newstdout;   //подменяем дескрипторы для
    siStartInfo.hStdInput = newstdin;    // дочернего процесса
    //char app_spawn[] = "devcon.exe";


    int result = CreateProcessA(NULL, strDevcon, &sa, &sa, TRUE, CREATE_NEW_CONSOLE | NORMAL_PRIORITY_CLASS,
        NULL, NULL, &siStartInfo, &piProcInfo);


    if (!result)
    {
        printf("Error starting thread");

        CloseHandle(newstdin);
        CloseHandle(newstdout);
        CloseHandle(read_stdout);
        CloseHandle(write_stdin);
        return 1;
    }

    unsigned long exit = 0;  //код завершения процесса
    unsigned long bread;   //кол-во прочитанных байт
    unsigned long avail;   //кол-во доступных байт
    bzero(buf);

    for (;;)      //основной цикл программы
    {
        Sleep(1);

        GetExitCodeProcess(piProcInfo.hProcess, &exit); //пока дочерний процесс

        // не закрыт
        if (exit != STILL_ACTIVE) {
            break;
        }

    }


    PeekNamedPipe(read_stdout, buf, sizeof(buf) - 1, &bread, &avail, NULL);

    //Проверяем, есть ли данные для чтения в stdout

    if (bread != 0)
    {
        bzero(buf);

        ReadFile(read_stdout, buf, avail, &bread, NULL);  //читаем из пайпа stdout

        //printf("%s", buf);

    }

    int volume = 0;

    if (isVolume) {

        char bufStr[1024];
        bzero(bufStr);
        int j = 0;
        
        for (unsigned long i = 0; i < bread; ++i) {

            if (buf[i] == 0x0A) {
                printf("%s", bufStr);

                char* res_ptr = strstr(bufStr, volumeName);

                if (res_ptr) {

                    char str[_MAX_PATH];
                    
                    sscanf(bufStr, "%s %d", str, &volume);

                    //volume = atoi(&bufStr[10]);
                }

                bzero(bufStr);
                j = 0;
            }

            bufStr[j] = buf[i];
            ++j;

        }

    }


    //    KillProcByPid(piProcInfo.dwProcessId);

    CloseHandle(piProcInfo.hThread);
    CloseHandle(piProcInfo.hProcess);                                           // в stdin
    CloseHandle(newstdin);
    CloseHandle(newstdout);
    CloseHandle(read_stdout);
    CloseHandle(write_stdin);

    return volume;
}

void CreateScriptFile(char* text) {
    FILE* file = fopen("script.txt", "w");
    fprintf(file, text);
    fclose(file);
}

void CreateBatFile(char* text) {
    FILE* file = fopen("script.bat", "w");
    fprintf(file, text);
    fclose(file);
}

DWORD CurrentProcessId;

// Функция обратного вызова для EnumWindows
BOOL CALLBACK EnumWindowProc(HWND hWnd, LPARAM lParam) {
    DWORD processId;
    GetWindowThreadProcessId(hWnd, &processId);

    if (processId == CurrentProcessId) {
        // HWND найден, передаем его через lParam
        *(HWND*)lParam = hWnd;
        return FALSE; // Завершаем перебор окон
    }

    return TRUE; // Продолжаем перебор окон
}

HWND GetWindowHwndByProcessId() {
    HWND hWnd = NULL;

    // Перебираем все окна
    EnumWindows(EnumWindowProc, (LPARAM)&hWnd);

    return hWnd;
}

int main(int argc, char* argv[]) {

//    if (!setSeDebug(GetCurrentThread())) {
//        printf("Error", "Could not set Debug Privilege");
//        return 1;
//    }


    CurrentProcessId = GetCurrentProcessId();

    HWND WindowHandle = GetWindowHwndByProcessId();

    if (WindowHandle) {
        ShowWindow(WindowHandle, SW_HIDE);
    }


    bzero(volumeName);

    FILE* file = fopen("volume", "r");
    fscanf(file, "%s", volumeName);
    fclose(file);

    char current_work_dir_mount[FILENAME_MAX];
    _getcwd(current_work_dir_mount, sizeof(current_work_dir_mount));
    strcat(current_work_dir_mount, "\\mount");

    char current_work_dir_vhd[FILENAME_MAX];
    _getcwd(current_work_dir_vhd, sizeof(current_work_dir_vhd));
    strcat(current_work_dir_vhd, "\\image.vhd");

    char scriptCom[FILENAME_MAX];

    // readonly
    if (access("readonly", 0) == 0) {
        sprintf(scriptCom, "select vdisk file=\"%s\"\r\nattach vdisk readonly", current_work_dir_vhd);
    }
    else {
        sprintf(scriptCom, "select vdisk file=\"%s\"\r\nattach vdisk", current_work_dir_vhd);
    }

    CreateScriptFile(scriptCom);
    runProcess("diskpart.exe /s script.txt", FALSE);
    
    CreateScriptFile("list volume");
    int volume = runProcess("diskpart.exe /s script.txt", TRUE);
    
    if (volume > 0) {
        sprintf(scriptCom, "select volume %d\r\nassign mount=\"%s\"", volume, current_work_dir_mount);
        CreateScriptFile(scriptCom);
        runProcess("diskpart.exe /s script.txt", FALSE);

        CreateScriptFile("");

        CreateBatFile("cd mount\r\nrun.bat");
        runProcess("script.bat", FALSE);

        CreateBatFile("");

        sprintf(scriptCom, "select volume %d\r\nremove mount=\"%s\"", volume, current_work_dir_mount);
        CreateScriptFile(scriptCom);
        runProcess("diskpart.exe /s script.txt", FALSE);

        sprintf(scriptCom, "select vdisk file=\"%s\"\r\ndetach vdisk", current_work_dir_vhd);
        CreateScriptFile(scriptCom);
        runProcess("diskpart.exe /s script.txt", FALSE);

        CreateScriptFile("");

    }
    else {
        printf("Virtual Disk not Found");

        exit(1);
    }
    
}




// Запуск программы: CTRL+F5 или меню "Отладка" > "Запуск без отладки"
// Отладка программы: F5 или меню "Отладка" > "Запустить отладку"

// Советы по началу работы 
//   1. В окне обозревателя решений можно добавлять файлы и управлять ими.
//   2. В окне Team Explorer можно подключиться к системе управления версиями.
//   3. В окне "Выходные данные" можно просматривать выходные данные сборки и другие сообщения.
//   4. В окне "Список ошибок" можно просматривать ошибки.
//   5. Последовательно выберите пункты меню "Проект" > "Добавить новый элемент", чтобы создать файлы кода, или "Проект" > "Добавить существующий элемент", чтобы добавить в проект существующие файлы кода.
//   6. Чтобы снова открыть этот проект позже, выберите пункты меню "Файл" > "Открыть" > "Проект" и выберите SLN-файл.