239 lines
9.2 KiB
C++
239 lines
9.2 KiB
C++
#include <windows.h>
|
|
#include <shellapi.h>
|
|
#include <wtsapi32.h>
|
|
#include <cwchar>
|
|
#include <initializer_list>
|
|
#include <string>
|
|
|
|
namespace {
|
|
#ifdef STAYHERE_TEST_INSTANCE
|
|
constexpr wchar_t ClassName[] = L"StayHere.Test.Tray.1";
|
|
constexpr wchar_t MutexName[] = L"Local\\StayHere.Test.Tray.1";
|
|
#else
|
|
constexpr wchar_t ClassName[] = L"StayHere.Tray.1";
|
|
constexpr wchar_t MutexName[] = L"Local\\StayHere.Tray.1";
|
|
#endif
|
|
constexpr UINT TrayMessage = WM_APP + 1;
|
|
constexpr UINT ControlMessage = WM_APP + 2;
|
|
constexpr UINT Toggle = 1, Quit = 2, Pause = 3, Resume = 4, Status = 5;
|
|
constexpr DWORD IdleThreshold = 180000;
|
|
bool enabled = true;
|
|
bool sessionLocked = false;
|
|
bool powerOk = true;
|
|
bool inputOk = true;
|
|
UINT taskbarCreated;
|
|
HICON activeIcon, pausedIcon;
|
|
NOTIFYICONDATAW tray{};
|
|
|
|
constexpr bool shouldSignal(bool active, bool unlocked, DWORD now, DWORD lastInput) {
|
|
return active && unlocked && DWORD(now - lastInput) >= IdleThreshold;
|
|
}
|
|
static_assert(!shouldSignal(true, true, 179999, 0));
|
|
static_assert(shouldSignal(true, true, 180000, 0));
|
|
static_assert(!shouldSignal(false, true, 180000, 0));
|
|
static_assert(!shouldSignal(true, false, 180000, 0));
|
|
static_assert(shouldSignal(true, true, 179000, 0xfffffc18)); // tick counter wrap
|
|
|
|
bool interactiveDesktop() {
|
|
if (sessionLocked) return false;
|
|
HDESK desktop = OpenInputDesktop(0, FALSE, DESKTOP_READOBJECTS);
|
|
if (!desktop) return false;
|
|
wchar_t name[128]{};
|
|
DWORD needed = 0;
|
|
bool ok = GetUserObjectInformationW(desktop, UOI_NAME, name, sizeof(name), &needed)
|
|
&& _wcsicmp(name, L"Default") == 0;
|
|
CloseDesktop(desktop);
|
|
if (!ok) return false;
|
|
DWORD session;
|
|
LPWSTR buffer = nullptr;
|
|
DWORD bytes = 0;
|
|
if (!ProcessIdToSessionId(GetCurrentProcessId(), &session)) return false;
|
|
if (!WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, session, WTSConnectState, &buffer, &bytes)) return false;
|
|
ok = bytes >= sizeof(WTS_CONNECTSTATE_CLASS)
|
|
&& *reinterpret_cast<WTS_CONNECTSTATE_CLASS*>(buffer) == WTSActive;
|
|
WTSFreeMemory(buffer);
|
|
return ok;
|
|
}
|
|
|
|
HICON makeIcon(COLORREF color) {
|
|
HDC screen = GetDC(nullptr);
|
|
HDC dc = CreateCompatibleDC(screen);
|
|
HBITMAP bitmap = CreateCompatibleBitmap(screen, 32, 32);
|
|
HGDIOBJ previous = SelectObject(dc, bitmap);
|
|
RECT rect{0, 0, 32, 32};
|
|
FillRect(dc, &rect, static_cast<HBRUSH>(GetStockObject(BLACK_BRUSH)));
|
|
HBRUSH brush = CreateSolidBrush(color);
|
|
HGDIOBJ oldBrush = SelectObject(dc, brush);
|
|
HGDIOBJ oldPen = SelectObject(dc, GetStockObject(NULL_PEN));
|
|
Ellipse(dc, 3, 3, 29, 29);
|
|
SelectObject(dc, oldBrush);
|
|
SelectObject(dc, oldPen);
|
|
DeleteObject(brush);
|
|
SelectObject(dc, previous);
|
|
HBITMAP mask = CreateBitmap(32, 32, 1, 1, nullptr);
|
|
HGDIOBJ maskPrevious = SelectObject(dc, mask);
|
|
PatBlt(dc, 0, 0, 32, 32, WHITENESS);
|
|
oldBrush = SelectObject(dc, GetStockObject(BLACK_BRUSH));
|
|
oldPen = SelectObject(dc, GetStockObject(NULL_PEN));
|
|
Ellipse(dc, 3, 3, 29, 29);
|
|
SelectObject(dc, oldPen);
|
|
SelectObject(dc, oldBrush);
|
|
SelectObject(dc, maskPrevious);
|
|
ICONINFO info{TRUE, 0, 0, mask, bitmap};
|
|
HICON icon = CreateIconIndirect(&info);
|
|
DeleteObject(bitmap);
|
|
DeleteObject(mask);
|
|
DeleteDC(dc);
|
|
ReleaseDC(nullptr, screen);
|
|
return icon;
|
|
}
|
|
|
|
void updateTray() {
|
|
tray.hIcon = enabled ? activeIcon : pausedIcon;
|
|
const wchar_t* tip = !enabled ? L"StayHere — призупинено"
|
|
: !powerOk ? L"StayHere — помилка утримання від сну"
|
|
: !inputOk ? L"StayHere — не вдалося надіслати F15"
|
|
: L"StayHere — активно (F15 після 3 хв простою)";
|
|
lstrcpynW(tray.szTip, tip, ARRAYSIZE(tray.szTip));
|
|
Shell_NotifyIconW(NIM_MODIFY, &tray);
|
|
}
|
|
|
|
void setEnabled(bool value) {
|
|
enabled = value;
|
|
inputOk = true;
|
|
powerOk = SetThreadExecutionState(ES_CONTINUOUS |
|
|
(enabled && interactiveDesktop() ? ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED : 0)) != 0;
|
|
updateTray();
|
|
}
|
|
|
|
void tick() {
|
|
const bool unlocked = interactiveDesktop();
|
|
powerOk = SetThreadExecutionState(ES_CONTINUOUS |
|
|
(enabled && unlocked ? ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED : 0)) != 0;
|
|
LASTINPUTINFO idle{sizeof(idle), 0};
|
|
if (GetLastInputInfo(&idle) && shouldSignal(enabled, unlocked, GetTickCount(), idle.dwTime)) {
|
|
// Avoid inserting a key into a chord being held while the user is away.
|
|
for (int key : {VK_SHIFT, VK_CONTROL, VK_MENU, VK_LWIN, VK_RWIN,
|
|
VK_LBUTTON, VK_RBUTTON, VK_MBUTTON}) {
|
|
if (GetAsyncKeyState(key) & 0x8000) return;
|
|
}
|
|
INPUT input[2]{};
|
|
input[0].type = input[1].type = INPUT_KEYBOARD;
|
|
input[0].ki.wVk = input[1].ki.wVk = VK_F15;
|
|
input[1].ki.dwFlags = KEYEVENTF_KEYUP;
|
|
inputOk = SendInput(2, input, sizeof(INPUT)) == 2;
|
|
if (!inputOk) { // If only key-down succeeded, make a best effort to release it.
|
|
SendInput(1, &input[1], sizeof(INPUT));
|
|
}
|
|
}
|
|
updateTray();
|
|
}
|
|
|
|
void showMenu(HWND window) {
|
|
HMENU menu = CreatePopupMenu();
|
|
AppendMenuW(menu, MF_STRING | (enabled ? MF_CHECKED : 0), Toggle, L"Підтримувати активність");
|
|
AppendMenuW(menu, MF_SEPARATOR, 0, nullptr);
|
|
AppendMenuW(menu, MF_STRING, Quit, L"Вийти");
|
|
POINT pos;
|
|
GetCursorPos(&pos);
|
|
SetForegroundWindow(window);
|
|
UINT command = TrackPopupMenu(menu, TPM_RETURNCMD | TPM_RIGHTBUTTON, pos.x, pos.y, 0, window, nullptr);
|
|
DestroyMenu(menu);
|
|
if (command) SendMessageW(window, ControlMessage, command, 0);
|
|
PostMessageW(window, WM_NULL, 0, 0);
|
|
}
|
|
|
|
LRESULT CALLBACK windowProc(HWND window, UINT message, WPARAM wp, LPARAM lp) {
|
|
if (message == taskbarCreated) {
|
|
Shell_NotifyIconW(NIM_ADD, &tray);
|
|
return 0;
|
|
}
|
|
switch (message) {
|
|
case TrayMessage:
|
|
if (lp == WM_RBUTTONUP || lp == WM_CONTEXTMENU) showMenu(window);
|
|
else if (lp == WM_LBUTTONUP) setEnabled(!enabled);
|
|
return 0;
|
|
case ControlMessage:
|
|
if (wp == Status) return enabled ? 1 : 2;
|
|
if (wp == Toggle) setEnabled(!enabled);
|
|
else if (wp == Pause) setEnabled(false);
|
|
else if (wp == Resume) setEnabled(true);
|
|
else if (wp == Quit) DestroyWindow(window);
|
|
return 0;
|
|
case WM_TIMER: tick(); return 0;
|
|
case WM_WTSSESSION_CHANGE:
|
|
if (wp == WTS_SESSION_LOCK || wp == WTS_REMOTE_DISCONNECT || wp == WTS_CONSOLE_DISCONNECT)
|
|
sessionLocked = true;
|
|
else if (wp == WTS_SESSION_UNLOCK || wp == WTS_REMOTE_CONNECT || wp == WTS_CONSOLE_CONNECT)
|
|
sessionLocked = false;
|
|
setEnabled(enabled);
|
|
return 0;
|
|
case WM_DESTROY:
|
|
KillTimer(window, 1);
|
|
WTSUnRegisterSessionNotification(window);
|
|
SetThreadExecutionState(ES_CONTINUOUS);
|
|
Shell_NotifyIconW(NIM_DELETE, &tray);
|
|
PostQuitMessage(0);
|
|
return 0;
|
|
}
|
|
return DefWindowProcW(window, message, wp, lp);
|
|
}
|
|
}
|
|
|
|
int WINAPI wWinMain(HINSTANCE instance, HINSTANCE, PWSTR, int) {
|
|
int argc{};
|
|
auto argv=CommandLineToArgvW(GetCommandLineW(),&argc);
|
|
if(!argv)return 10;
|
|
const std::wstring argument=argc==2?argv[1]:L"";
|
|
LocalFree(argv);
|
|
if(argc>2)return 10;
|
|
const auto commandLine=argument.c_str();
|
|
UINT command = 0;
|
|
if (wcscmp(commandLine, L"--pause") == 0) command = Pause;
|
|
else if (wcscmp(commandLine, L"--resume") == 0) command = Resume;
|
|
else if (wcscmp(commandLine, L"--exit") == 0) command = Quit;
|
|
else if (wcscmp(commandLine, L"--status") == 0) command = Status;
|
|
else if (*commandLine) return 10;
|
|
HWND existing = FindWindowW(ClassName, nullptr);
|
|
if (command) {
|
|
if (!existing) return 3;
|
|
DWORD_PTR result = 0;
|
|
if (!SendMessageTimeoutW(existing, ControlMessage, command, 0, SMTO_ABORTIFHUNG, 2000, &result)) return 4;
|
|
return command == Status ? static_cast<int>(result) : 0;
|
|
}
|
|
HANDLE mutex = CreateMutexW(nullptr, FALSE, MutexName);
|
|
if (!mutex) return 5;
|
|
if (GetLastError() == ERROR_ALREADY_EXISTS) { CloseHandle(mutex); return 0; }
|
|
taskbarCreated = RegisterWindowMessageW(L"TaskbarCreated");
|
|
activeIcon = makeIcon(RGB(40, 185, 100));
|
|
pausedIcon = makeIcon(RGB(135, 140, 150));
|
|
WNDCLASSW wc{};
|
|
wc.lpfnWndProc = windowProc;
|
|
wc.hInstance = instance;
|
|
wc.lpszClassName = ClassName;
|
|
if (!RegisterClassW(&wc)) return 6;
|
|
HWND window = CreateWindowExW(0, ClassName, L"StayHere", WS_OVERLAPPED,
|
|
0, 0, 0, 0, nullptr, nullptr, instance, nullptr);
|
|
if (!window) return 7;
|
|
tray.cbSize = sizeof(tray);
|
|
tray.hWnd = window;
|
|
tray.uID = 1;
|
|
tray.uFlags = NIF_ICON | NIF_MESSAGE | NIF_TIP;
|
|
tray.uCallbackMessage = TrayMessage;
|
|
tray.hIcon = activeIcon;
|
|
lstrcpyW(tray.szTip, L"StayHere");
|
|
if (!Shell_NotifyIconW(NIM_ADD, &tray)) { DestroyWindow(window); return 8; }
|
|
WTSRegisterSessionNotification(window, NOTIFY_FOR_THIS_SESSION);
|
|
setEnabled(true);
|
|
if (!SetTimer(window, 1, 15000, nullptr)) { DestroyWindow(window); return 9; }
|
|
MSG msg{};
|
|
while (GetMessageW(&msg, nullptr, 0, 0) > 0) {
|
|
TranslateMessage(&msg);
|
|
DispatchMessageW(&msg);
|
|
}
|
|
DestroyIcon(activeIcon);
|
|
DestroyIcon(pausedIcon);
|
|
CloseHandle(mutex);
|
|
return 0;
|
|
}
|