1  // dllmain.cpp : Defines the entry point for the DLL application.
  2  #include "pch.h"
  3  
  4  #include <common/telemetry/EtwTrace/EtwTrace.h>
  5  #include <common/utils/process_path.h>
  6  #include <common/utils/resources.h>
  7  #include <common/utils/elevation.h>
  8  
  9  #include "FileLocksmithLib/IPC.h"
 10  #include "FileLocksmithLib/Settings.h"
 11  #include "FileLocksmithLib/Trace.h"
 12  
 13  #include <Shlwapi.h>
 14  #include <shobjidl_core.h>
 15  #include <string>
 16  #include <wrl/module.h>
 17  
 18  #include "Generated Files/resource.h"
 19  
 20  
 21  using namespace Microsoft::WRL;
 22  
 23  HINSTANCE g_hInst = 0;
 24  Shared::Trace::ETWTrace trace(L"FileLocksmithContextMenu");
 25  
 26  BOOL APIENTRY DllMain( HMODULE hModule,
 27                         DWORD  ul_reason_for_call,
 28                         LPVOID lpReserved
 29                       )
 30  {
 31      switch (ul_reason_for_call)
 32      {
 33      case DLL_PROCESS_ATTACH:
 34          g_hInst = hModule;
 35          Trace::RegisterProvider();
 36          break;
 37      case DLL_PROCESS_DETACH:
 38          Trace::UnregisterProvider();
 39          break;
 40      }
 41      return TRUE;
 42  }
 43  
 44  class __declspec(uuid("AAF1E27D-4976-49C2-8895-AAFA743C0A7E")) FileLocksmithContextMenuCommand final : public RuntimeClass<RuntimeClassFlags<ClassicCom>, IExplorerCommand, IObjectWithSite>
 45  {
 46  public:
 47      virtual const wchar_t* Title() { return L"File Locksmith"; }
 48      virtual const EXPCMDFLAGS Flags() { return ECF_DEFAULT; }
 49      virtual const EXPCMDSTATE State(_In_opt_ IShellItemArray* selection) { return ECS_ENABLED; }
 50  
 51      // IExplorerCommand
 52      IFACEMETHODIMP GetTitle(_In_opt_ IShellItemArray* items, _Outptr_result_nullonfailure_ PWSTR* name)
 53      {
 54          return SHStrDup(context_menu_caption.c_str(), name);
 55      }
 56  
 57      IFACEMETHODIMP GetIcon(_In_opt_ IShellItemArray*, _Outptr_result_nullonfailure_ PWSTR* icon)
 58      {
 59          std::wstring iconResourcePath = get_module_folderpath(g_hInst);
 60          iconResourcePath += L"\\Assets\\FileLocksmith\\";
 61          iconResourcePath += L"FileLocksmith.ico";
 62          return SHStrDup(iconResourcePath.c_str(), icon);
 63      }
 64  
 65      IFACEMETHODIMP GetToolTip(_In_opt_ IShellItemArray*, _Outptr_result_nullonfailure_ PWSTR* infoTip)
 66      {
 67          *infoTip = nullptr;
 68          return E_NOTIMPL;
 69      }
 70  
 71      IFACEMETHODIMP GetCanonicalName(_Out_ GUID* guidCommandName)
 72      {
 73          *guidCommandName = __uuidof(this);
 74          return S_OK;
 75      }
 76  
 77      IFACEMETHODIMP GetState(_In_opt_ IShellItemArray* selection, _In_ BOOL okToBeSlow, _Out_ EXPCMDSTATE* cmdState)
 78      {
 79          *cmdState = ECS_ENABLED;
 80  
 81          if (!FileLocksmithSettingsInstance().GetEnabled())
 82          {
 83              *cmdState = ECS_HIDDEN;
 84          }
 85  
 86          if (FileLocksmithSettingsInstance().GetShowInExtendedContextMenu())
 87          {
 88              *cmdState = ECS_HIDDEN;
 89          }
 90  
 91          return S_OK;
 92      }
 93  
 94      IFACEMETHODIMP Invoke(_In_opt_ IShellItemArray* selection, _In_opt_ IBindCtx*) noexcept
 95      {
 96          trace.UpdateState(true);
 97  
 98          Trace::Invoked();
 99          ipc::Writer writer;
100  
101          if (selection == nullptr)
102          {
103              return S_OK;
104          }
105  
106          if (HRESULT result = writer.start(); FAILED(result))
107          {
108              Trace::InvokedRet(result);
109  
110              trace.Flush();
111              trace.UpdateState(false);
112              return result;
113          }
114  
115          std::wstring path = get_module_folderpath(g_hInst);
116          path = path + L"\\PowerToys.FileLocksmithUI.exe";
117  
118          HRESULT result;
119  
120          if (!RunNonElevatedEx(path.c_str(), L"", get_module_folderpath(g_hInst)))
121          {
122              result = E_FAIL;
123              Trace::InvokedRet(result);
124  
125              trace.Flush();
126              trace.UpdateState(false);
127  
128              return result;
129          }
130  
131          DWORD num_items;
132          selection->GetCount(&num_items);
133  
134          for (DWORD i = 0; i < num_items; i++)
135          {
136              IShellItem* item;
137              result = selection->GetItemAt(i, &item);
138              if (SUCCEEDED(result))
139              {
140                  LPWSTR file_path;
141                  result = item->GetDisplayName(SIGDN_FILESYSPATH, &file_path);
142                  if (SUCCEEDED(result))
143                  {
144                      // TODO Aggregate items and send to UI
145                      writer.add_path(file_path);
146                      CoTaskMemFree(file_path);
147                  }
148  
149                  item->Release();
150              }
151          }
152  
153          Trace::InvokedRet(S_OK);
154  
155          trace.Flush();
156          trace.UpdateState(false);
157  
158          return S_OK;
159      }
160  
161      IFACEMETHODIMP GetFlags(_Out_ EXPCMDFLAGS* flags)
162      {
163          *flags = Flags();
164          return S_OK;
165      }
166      IFACEMETHODIMP EnumSubCommands(_COM_Outptr_ IEnumExplorerCommand** enumCommands)
167      {
168          *enumCommands = nullptr;
169          return E_NOTIMPL;
170      }
171  
172      // IObjectWithSite
173      IFACEMETHODIMP SetSite(_In_ IUnknown* site) noexcept
174      {
175          m_site = site;
176          return S_OK;
177      }
178      IFACEMETHODIMP GetSite(_In_ REFIID riid, _COM_Outptr_ void** site) noexcept { return m_site.CopyTo(riid, site); }
179  
180  protected:
181      ComPtr<IUnknown> m_site;
182  
183  private:
184      std::wstring context_menu_caption = GET_RESOURCE_STRING_FALLBACK(IDS_FILE_LOCKSMITH_CONTEXT_MENU_ENTRY, L"Unlock with File Locksmith");
185  };
186  
187  CoCreatableClass(FileLocksmithContextMenuCommand)
188  CoCreatableClassWrlCreatorMapInclude(FileLocksmithContextMenuCommand)
189  
190  
191  STDAPI DllGetActivationFactory(_In_ HSTRING activatableClassId, _COM_Outptr_ IActivationFactory** factory)
192  {
193      return Module<ModuleType::InProc>::GetModule().GetActivationFactory(activatableClassId, factory);
194  }
195  
196  STDAPI DllCanUnloadNow()
197  {
198      return Module<InProc>::GetModule().GetObjectCount() == 0 ? S_OK : S_FALSE;
199  }
200  
201  STDAPI DllGetClassObject(_In_ REFCLSID rclsid, _In_ REFIID riid, _COM_Outptr_ void** instance)
202  {
203      return Module<InProc>::GetModule().GetClassObject(rclsid, riid, instance);
204  }