#include"ZamknijSystem.h"

#include <exception>

bool CzyWinNT()
{
	OSVERSIONINFO osInfo;
	osInfo.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);

	if(GetVersionEx(&osInfo) == 0)
	{
		throw std::exception("Nie udało się pobrać informacji o systemie");		
	}

	return osInfo.dwPlatformId == VER_PLATFORM_WIN32_NT;
}

bool ZamknijSystem1(int Action)
{
	//Jezeli tylko logoff (do niego nie trzeba uprawnien), to robi i wychodzi
	if((Action & ~(EWX_FORCE | EWX_FORCEIFHUNG))==EWX_LOGOFF)
	{
		return ExitWindowsEx(Action,0)!=0;
	}

	//Jezeli nie NT, to robi od razu (funkcje pytajace o uprawnienia sa nieobecne w Win95/98/ME)
	if (!CzyWinNT())
	{
		return ExitWindowsEx(Action,0)!=0;
	}

	//Jezeli jednak NT, to pobiera odpowiednie uprawnienia

	//Otwieranie tokenu dostępu
	HANDLE hToken;
	if (!OpenProcessToken(GetCurrentProcess(),TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,&hToken))
	{
		MessageBox(0,L"Nie udało się otworzyć tokenu procesu",L"",MB_OK);
		return false;
	}

	//Pobranie LUID dla uprawnienia zamykania systemu (SE_SHUTDOWN_NAME) dla pierwszego (i jednynego) uprawnienia umieszczonego w tkp
	LUID Luid; //zmienna 64-bitowa, w której przechowywany bedzie kod LUID uprawnienia
	LookupPrivilegeValue(NULL,SE_SHUTDOWN_NAME,&Luid);

	//Ustalenie wartosci dla uprawnienia
	TOKEN_PRIVILEGES tkp; // struktura tokenu
	tkp.PrivilegeCount=1; //ile uprawnien do ustalenia
	tkp.Privileges[0].Luid=Luid;
	tkp.Privileges[0].Attributes=SE_PRIVILEGE_ENABLED; //uaktywanienie uprawnienia (nadanie wartosści elementowi Attributes)

	//Modyfikacja uprawnienia do zamkniecia systemu dla biezacego procesu (shutdown privilege)
	//nie pobieramy informacji o stanie uprawnien sprzed modyfikacji
	//ustalamy wartosc uprawnienia dla hToken na ustalona w &tkp
	AdjustTokenPrivileges(hToken,false,&tkp,0,(TOKEN_PRIVILEGES*) NULL,NULL);

	//Jezeli uruchomienie AdjustTokenPrivileges nie powidiodlo sie (np. system nie jest NT), wyswietlany jest komunikat
	if (GetLastError()!=ERROR_SUCCESS) MessageBox(0,L"Przyznanie uprawnienia nie powiodło się.",L"",MB_OK);

	//Uruchomienie funkcji ExitWindowsEx (zadziala w NT tylko jezeli posiadamy (dokladniej jezeli proces posiada) uprawnienia do zamykania systemu)
	bool Wynik=ExitWindowsEx(Action,0)!=0;  //w Wynik bedzie przechowywany wynik uruchomienia ExitWindowsEx (ewentualny kod bledu)
	if (!Wynik) MessageBox(0,L"Zamknięcie systemu nie powiodło się.",L"",MB_OK);

	//Odbieranie procesowi uprawnienia do zamykania systemu
	tkp.Privileges[0].Attributes=0; //atrybut odebrania uprawnienia
	AdjustTokenPrivileges(hToken,false,&tkp,NULL,NULL,NULL); //modyfikacja uprawnienia bez pobrania informacji, korzystamy z poprzednio pobranego LUID

	//Informacja jeżeli odbieranie uprawnienia nie powiodło się
	if (GetLastError()!=ERROR_SUCCESS) MessageBox(0,L"Odebranie uprawnienia nie powiodło się.",L"",MB_OK);

	CloseHandle(hToken);

	return Wynik;
}


//---------------------------------------------------------------------------

bool ZamknijSystem2(bool inicjujLubPrzerwij,CString nazwaKomputera,CString komunikat,int opoznienie,bool force,bool reboot)
{
	//Zob. komentarze w ZamknijSystem1

	if (!CzyWinNT())
	{
		MessageBox(0,L"Funkcja ZamknijSystem2 obsługuje jedynie systemy oparte na technologii NT",L"",MB_OK);
		return false;
	}

	//Jezeli nazwa komputera nie jest wpisana, to system lokalny
	wchar_t * pNazwaKomputera= nazwaKomputera.MakeUpper().GetBuffer();	
	if (nazwaKomputera.IsEmpty() || nazwaKomputera.GetBuffer() == NULL) pNazwaKomputera=NULL;


	//Jezeli wpisana nazwa identyfikuje lokalny komputer, to zmiana na NULL
	//aby uniknac pobierania sieciowych uprawnien
	wchar_t nazwaLokalnegoKomputera[MAX_COMPUTERNAME_LENGTH+1];
	unsigned long dlg_nazwaLokalnegoKomputera=MAX_COMPUTERNAME_LENGTH+1;
	GetComputerName(nazwaLokalnegoKomputera,&dlg_nazwaLokalnegoKomputera);
	if (pNazwaKomputera!=NULL &&  wcscmp(pNazwaKomputera,_wcsupr(nazwaLokalnegoKomputera))==0) pNazwaKomputera=NULL;


	HANDLE hToken;
	if (!OpenProcessToken(GetCurrentProcess(),TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,&hToken))
	{
		MessageBox(0,L"Nie udało się otworzyć tokenu procesu",L"",MB_OK);
		return false;
	}


	//Teraz podawana jest tu nazwa komputera pobierana przez glowe funkcji ZamknijSystem2
	//Jezeli zamykamy lokalny komputer pobieramy LUID uprawnienia SE_SHUTDOWN_NAME na lokalnym komputerze, jezeli zdalny to SE_REMOTE_SHUTDOWN_NAME na zdalnym komputerze
	LUID Luid;
	bool wynik;
	if (pNazwaKomputera==NULL)
		wynik=LookupPrivilegeValue(NULL,SE_SHUTDOWN_NAME,&Luid)!=0;
	else
		wynik=LookupPrivilegeValue(NULL,SE_REMOTE_SHUTDOWN_NAME,&Luid)!=0;

	if(!wynik)
	{
		MessageBox(0,L"Odnalezienie kodu LUID uprawnienia nie powiodlo sie.",L"",MB_OK);
		return false;
	}

	TOKEN_PRIVILEGES tkp;
	tkp.PrivilegeCount=1;
	tkp.Privileges[0].Luid=Luid;
	tkp.Privileges[0].Attributes=SE_PRIVILEGE_ENABLED;

	if (!AdjustTokenPrivileges(hToken,false,&tkp,0,NULL,NULL))
	{
		MessageBox(0,L"Przyznanie uprawnienia nie powiodło się.",L"",MB_OK);
		return false;
	}

	if (inicjujLubPrzerwij)
	{
		if (!InitiateSystemShutdown(pNazwaKomputera,komunikat.GetBuffer(),opoznienie,force,reboot))
		{
			MessageBox(0,L"Zamknięcie systemu nie jest możliwe.",L"",MB_OK);
			return false;
		}
	}
	else
	{
		if (!AbortSystemShutdown(pNazwaKomputera))
		{
			MessageBox(0,L"Odwołanie zamknięcia zdalnego systemu nie jest możliwe.",L"",MB_OK);
			return false;
		}
	}

	tkp.Privileges[0].Attributes=0;
	AdjustTokenPrivileges(hToken,false,&tkp,0,NULL,NULL);
	if (GetLastError()!=ERROR_SUCCESS)
	{
		MessageBox(0,L"Odebranie uprawnienia nie powiodło się.",L"",MB_OK);
		return false;
	}

	return true;
}

//---------------------------------------------------------------------------

bool UspijSystem(bool przelaczWStanOczekiwania,bool wymuszenie)
{
	//Zob. komentarze w ZamknijSystem1
	
	if (!CzyWinNT())
	{
		return SetSystemPowerState(przelaczWStanOczekiwania,wymuszenie)!=0;
	}

	HANDLE hToken;
	if (!OpenProcessToken(GetCurrentProcess(),TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,&hToken))
	{
		MessageBox(0,L"Nie udało się otworzyć tokenu procesu",L"",MB_OK);
		return false;
	}
	
	LUID Luid; 
	LookupPrivilegeValue(NULL,SE_SHUTDOWN_NAME,&Luid);

	TOKEN_PRIVILEGES tkp;
	tkp.PrivilegeCount=1;
	tkp.Privileges[0].Luid=Luid;
	tkp.Privileges[0].Attributes=SE_PRIVILEGE_ENABLED;

	AdjustTokenPrivileges(hToken,false,&tkp,0,(TOKEN_PRIVILEGES*) NULL,NULL);

	if (GetLastError()!=ERROR_SUCCESS) MessageBox(0,L"Przyznanie uprawnienia nie powiodło się.",L"",MB_OK);

	//Uruchomienie funkcji SetSystemPowerState
	bool Wynik=SetSystemPowerState(przelaczWStanOczekiwania,wymuszenie)!=0;
	if (!Wynik) MessageBox(0,L"Wstrzymanie systemu lub hibernacja nie powiodła się.",L"",MB_OK);

	tkp.Privileges[0].Attributes=0;
	AdjustTokenPrivileges(hToken,false,&tkp,NULL,NULL,NULL);

	if (GetLastError()!=ERROR_SUCCESS) MessageBox(0,L"Odebranie uprawnienia nie powiodło się.",L"",MB_OK);

	CloseHandle(hToken);

	return Wynik;
}

