#include "interprocessSSP.h"
#include "connectorSSP.h"

EXISTING_RDP_SESSION remote_interactive_sessions[MAX_RDP_SESSIONS] = { 0 };

RDP_SESSION_ARRAY rdp_sessions = { 0 };


PSID convert_sidstring_to_sid(LPCWSTR sid_str) {
    PSID p_sid = NULL;
    if (!ConvertStringSidToSidW(sid_str, &p_sid)) {
        WCHAR function_name[64] = { 0 };
        MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));
        log_line(LOG_TYPE_WARNING,
            L"[%s] Blad ConvertStringSidToSidW: %lu", function_name, GetLastError()
        );
        return NULL;
    }
    return p_sid;
}

void get_LUID_string(const PLUID luid, PWSTR out, size_t out_len) {
    if (!luid) {
        StringCchCopyW(out, out_len, L"<NULL>");
        return;
    }
    StringCchPrintfW(out, out_len, L"%08x-%08x", luid->HighPart, luid->LowPart);
}

BOOL get_PIDs_from_sessionID(DWORD in_session_id, RELATED_PROCESSES *session_processes, BOOL update_flag) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnap == INVALID_HANDLE_VALUE) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad CreateToolhelp32Snapshot", function_name);
        return FALSE;
    }

    PROCESSENTRY32W pe;
    pe.dwSize = sizeof(pe);

    if (Process32FirstW(hSnap, &pe)) {
        do {
            /*
            *   w przypadku gdy procesow jest wiecej niz 5
            *   zrob realloc - powieksz 2 razy, powinno wystarczyc
            */
            if (session_processes->count >= INITIAL_PROCESSES_COUNT) {
                RELATED_PROCESS* new_processes = (RELATED_PROCESS*)HeapReAlloc(
                    GetProcessHeap(),
                    HEAP_ZERO_MEMORY,
                    session_processes->process,
                    INITIAL_PROCESSES_COUNT * 2 * sizeof(RELATED_PROCESS));

                if (!new_processes) {
                    log_line(LOG_TYPE_ERROR, L"[%s] Blad w HeapReAlloc", function_name);
                    return FALSE;
                }
                session_processes->process = new_processes;
            }
            DWORD procSessionId = 0;
            if (ProcessIdToSessionId(pe.th32ProcessID, &procSessionId)) {
                if (procSessionId == in_session_id) {
                    // jesli PID jest jednym z istotnych procesow systemowych - NIE UBIJAJ
                    if (CompareStringOrdinal(pe.szExeFile, -1, L"services.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"smss.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"wininit.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"lsass.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"csrss.exe", -1, TRUE) == CSTR_EQUAL)
                    {
                        continue;
                    }
                    if (update_flag) {
                        for (DWORD i = 0; i < session_processes->count; i++) {
                            if (CompareStringOrdinal(pe.szExeFile, -1, 
                                session_processes->process[session_processes->count].process_name, 
                                -1, TRUE) == CSTR_EQUAL && 
                                session_processes->process[session_processes->count].pid == pe.th32ProcessID) 
                            {
                                log_line(LOG_TYPE_DEBUG,
                                    L"[%s] Dla sesji ID = %lu, proces = %s, PID = %lu. juz istnieje...",
                                    function_name,
                                    session_processes->process[session_processes->count].process_name,
                                    session_processes->process[session_processes->count].pid
                                );
                                continue;
                            }
                            session_processes->process[session_processes->count].pid = pe.th32ProcessID;
                            session_processes->process[session_processes->count].pid;
                            StringCchCopyW(session_processes->process[session_processes->count].process_name,
                                ARRAYSIZE(session_processes->process[session_processes->count].process_name),
                                pe.szExeFile);
                            log_line(LOG_TYPE_DEBUG,
                                L"[%s] Wykryto dodatkowy proces dla sesji = %s, PID = %lu.",
                                function_name,
                                session_processes->process[session_processes->count].process_name,
                                session_processes->process[session_processes->count].pid
                            );
                            session_processes->count++;
                        }
                        
                    }
                    else {
                        session_processes->process[session_processes->count].pid = pe.th32ProcessID;
                        session_processes->process[session_processes->count].pid;
                        StringCchCopyW(session_processes->process[session_processes->count].process_name,
                            ARRAYSIZE(session_processes->process[session_processes->count].process_name),
                            pe.szExeFile);
                        log_line(LOG_TYPE_DEBUG,
                            L"[%s] Wykryto proces dla sesji = %s, PID = %lu.",
                            function_name,
                            session_processes->process[session_processes->count].process_name,
                            session_processes->process[session_processes->count].pid
                        );
                        session_processes->count++;
                    }
                }
            }
        } while (Process32NextW(hSnap, &pe));
    }

    return TRUE;
}

BOOL get_PID_from_SessionID(DWORD in_session_id, DWORD *pid, WCHAR *pid_exe_name) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    DWORD processes_count = 0;

    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnap == INVALID_HANDLE_VALUE) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad CreateToolhelp32Snapshot", function_name);
        return FALSE;
    }

    PROCESSENTRY32W pe;
    pe.dwSize = sizeof(pe);

    if (Process32FirstW(hSnap, &pe)) {
        do {
            DWORD procSessionId = 0;
            if (ProcessIdToSessionId(pe.th32ProcessID, &procSessionId)) {
                if (procSessionId == in_session_id) {
                    // jesli PID jest jednym z istotnych procesow systemowych - NIE UBIJAJ
                    if (CompareStringOrdinal(pe.szExeFile, -1, L"services.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"smss.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"wininit.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"lsass.exe", -1, TRUE) == CSTR_EQUAL ||
                        CompareStringOrdinal(pe.szExeFile, -1, L"csrss.exe", -1, TRUE) == CSTR_EQUAL)
                    {
                        continue;
                    }
                    *pid = pe.th32ProcessID;
                    StringCchCopyW(pid_exe_name, 64, pe.szExeFile);
                    processes_count++;
                    break;
                }
            }
        } while (Process32NextW(hSnap, &pe));
    }
    if (processes_count > 1) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Odnaleziono %lu procesow.",
            function_name,
            processes_count
        );
    }
    log_line(LOG_TYPE_DEBUG,
        L"[%s] Skorelowano sesje RDP z procesem: %s, PID: %lu, ID sesji RDP: %lu",
        function_name,
        pid_exe_name,
        *pid,
        in_session_id
    );
    return TRUE;
}
/*
FIND_SESSION_STATUS find_remote_domain_user_session(PSID user_sid, PUNICODE_STRING domain_username, DWORD *out_session_id) {
    PWTS_SESSION_INFO session_info = NULL;
    DWORD session_count = 0;

    WCHAR function_name[64] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    if (!WTSEnumerateSessionsW(WTS_CURRENT_SERVER_HANDLE, 0, 1, &session_info, &session_count)) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad pobierania sesji WTSEnumerateSessions", function_name);
        return FIND_SESSION_ERROR;
    }

    BOOL is_rdp_session_empty = TRUE;
    for (int i = 0; i < MAX_RDP_SESSIONS; i++) {
        if (remote_interactive_sessions[i].active) {
            is_rdp_session_empty = FALSE;
        }
    }
    WCHAR sid_string[128] = { 0 };
    check_SID(user_sid, sid_string, ARRAYSIZE(sid_string));

    if (!is_rdp_session_empty) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Wykryto istniejace sesje RDP...", function_name);
        for (int i = 0; i < MAX_RDP_SESSIONS; i++) {
            if (compare_unicode_with_wchar(domain_username, remote_interactive_sessions[i].domain_username)
                && wcscmp(sid_string, remote_interactive_sessions[i].user_sid) == 0)
            {
                
                for (DWORD j = 0; j < session_count; j++) {
                    DWORD session_Id = session_info[j].SessionId;
                    USHORT* protocol = NULL;
                    DWORD bytesReturned = 0;

                    if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_Id, WTSClientProtocolType, (LPWSTR*)&protocol, &bytesReturned)) {
                        if (*protocol == WTS_PROTOCOL_RDP) {
                            
                            WTS_CLIENT_ADDRESS* address_ptr = NULL;
                            if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_Id, WTSClientAddress, (LPWSTR*)&address_ptr, &bytesReturned)) {

                                WCHAR client_ip[INET_ADDRSTRLEN] = { 0 };

                                struct sockaddr_in sa;
                                ZeroMemory(&sa, sizeof(sa));
                                sa.sin_family = AF_INET;
                                CopyMemory(&sa.sin_addr, &address_ptr->Address[2], 4);

                                if (!InetNtopW(AF_INET, &sa.sin_addr, client_ip, INET_ADDRSTRLEN)) {
                                    log_line(LOG_TYPE_ERROR, L"[%s] Blad konwersji IP w InetNtopW", function_name);
                                    WTSFreeMemory(address_ptr);
                                    WTSFreeMemory(protocol);
                                    return FIND_SESSION_ERROR;
                                }
                                /*
                                log_line(LOG_TYPE_DEBUG,
                                    L"[%s] Przed konwersja IP...", function_name);
                                    

                                if (wcscmp(client_ip, remote_interactive_sessions[i].ip_address) != 0) {
                                    WTSFreeMemory(address_ptr);
                                    WTSFreeMemory(protocol);
                                    return FIND_SESSION_ERROR;
                                }
                                /*log_line(LOG_TYPE_DEBUG,
                                    L"[%s] Po konwersji IP...", function_name);

                                WCHAR user_name[64] = { 0 };
                                PWTSCLIENTW pClient = NULL;
                                DWORD bytes = 0;

                                if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_Id, WTSClientInfo, (LPWSTR*)&pClient, &bytes) && pClient) {
                                    log_line(LOG_TYPE_DEBUG,
                                        L"[%s] Dla konta:    %s, z IP:   %s, SID:   %s juz istnieje zapoczatkowana sesja RDP!",
                                        function_name,
                                        pClient->UserName,
                                        client_ip,
                                        sid_string);
                                    WTSFreeMemory(pClient);
                                }
                                WTSFreeMemory(address_ptr);
                                WTSFreeMemory(protocol);
                                
                                return FIND_SESSION_FOUND;
                            }
                        }
                        WTSFreeMemory(protocol);
                    }
                }
            }
        }
    }

    for (DWORD i = 0; i < session_count; i++) {
        DWORD session_Id = session_info[i].SessionId;

        USHORT* protocol = NULL;
        DWORD bytesReturned = 0;

        if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_Id, WTSClientProtocolType, (LPWSTR*)&protocol, &bytesReturned)) {
            if (*protocol == WTS_PROTOCOL_RDP) {
                LPWSTR data_buffer = NULL;

                WCHAR user_name[64] = { 0 };
                WCHAR domain_name[64] = { 0 };
                WCHAR client_name[64] = { 0 };
                WCHAR protocol_type[16] = { 0 };
                WCHAR client_ip[INET_ADDRSTRLEN] = { 0 };

                WTS_CLIENT_ADDRESS* address_ptr = NULL;
                
                copy_lpwstr_string(remote_protocol_type_to_string(*protocol), protocol_type, ARRAYSIZE(protocol_type));

                retrieve_session_data(session_Id, WTSUserName, user_name, ARRAYSIZE(user_name));
                retrieve_session_data(session_Id, WTSDomainName, domain_name, ARRAYSIZE(domain_name));
                retrieve_session_data(session_Id, WTSClientName, client_name, ARRAYSIZE(client_name));

                if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_Id, WTSClientAddress, (LPWSTR*)&address_ptr, &bytesReturned)) {
                    if (address_ptr->AddressFamily == AF_INET) {
                        struct sockaddr_in sa;
                        ZeroMemory(&sa, sizeof(sa));
                        sa.sin_family = AF_INET;
                        CopyMemory(&sa.sin_addr, &address_ptr->Address[2], 4);

                        if (!InetNtopW(AF_INET, &sa.sin_addr, client_ip, INET_ADDRSTRLEN)) {
                            log_line(LOG_TYPE_ERROR, L"[%s] Blad konwersji IP w InetNtopW", function_name);
                            WTSFreeMemory(protocol);
                            WTSFreeMemory(address_ptr);
                            return FIND_SESSION_ERROR;
                        }
                    }
                    WTSFreeMemory(address_ptr);
                }

                PWTSCLIENTW pClient = NULL;
                DWORD bytes = 0;

                if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_Id, WTSClientInfo, (LPWSTR*)&pClient, &bytes) && pClient) {
                    log_line(LOG_TYPE_DEBUG,
                        L"[%s] Przechwycono sesje RemoteInteractive:  Nazwa konta:   %s, Domena:   %s, Nazwa maszyny:   %s, IP:   %s, Protokol:   %s",
                        function_name,
                        pClient->UserName,
                        pClient->Domain,
                        pClient->ClientName,
                        client_ip,
                        protocol_type);
                }

                for (int i = 0; i < MAX_RDP_SESSIONS; i++) {
                    if (!remote_interactive_sessions[i].active) {
                        remote_interactive_sessions[i].active = TRUE;
                        StringCchCopyW(remote_interactive_sessions[i].user_sid, ARRAYSIZE(remote_interactive_sessions[i].user_sid), sid_string);
                        StringCchCopyW(remote_interactive_sessions[i].domain_username, ARRAYSIZE(remote_interactive_sessions[i].domain_username), pClient->UserName);
                        StringCchCopyW(remote_interactive_sessions[i].ip_address, ARRAYSIZE(remote_interactive_sessions[i].ip_address), client_ip);

                        log_line(LOG_TYPE_DEBUG, L"[%s] Dodano sesje RDP do kolejki... Dane sesji: User SID: %s, Nazwa konta: %s, IP: %s", 
                            function_name,
                            remote_interactive_sessions[i].user_sid,
                            remote_interactive_sessions[i].domain_username,
                            remote_interactive_sessions[i].ip_address);
                        WTSFreeMemory(pClient);
                        WTSFreeMemory(protocol);
                        WTSFreeMemory(session_info);

                        *out_session_id = session_Id;

                        return FIND_SESSION_NOT_FOUND;
                    }
                }
                WTSFreeMemory(pClient);
            }
            WTSFreeMemory(protocol);
        }
    }
    WTSFreeMemory(session_info);
    return FIND_SESSION_ERROR;
}
*/
BOOL retrieve_session_data(DWORD session_id, WTS_INFO_CLASS info, WCHAR* out_buff, size_t out_size) {
    LPWSTR data_buffer = NULL;
    DWORD bytes_returned = 0;

    if (!WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_id, info, &data_buffer, &bytes_returned)) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad przy pozyskiwaniu WTS_INFO_CLASS, kod: %lu...", 
            L"retrieve_session_data", 
            DWORD(info));
        return FALSE;
    }

    if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_id, info, &data_buffer, &bytes_returned)) {
        if (data_buffer != NULL && bytes_returned > sizeof(WCHAR)) {
            copy_lpwstr_string(data_buffer, out_buff, out_size);
        }
        else {
            copy_lpwstr_string((LPWSTR)NULL, out_buff, out_size);
        }
        WTSFreeMemory(data_buffer);
        return TRUE;
    }
    
    copy_lpwstr_string((LPWSTR)NULL, out_buff, out_size);
    return TRUE;
}

BOOL convert_ip_addr_to_string(WTS_CLIENT_ADDRESS *ip, WCHAR *ip_data) {
    if (ip->AddressFamily == AF_INET) {
        struct sockaddr_in sa;
        ZeroMemory(&sa, sizeof(sa));
        sa.sin_family = AF_INET;
        CopyMemory(&sa.sin_addr, &ip->Address[2], 4);

        if (!InetNtopW(AF_INET, &sa.sin_addr, ip_data, INET_ADDRSTRLEN)) {
            log_line(LOG_TYPE_ERROR, L"[%s] Blad konwersji IP w InetNtopW", L"convert_ip_addr_to_string");
            return FALSE;
        }
    }
    else {
        WCHAR buffer[512] = { 0 };
        WCHAR temp[8];

        for (int i = 0; i < 20; i++) {
            StringCchPrintfW(temp, ARRAYSIZE(temp), L"%02X ", ip->Address[i]);
            StringCchCatW(buffer, ARRAYSIZE(buffer), temp);
        }

        log_line(LOG_TYPE_WARNING, L"[%s] Nieobslugiwany format adresu, kod: %lu, zawartosc addressFamily: %s...", 
            L"convert_ip_addr_to_string", 
            ip->AddressFamily,
            buffer);
        return FALSE;
    }
    return TRUE;
}


void print_kerberos_module_functions(HMODULE kerberos_module) {
    WCHAR function_name[64] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    BYTE* base_address = (BYTE*)kerberos_module;
    IMAGE_DOS_HEADER* dos_header = (IMAGE_DOS_HEADER*)base_address;

    if (dos_header->e_magic != IMAGE_DOS_SIGNATURE)
        return;

    IMAGE_NT_HEADERS* ntHeaders = (IMAGE_NT_HEADERS*)(base_address + dos_header->e_lfanew);
    if (ntHeaders->Signature != IMAGE_NT_SIGNATURE) 
        return;

    IMAGE_DATA_DIRECTORY exportDir = ntHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    if (exportDir.VirtualAddress == 0) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Brak eksportow w kerberos.dll",
            function_name
        );
        return;
    }
    IMAGE_EXPORT_DIRECTORY* exportDirectory = (IMAGE_EXPORT_DIRECTORY*)(base_address + exportDir.VirtualAddress);

    DWORD* namesRVA = (DWORD*)(base_address + exportDirectory->AddressOfNames);
    WORD* ordinals = (WORD*)(base_address + exportDirectory->AddressOfNameOrdinals);
    DWORD* functions = (DWORD*)(base_address + exportDirectory->AddressOfFunctions);

    log_line(LOG_TYPE_DEBUG,
        L"[%s] Lista eksportowanych funkcji w kerberos.dll: ",
        function_name
    );

    for (DWORD i = 0; i < exportDirectory->NumberOfNames; i++) {
        char* funcName = (char*)(base_address + namesRVA[i]);

        WCHAR funcNameW[256];
        MultiByteToWideChar(CP_ACP, 0, funcName, -1, funcNameW, ARRAYSIZE(funcNameW));

        WORD ordinal = ordinals[i];
        DWORD funcRVA = functions[ordinal];
        void* funcAddress = base_address + funcRVA;

        log_line(LOG_TYPE_DEBUG,
            L"[%s] %s at ordinal %d, adres: %p",
            function_name,
            funcNameW,
            ordinal + exportDirectory->Base,
            funcAddress
        );
    }
}

void test_load_library(LPCWSTR dll_name, LPCWSTR rust_library) {
    HMODULE h_lib = LoadLibraryW(dll_name);
    if (h_lib) {
        log_line(LOG_TYPE_INFO,
            L"[%s] Zaladowano modul DLL %s",
            L"test_load_library",
            dll_name
        );
        typedef int (*PFN_Run)();
        PFN_Run pRun = (PFN_Run)GetProcAddress(h_lib, "Run");
        if (pRun) {
            int result = pRun();
            log_line(LOG_TYPE_INFO,
                L"[%s]rezultat z zaladowanej DLL: %d",
                L"test_load_library",
                result
            );
        }
        FreeLibrary(h_lib);
    }
    else {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad w LoadLibraryW: %lu", L"test_load_library", GetLastError());
    }
    HMODULE h_rust_lib = LoadLibraryW(rust_library);

    if (h_rust_lib) {
        log_line(LOG_TYPE_INFO,
            L"[%s] Zaladowano DLL stworzona w Rust %s",
            L"test_load_library",
            rust_library
        );
        typedef int (*PFN_Run)();
        PFN_Run rust_run = (PFN_Run)GetProcAddress(h_rust_lib, "Run");
        if (rust_run) {
            int result = rust_run();
            log_line(LOG_TYPE_INFO,
                L"[%s]rezultat z biblioteki w Rust: %d",
                L"test_load_library",
                result
            );
        }
        FreeLibrary(h_rust_lib);
    }
    else {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad w LoadLibraryW dla biblioteki Rust: %lu", L"test_load_library", GetLastError());
    }
}


BOOL initialize_rdp_sessions_array() {
    rdp_sessions.session_data = (RDP_SESSION_DATA*)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY,
        INITIAL_RDP_SESSIONS * sizeof(RDP_SESSION_DATA));

    if (!rdp_sessions.session_data) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad alokacji pamieci dla sesji RDP", L"initialize_rdp_sessions_array");
        return FALSE;
    }

    rdp_sessions.session_count = 0;
    rdp_sessions.capacity = INITIAL_RDP_SESSIONS;

    log_line(LOG_TYPE_INFO, L"[%s] Poprawnie zaalokowano pamiec dla sesji RDP...", L"initialize_rdp_sessions_array");

    return TRUE;
}

void free_rdp_sessions_array() {
    if (rdp_sessions.session_data) {
        if (rdp_sessions.session_data->processes) {
            HeapFree(GetProcessHeap(), 0, rdp_sessions.session_data->processes);
            log_line(LOG_TYPE_INFO, L"[%s] Zwalnianie pamieci po sesjach RDP - procesy...", L"free_rdp_sessions_array");
        }
        HeapFree(GetProcessHeap(), 0, rdp_sessions.session_data);
        log_line(LOG_TYPE_INFO, L"[%s] Zwalnianie pamieci po sesjach RDP...", L"free_rdp_sessions_array");
    }
    ZeroMemory(&rdp_sessions, sizeof(rdp_sessions));
}

MATCH_SESSION_STATUS match_existing_rdp_sessions(DWORD in_session_id, DWORD *out_session_id) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    log_line(LOG_TYPE_INFO,
        L"[%s] ilosc sesji RDP = %lu",
        function_name,
        rdp_sessions.session_count);

    if (rdp_sessions.session_count == 0) {
        //utworz nowa sesje
        log_line(LOG_TYPE_INFO, 
            L"[%s] Wykryto nowa sesje RDP - Brak aktywnych sesji RDP uzytkownikow domenowych, dodawanie nowej...", 
            function_name);
        *out_session_id = in_session_id;

        return SESSION_CREATE_NEW;
    }
    if (rdp_sessions.session_count != 0) {
        for (DWORD i = 0; i < rdp_sessions.session_count; i++) {
            RDP_SESSION_DATA* sess = &rdp_sessions.session_data[i];

            /*  
            *   w przypadku gdy id sesji zgadza sie z juz istniejaca :
            *   sprawdz czy jest valid - jesli jest valid, to znaczy ze MFA 
            *   zostalo we wczesniejszej zaakceptowane i zaktualizuj proces wazny dla sesji RDP
            *   jesli nie zostalo zaakceptowane - ubij procesy
            */

            if (sess->session_id == in_session_id) {
                *out_session_id = sess->session_id;
                if (sess->valid) {
                    log_line(LOG_TYPE_INFO,
                        L"[%s] Wykryto istniejaca sesje RDP, ID sesji = %lu, aktualizacja stanu polaczenia...",
                        function_name,
                        in_session_id);
                    return SESSION_UPDATE_EXISTING;
                }
                else { 
                    /*
                    DWORD pid = 0;
                    WCHAR proc_name[128] = { 0 };

                    DWORD processes_count = 0;

                    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
                    if (hSnap == INVALID_HANDLE_VALUE) {
                        log_line(LOG_TYPE_ERROR, L"[%s] Blad CreateToolhelp32Snapshot", function_name);
                        //return FALSE;
                    }

                    PROCESSENTRY32W pe;
                    pe.dwSize = sizeof(pe);

                    if (Process32FirstW(hSnap, &pe)) {
                        do {
                            DWORD procSessionId = 0;
                            if (ProcessIdToSessionId(pe.th32ProcessID, &procSessionId)) {
                                if (procSessionId == in_session_id) {
                                    // jesli PID jest jednym z istotnych procesow systemowych - NIE UBIJAJ
                                    if (CompareStringOrdinal(pe.szExeFile, -1, L"services.exe", -1, TRUE) == CSTR_EQUAL ||
                                        CompareStringOrdinal(pe.szExeFile, -1, L"smss.exe", -1, TRUE) == CSTR_EQUAL ||
                                        CompareStringOrdinal(pe.szExeFile, -1, L"wininit.exe", -1, TRUE) == CSTR_EQUAL ||
                                        CompareStringOrdinal(pe.szExeFile, -1, L"lsass.exe", -1, TRUE) == CSTR_EQUAL ||
                                        CompareStringOrdinal(pe.szExeFile, -1, L"csrss.exe", -1, TRUE) == CSTR_EQUAL)
                                    {
                                        continue;
                                    }
                                    //processes_count++;
                                    log_line(LOG_TYPE_DEBUG,
                                        L"[%s] Odnaleziono proces: %s, PID: %lu",
                                        function_name,
                                        pe.szExeFile,
                                        pe.th32ProcessID
                                    );
                                }
                            }
                        } while (Process32NextW(hSnap, &pe));
                    }
                    if (processes_count > 1) {
                        log_line(LOG_TYPE_DEBUG,
                            L"[%s] Odnaleziono %lu procesow.",
                            function_name,
                            processes_count
                        );
                    }
                    */
                    log_line(LOG_TYPE_INFO,
                        L"[%s] Sesja RDP ID = %lu juz istnieje i zostala odrzucona...",
                        function_name,
                        in_session_id);
                    return SESSION_TERMINATE_EXISTING;
                }
            }
        }
    }
    log_line(LOG_TYPE_ERROR,
        L"[%s] Wystapil blad sesji RDP, ID sesji: %lu, sesja o nieznanym statusie...",
        function_name,
        in_session_id);
    return SESSION_STATUS_UNKNOWN;
}

BOOL create_new_rdp_session(PSID user_psid, DWORD session_id) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    RDP_SESSION_DATA new_session = { 0 };

    new_session.processes = (RELATED_PROCESSES*)HeapAlloc(
        GetProcessHeap(),
        HEAP_ZERO_MEMORY,
        sizeof(RELATED_PROCESSES));

    if (!new_session.processes) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad alokacji pamieci dla sesji RDP - alokacja danych procesu", function_name);
        return FALSE;
    }

    new_session.processes->process = (RELATED_PROCESS*)HeapAlloc(
        GetProcessHeap(),
        HEAP_ZERO_MEMORY,
        INITIAL_PROCESSES_COUNT * sizeof(RELATED_PROCESS));
    if (!new_session.processes->process) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad alokacji pamieci dla sesji RDP - alokacja danych PID", function_name);
        HeapFree(GetProcessHeap(), 0, new_session.processes);
        return FALSE;
    }
    new_session.processes->count = 0;

    //log_line(LOG_TYPE_DEBUG, L"[%s] processes ptr: %p", function_name, new_session.processes);
    //log_line(LOG_TYPE_DEBUG, L"[%s] processes->process ptr: %p", function_name, new_session.processes->process);
    //log_line(LOG_TYPE_DEBUG, L"[%s] proc_name ptr: %p", function_name, proc_name);

    if (!retrieve_rdp_session_info(session_id, &new_session, user_psid)) {
        log_line(LOG_TYPE_ERROR,
            L"[%s] Nie utworzono nowej sesji RDP...",
            function_name);
        return FALSE;
    }

    //zinkrementuj licznik sesji
    rdp_sessions.session_count++;

    log_line(LOG_TYPE_INFO,
        L"[%s] Przechwycono sesje %s  --  SID konta: %s, IP maszyny: %s, Nazwa konta: %s, Nazwa DNS: %s, Nazwa klienta: %s, ID sesji: %lu",
        function_name,
        L"RemoteInteractive",
        new_session.user_sid,
        new_session.client_ip,
        new_session.user_name,
        new_session.domain_name,
        new_session.client_name,
        new_session.session_id
    );

    /*
    *   dane z sesji uzyskane - teraz czas na mfa
    *   najpierw przygotuj bufor JSON do wyslania do serwera
    */

    char buff[192] = { 0 };
    format_data_for_connection(new_session.session_id, buff);
    NTSTATUS mfa_result = manage_mfa_session(buff);

    /*
    *   dalsza obsluga sesji zalezy od rezultatu mfa
    *   jesli nastapi zaakceptowanie polaczenia - ustaw flage valid na true
    *   ponadto - dodaj informacje o procesie ktory kontroluje RDP
    *   jesli nastapi odrzucenie w mfa - ustaw flage valid na false i ubij polaczenie i nie dodawaj sesji do listy
    *   kazdy status inny niz STATUS_SUCCESS (mozliwe jest np. zerwanie polaczenia w trakcie negocjacji tcp, blad w http, itp.
    *   powoduje zamkniecie sesji)
    */

    if (mfa_result != STATUS_SUCCESS) {
        log_line(LOG_TYPE_ERROR,
            L"[%s] MFA nie powiodlo sie, odrzucanie polaczenia...",
            function_name);
        new_session.valid = FALSE;
        /*
        if (get_PIDs_from_sessionID(session_id, new_session.processes, FALSE)) {
            //terminate_remaining_processes(new_session.processes, session_id);
        }*/
    } else {
        log_line(LOG_TYPE_ERROR,
            L"[%s] MFA zaakceptowane, procedowanie polaczenia...",
            function_name);
        new_session.valid = TRUE; 
    }

    if (get_PIDs_from_sessionID(session_id, new_session.processes, FALSE)) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Dodano procesy zwiazane z sesja RDP...",
            function_name
        );
    }
    /*
    *   Procesowanie sesji, dodaj sesje do listy
    */
    if (add_session_to_list(&new_session)) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Dodano poprawnie sesje do listy, ID sesji RDP = %lu...",
            function_name,
            new_session.session_id
        );
    }

    return TRUE;
}

BOOL update_existing_rdp_session(DWORD session_id) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    for (DWORD i = 0; i < rdp_sessions.session_count; i++) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Liczba sesji RDP: %lu",
            function_name,
            rdp_sessions.session_count
        );
        RDP_SESSION_DATA* sess = &rdp_sessions.session_data[i];
        if (sess->session_id == session_id) {
            /*
            *   sprawdzam czy dla sesji sa juz zarejestrowane procesy
            */

            if (get_PIDs_from_sessionID(session_id, sess->processes, TRUE))
            {

            }

            if (!sess->valid) {
                terminate_remaining_processes(sess->processes, session_id);
            }
        }
    }
    return TRUE;
}

BOOL terminate_remaining_processes(RELATED_PROCESSES* session_processes, DWORD sess_id) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    /*
    *   zakoncz procesy dla danej sesji
    */
    if (session_processes->count != 0) {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Konczenie procesow dla sesji o ID = %lu, ilosc procesow: %lu",
            function_name,
            sess_id,
            session_processes->count
        );
        for (DWORD i = 0; i < session_processes->count; i++) {
            RELATED_PROCESS proc = session_processes->process[i];
            DWORD proc_id = proc.pid;
            if (WTSTerminateProcess(WTS_CURRENT_SERVER_HANDLE, proc_id, 1)) {
                log_line(LOG_TYPE_DEBUG,
                    L"[%s] Zakonczono - %s -> PID: %lu",
                    function_name,
                    proc.process_name,
                    proc_id
                );
            }
            else {
                log_line(LOG_TYPE_ERROR,
                    L"[%s] Blad WTSTerminateProcess...",
                    function_name
                );
                return FALSE;
            }
        }
    }
    else {
        log_line(LOG_TYPE_DEBUG,
            L"[%s] Nie odnaleziono procesow dla sesji RDP o ID = %lu",
            function_name
        );
        return FALSE;
    }
    return TRUE;
}

BOOL add_session_to_list(const RDP_SESSION_DATA *session) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    if (!rdp_sessions.session_data) {
        log_line(LOG_TYPE_ERROR,
            L"[%s] Blad w dostepie do listy sesji RDP...",
            function_name);
        return FALSE;
    }
    if (rdp_sessions.session_count >= rdp_sessions.capacity) {
        DWORD new_capacity = rdp_sessions.capacity * 2;
        RDP_SESSION_DATA* resized = (RDP_SESSION_DATA*)HeapReAlloc(
            GetProcessHeap(),
            HEAP_ZERO_MEMORY,
            rdp_sessions.session_data,
            new_capacity * sizeof(RDP_SESSION_DATA));
        if (!resized) {
            log_line(LOG_TYPE_ERROR,
                L"[%s] Nie udalo sie zrealokowac pamieci, blad HeapReAlloc...",
                function_name);
            return FALSE;
        }

        rdp_sessions.session_data = resized;
        rdp_sessions.capacity = new_capacity;
    }

    rdp_sessions.session_data[rdp_sessions.session_count++] = *session;

    return TRUE;
}

BOOL remove_session_from_list(DWORD sess_id) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    for (DWORD i = 0; i < rdp_sessions.session_count; i++) {
        RDP_SESSION_DATA* sess = &rdp_sessions.session_data[i];
        if (sess->session_id == sess_id) {
            log_line(LOG_TYPE_DEBUG,
                L"[%s] Usuwanie sesji RDP: %lu z listy --> PSID: %s, IP: %s, Nazwa konta: %s...",
                function_name,
                sess_id,
                sess->user_sid,
                sess->client_ip,
                sess->domain_name
            );
            if (!terminate_remaining_processes(sess->processes, sess_id)) {
                return FALSE;
            }

            /*
            *   procesy ubite, wyczysc dane po sesji
            */
            if (sess->processes) {
                if (sess->processes->process) {
                    log_line(LOG_TYPE_DEBUG, L"[%s] Wyczyszczono sess->processes->process...", function_name);
                    HeapFree(GetProcessHeap(), 0, sess->processes->process);
                }
                log_line(LOG_TYPE_DEBUG, L"[%s] Wyczyszczono sess->processes...", function_name);
                HeapFree(GetProcessHeap(), 0, sess->processes);
            }

            rdp_sessions.session_count--;
            break;
        }
    }
    return TRUE;
}

BOOL retrieve_rdp_session_info(DWORD session_id, RDP_SESSION_DATA *session_data, PSID user_psid) {
    WCHAR function_name[40] = { 0 };
    MultiByteToWideChar(CP_ACP, 0, __FUNCTION__, -1, function_name, ARRAYSIZE(function_name));

    if (!session_data) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad przy pozyskiwaniu danych nowej sesji RDP...", function_name);
        return FALSE;
    }
    //ZeroMemory(session_data, sizeof(RDP_SESSION_DATA));

    DWORD bytesReturned = 0;

    session_data->session_id = session_id;

    session_data->protocol_type = WTS_PROTOCOL_RDP;

    check_SID(user_psid, session_data->user_sid, ARRAYSIZE(session_data->user_sid));

    if (!retrieve_session_data(session_id, WTSUserName, session_data->user_name, ARRAYSIZE(session_data->user_name)))
        return FALSE;

    if (!retrieve_session_data(session_id, WTSDomainName, session_data->domain_name, ARRAYSIZE(session_data->domain_name)))
        return FALSE;

    if (!retrieve_session_data(session_id, WTSClientName, session_data->client_name, ARRAYSIZE(session_data->client_name)))
        return FALSE;

    WTS_CLIENT_ADDRESS* address_ptr = NULL;
    if (!WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session_id, WTSClientAddress, (LPWSTR*)&address_ptr, &bytesReturned)) {
        log_line(LOG_TYPE_ERROR, L"[%s] Blad przy pozyskiwaniu adresu IP...", function_name);
        return FALSE;
    }

    if (!convert_ip_addr_to_string(address_ptr, session_data->client_ip))
        return FALSE;


    //kiedys trzeba dodac szersza obsluge Console

    return TRUE;
}

void format_data_for_connection(DWORD session_id, char *buffer) {
    char dns[64] = { 0 };
    char user[64] = { 0 };
    for (DWORD i = 0; i < rdp_sessions.session_count; i++) {
        RDP_SESSION_DATA* sess = &rdp_sessions.session_data[i];
        if (sess->session_id == session_id) {
            WideCharToMultiByte(CP_UTF8, 0, sess->domain_name, -1, dns, sizeof(dns), NULL, NULL);
            WideCharToMultiByte(CP_UTF8, 0, sess->user_name, -1, user, sizeof(user), NULL, NULL);
        }
    }
    StringCchPrintfA(buffer, 192, "{\"Domena\":\"%s\",\"Nazwa konta\":\"%s\"}", dns, user);
}