150 lines
7.8 KiB
C++
150 lines
7.8 KiB
C++
#include "screenshot_notifications.h"
|
|
#include <atomic>
|
|
#include <shlobj.h>
|
|
#include <propkey.h>
|
|
#include <propvarutil.h>
|
|
#include <wrl/client.h>
|
|
namespace ks {
|
|
namespace {
|
|
constexpr wchar_t AppId[]=L"HrynCo.KSwitcher.Screenshot";
|
|
constexpr GUID ActivatorId{0xe3f78047,0xf0ad,0x49b6,{0x90,0x74,0x41,0x19,0x31,0x70,0x61,0x67}};
|
|
}
|
|
namespace {
|
|
using Microsoft::WRL::ComPtr;
|
|
void check(HRESULT result) {
|
|
if(FAILED(result)) throw std::runtime_error("Screenshot registration failed: " + std::to_string((unsigned long)result));
|
|
}
|
|
void registryString(const std::wstring &path,const wchar_t *name,const std::wstring &value)
|
|
{
|
|
HKEY key{};
|
|
if(RegCreateKeyExW(HKEY_CURRENT_USER,path.c_str(),0,nullptr,0,KEY_SET_VALUE,nullptr,&key,nullptr)!=ERROR_SUCCESS)
|
|
throw std::runtime_error("Cannot register screenshot notifications");
|
|
auto result=RegSetValueExW(key,name,0,REG_SZ,(const BYTE *)value.c_str(),DWORD((value.size()+1)*sizeof(wchar_t)));
|
|
RegCloseKey(key);
|
|
if(result!=ERROR_SUCCESS) throw std::runtime_error("Cannot register screenshot notification activation");
|
|
}
|
|
}
|
|
void registerScreenshotNotifications()
|
|
{
|
|
auto exe=(std::filesystem::path(executable()).parent_path()/L"KSwitcher.Screenshot.exe").wstring();
|
|
if(!std::filesystem::exists(exe)) throw std::runtime_error("Screenshot helper is missing");
|
|
registryString(L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32",nullptr,L"\""+exe+L"\" --toast-activated");
|
|
registryString(L"Software\\Classes\\CLSID\\{e3f78047-f0ad-49b6-9074-411931706167}\\LocalServer32",L"ServerExecutable",exe);
|
|
auto identity=std::wstring(L"Software\\Classes\\AppUserModelId\\")+AppId;
|
|
registryString(identity,L"DisplayName",L"KSwitcher Screenshot");
|
|
registryString(identity,L"IconUri",exe);
|
|
registryString(identity,L"CustomActivator",L"{e3f78047-f0ad-49b6-9074-411931706167}");
|
|
// A shortcut supplies desktop notification identity and its COM activation callback.
|
|
ComPtr<IShellLinkW> link; check(CoCreateInstance(CLSID_ShellLink,nullptr,CLSCTX_INPROC_SERVER,IID_PPV_ARGS(&link)));
|
|
check(link->SetPath(exe.c_str())); check(link->SetIconLocation(exe.c_str(),0));
|
|
ComPtr<IPropertyStore> properties; check(link.As(&properties));
|
|
PROPVARIANT id{}; check(InitPropVariantFromString(AppId,&id));
|
|
auto result=properties->SetValue(PKEY_AppUserModel_ID,id); PropVariantClear(&id); check(result);
|
|
GUID activator=ActivatorId;
|
|
PROPVARIANT clsid{}; clsid.vt=VT_CLSID; clsid.puuid=&activator;
|
|
check(properties->SetValue(PKEY_AppUserModel_ToastActivatorCLSID,clsid));
|
|
check(properties->Commit());
|
|
PWSTR programs{}; check(SHGetKnownFolderPath(FOLDERID_Programs,0,nullptr,&programs));
|
|
auto shortcut=std::filesystem::path(programs)/L"KSwitcher Screenshot.lnk"; CoTaskMemFree(programs);
|
|
ComPtr<IPersistFile> file; check(link.As(&file)); check(file->Save(shortcut.c_str(),TRUE));
|
|
// Notify the shell that notification identity/activation registration changed.
|
|
SHChangeNotify(SHCNE_ASSOCCHANGED,SHCNF_IDLIST,nullptr,nullptr);
|
|
}
|
|
bool validScreenshotToken(const std::wstring &id)
|
|
{
|
|
if(id.size()!=36) return false;
|
|
for(size_t i=0;i<id.size();++i) {
|
|
if(i==8 || i==13 || i==18 || i==23) { if(id[i]!=L'-') return false; }
|
|
else if(!((id[i]>=L'0' && id[i]<=L'9') || (id[i]>=L'a' && id[i]<=L'f') || (id[i]>=L'A' && id[i]<=L'F'))) return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
namespace {
|
|
// Windows SDK INotificationActivationCallback ABI (notificationactivationcallback.h).
|
|
constexpr GUID CallbackId{0x53e31837,0x6600,0x4a81,{0x93,0x95,0x75,0xcf,0xfe,0x74,0x6f,0x94}};
|
|
struct NotificationInput { LPCWSTR key,value; };
|
|
struct NotificationCallback : IUnknown {
|
|
virtual HRESULT STDMETHODCALLTYPE Activate(LPCWSTR,LPCWSTR,const NotificationInput *,ULONG)=0;
|
|
};
|
|
class Activator final : public NotificationCallback {
|
|
std::atomic<ULONG> references{1};
|
|
DWORD thread=GetCurrentThreadId();
|
|
std::wstring target;
|
|
bool quit;
|
|
public:
|
|
Activator(std::wstring path,bool shouldQuit):target(std::move(path)),quit(shouldQuit) {}
|
|
HRESULT STDMETHODCALLTYPE QueryInterface(REFIID id,void **out) override {
|
|
if(!out) return E_POINTER;
|
|
*out=nullptr;
|
|
if(id!=IID_IUnknown && id!=CallbackId) return E_NOINTERFACE;
|
|
*out=static_cast<NotificationCallback *>(this); AddRef(); return S_OK;
|
|
}
|
|
ULONG STDMETHODCALLTYPE AddRef() override { return ++references; }
|
|
ULONG STDMETHODCALLTYPE Release() override { auto count=--references; if(!count) delete this; return count; }
|
|
HRESULT STDMETHODCALLTYPE Activate(LPCWSTR app,LPCWSTR args,const NotificationInput *,ULONG) override {
|
|
if(!app || std::wstring(app)!=AppId || !args) return E_INVALIDARG;
|
|
const bool update=std::wstring(args)==L"check-update";
|
|
if(!update && !validScreenshotToken(args)) return E_INVALIDARG;
|
|
auto path=update ? (std::filesystem::path(target).parent_path()/L"KSwitcher.exe").wstring() : target;
|
|
auto command=update ? L"\""+path+L"\" --check-update" :
|
|
L"\""+path+L"\" --open \"kswitcher-screenshot:"+args+L"\"";
|
|
STARTUPINFOW startup{sizeof(startup)}; PROCESS_INFORMATION process{};
|
|
auto ok=CreateProcessW(path.c_str(),command.data(),nullptr,nullptr,FALSE,CREATE_SUSPENDED,nullptr,nullptr,&startup,&process);
|
|
auto result=ok?S_OK:HRESULT_FROM_WIN32(GetLastError());
|
|
if(ok) {
|
|
AllowSetForegroundWindow(process.dwProcessId); ResumeThread(process.hThread);
|
|
CloseHandle(process.hThread); CloseHandle(process.hProcess);
|
|
}
|
|
if(quit) PostThreadMessageW(thread,WM_QUIT,0,0);
|
|
return result;
|
|
}
|
|
};
|
|
class ActivatorFactory final : public IClassFactory {
|
|
std::atomic<ULONG> references{1};
|
|
std::wstring target;
|
|
bool quit;
|
|
public:
|
|
ActivatorFactory(std::wstring path,bool shouldQuit):target(std::move(path)),quit(shouldQuit) {}
|
|
HRESULT STDMETHODCALLTYPE QueryInterface(REFIID id,void **out) override {
|
|
if(!out) return E_POINTER;
|
|
*out=nullptr;
|
|
if(id!=IID_IUnknown && id!=IID_IClassFactory) return E_NOINTERFACE;
|
|
*out=static_cast<IClassFactory *>(this); AddRef(); return S_OK;
|
|
}
|
|
ULONG STDMETHODCALLTYPE AddRef() override { return ++references; }
|
|
ULONG STDMETHODCALLTYPE Release() override { auto count=--references; if(!count) delete this; return count; }
|
|
HRESULT STDMETHODCALLTYPE CreateInstance(IUnknown *outer,REFIID id,void **out) override {
|
|
if(outer) return CLASS_E_NOAGGREGATION;
|
|
auto object=new Activator(target,quit); auto result=object->QueryInterface(id,out); object->Release(); return result;
|
|
}
|
|
HRESULT STDMETHODCALLTYPE LockServer(BOOL) override { return S_OK; }
|
|
};
|
|
}
|
|
int runScreenshotToastActivator()
|
|
{
|
|
if(FAILED(CoInitializeEx(nullptr,COINIT_APARTMENTTHREADED))) return 1;
|
|
MSG message{}; PeekMessageW(&message,nullptr,0,0,PM_NOREMOVE);
|
|
auto factory=new ActivatorFactory(executable(),true); DWORD registration{};
|
|
auto result=CoRegisterClassObject(ActivatorId,factory,CLSCTX_LOCAL_SERVER,REGCLS_MULTIPLEUSE,®istration);
|
|
factory->Release();
|
|
if(SUCCEEDED(result)) {
|
|
auto timer=SetTimer(nullptr,0,60000,[](HWND,UINT,UINT_PTR,DWORD){PostQuitMessage(0);});
|
|
while(GetMessageW(&message,nullptr,0,0)>0) { TranslateMessage(&message); DispatchMessageW(&message); }
|
|
if(timer) KillTimer(nullptr,timer);
|
|
CoRevokeClassObject(registration);
|
|
}
|
|
CoUninitialize(); return FAILED(result)?1:0;
|
|
}
|
|
DWORD startScreenshotActivationListener()
|
|
{
|
|
auto target=std::filesystem::path(executable()).parent_path()/L"KSwitcher.Screenshot.exe";
|
|
auto factory=new ActivatorFactory(target.wstring(),false);
|
|
DWORD registration{};
|
|
auto result=CoRegisterClassObject(ActivatorId,factory,CLSCTX_LOCAL_SERVER,REGCLS_MULTIPLEUSE,®istration);
|
|
factory->Release();
|
|
return SUCCEEDED(result)?registration:0;
|
|
}
|
|
|
|
}
|