#include "domainUtilsSSP.h"

#include "utilsSSP.h"

DOMAIN_INFO domain_info = { 0 };

USER_INFO* domain_users = NULL;

LONG domain_monitor_thread_counter = 0;


HRESULT get_domain_info() {
	DSROLE_PRIMARY_DOMAIN_INFO_BASIC* domain_role_info = NULL;

	DWORD status = DsRoleGetPrimaryDomainInformation(NULL, DsRolePrimaryDomainInfoBasic, (PBYTE*)&domain_role_info);
	if (status == ERROR_SUCCESS) {
		domain_info.is_domain_controller = (domain_role_info->MachineRole == DsRole_RolePrimaryDomainController ||
			domain_role_info->MachineRole == DsRole_RoleBackupDomainController);
		domain_info.is_machine_domain_joined = (domain_role_info->MachineRole != DsRole_RoleStandaloneWorkstation &&
			domain_role_info->MachineRole != DsRole_RoleStandaloneServer);
		StringCchCopyW(domain_info.NetBIOS_name, ARRAYSIZE(domain_info.NetBIOS_name), domain_role_info->DomainNameFlat);
		StringCchCopyW(domain_info.DNS_domain_name, ARRAYSIZE(domain_info.DNS_domain_name), domain_role_info->DomainNameDns);
		DsRoleFreeMemory(domain_role_info);
	}
	else {
		return HRESULT_FROM_WIN32(status);
	}

	SID_NAME_USE pe_use;
	BYTE sid_buffer[SECURITY_MAX_SID_SIZE] = { 0 };
	DWORD sid_size = sizeof(sid_buffer);
	WCHAR referenced_domain[MAX_PATH];
	DWORD domain_size = ARRAYSIZE(referenced_domain);

	if (!LookupAccountNameW(NULL, domain_info.DNS_domain_name, sid_buffer, &sid_size,
		referenced_domain, &domain_size, &pe_use
	)) {
		DWORD err = GetLastError();
		log_line(LOG_TYPE_ERROR, L"[%s] Blad LookupAccountNameW, kod bledu: %lu", L"get_domain_info", GetLastError());
		return HRESULT_FROM_WIN32(err);
	}

	DWORD len = GetLengthSid(sid_buffer);
	domain_info.domain_SID = (PSID)LocalAlloc(LMEM_FIXED, len);
	if (!domain_info.domain_SID) {
		log_line(LOG_TYPE_ERROR, L"[%s] Blad alokacji pamieci dla SID", L"get_domain_info");
		return E_OUTOFMEMORY;
	}
	if (!CopySid(len, domain_info.domain_SID, sid_buffer)) {
		DWORD err = GetLastError();
		LocalFree(domain_info.domain_SID);
		domain_info.domain_SID = NULL;
		log_line(LOG_TYPE_ERROR, L"[%s] Blad w kopiowaniu SID", L"get_domain_info");
		return HRESULT_FROM_WIN32(err);
	}
	return S_OK;
}

HRESULT get_domain_users() {
	DWORD total_users = 0;
	DWORD read = 0;

	WCHAR function_name[] = L"get_domain_users";

	NET_API_STATUS status = NetUserEnum(
		NULL, 0, FILTER_NORMAL_ACCOUNT,
		NULL, 0, &read, &total_users, NULL
	);
	if (status != ERROR_SUCCESS && status != ERROR_MORE_DATA) {
		log_line(LOG_TYPE_ERROR, L"[%s] Blad NetUserEnum, kod bledu: %lu", function_name, status);
		return HRESULT_FROM_WIN32(status);
	}
	if (total_users == 0) {
		log_line(LOG_TYPE_WARNING, L"[%s] Nie odnaleziono uzytkownikow domeny...", function_name);
		return S_OK;
	}

	HANDLE hHeap = GetProcessHeap();
	domain_users = (USER_INFO*)HeapAlloc(hHeap, HEAP_ZERO_MEMORY, sizeof(USER_INFO) * total_users);
	if (!domain_users)
	{
		log_line(LOG_TYPE_ERROR, L"[%s] Blad alokacji pamieci dla USER_INFO", function_name);
		return E_OUTOFMEMORY;
	}

	USER_INFO_0* user_info_buff = NULL;
	DWORD resume = 0, read_sec = 0, total_users_sec = 0;
	DWORD user_index = 0;
	int regular_users = 0;
	int system_users = 0;

	DWORD start = GetTickCount();

	do
	{
		status = NetUserEnum(NULL, 0, FILTER_NORMAL_ACCOUNT, (LPBYTE*)&user_info_buff,
			MAX_PREFERRED_LENGTH, &read_sec, &total_users_sec, &resume
		);

		if ((status == NERR_Success || status == ERROR_MORE_DATA) && user_info_buff)
		{
			for (DWORD i = 0; i < read_sec && user_index < total_users; ++i)
			{
				LPCWSTR name = user_info_buff[i].usri0_name;

				BYTE sidBuffer[512];
				DWORD sidSize = sizeof(sidBuffer);
				WCHAR domain[256];
				DWORD domainSize = ARRAYSIZE(domain);
				SID_NAME_USE sidType;

				if (LookupAccountNameW(NULL, name, sidBuffer, &sidSize, domain, &domainSize, &sidType))
				{
					DWORD rid = *GetSidSubAuthority((PSID)sidBuffer, *GetSidSubAuthorityCount((PSID)sidBuffer) - 1);
					domain_users[user_index].user_rid = rid;
					(rid < 1000) ? ++system_users : ++regular_users;

					StringCchCopyW(domain_users[user_index].user_name, ARRAYSIZE(domain_users[user_index].user_name), name);
					++user_index;
				}
			}

			NetApiBufferFree(user_info_buff);
			user_info_buff = NULL;
		}

	} while (status == ERROR_MORE_DATA && user_index < total_users_sec);
	
	DWORD totalTime = GetTickCount() - start;
	log_line(LOG_TYPE_INFO,
		L"[%s] Pobrano informacje o uzytkownikach domenowych: Uzytkownikow regularnych: %d, uzytkownikow systemowych: %d, czas wykonania: %.2f s.",
		function_name,
		regular_users,
		system_users,
		totalTime / 1000.0);

	return S_OK;
}

DWORD WINAPI run_dc_monitor(LPVOID) {
	const DWORD maxRetries = 30;
	const DWORD retryDelayMs = 1000;

	WCHAR function_name[] = L"DomainMonitor";

	log_line(LOG_TYPE_INFO, 
		L"[%s] Start monitora gotowosci domeny...",
		function_name
	);

	DWORD start = GetTickCount();

	for (DWORD i = 0; i < maxRetries; ++i) {
		HRESULT res = get_domain_info();
		if (SUCCEEDED(res)) {
			LPWSTR sid_string = NULL;
			if (!ConvertSidToStringSidW(domain_info.domain_SID, &sid_string)) {
				log_line(LOG_TYPE_ERROR, L"[%s] Blad w konwersji SID", function_name);
			}
			LocalFree(sid_string);

			DWORD elapsed_ms = GetTickCount() - start;
			log_line(LOG_TYPE_INFO, 
				L"[%s] Informacje o domenie gotowe po %.2f s", 
				function_name, 
				elapsed_ms / 1000.0);

			log_line(LOG_TYPE_INFO,
				L"[%s] Status maszyny  --  Nazwa DNS domeny: %s, Nazwa NetBIOS: %s, PSID domeny: %s, Czy w domenie: %s, Czy kontroler domeny: %s",
				function_name,
				domain_info.DNS_domain_name,
				domain_info.NetBIOS_name,
				sid_string,
				(domain_info.is_machine_domain_joined) ? L"TAK" : L"NIE",
				(domain_info.is_domain_controller) ? L"TAK" : L"NIE"
			);

			domain_info.initialized = TRUE;

			return 0;
		}

		if (res != HRESULT_FROM_WIN32(RPC_S_SERVER_UNAVAILABLE)) {
			log_line(LOG_TYPE_ERROR,
				L"[%s] Nie mozna pobrac danych do ustalenia statusu DC: 0x%08X",
				function_name,
				res);
			return 1;
		}
		Sleep(retryDelayMs);
	}

	DWORD totalTime = GetTickCount() - start;
	log_line(LOG_TYPE_WARNING, 
		L"[%s] Timeout — nie pobrano danych o domenie po %lu ms.",
		function_name,
		totalTime);
	return 2;
}