/*

* Copyright (c) 2009 Secure Endpoints Inc.

*

* Permission is hereby granted, free of charge, to any person

* obtaining a copy of this software and associated documentation

* files (the "Software"), to deal in the Software without

* restriction, including without limitation the rights to use, copy,

* modify, merge, publish, distribute, sublicense, and/or sell copies

* of the Software, and to permit persons to whom the Software is

* furnished to do so, subject to the following conditions:

*

* The above copyright notice and this permission notice shall be

* included in all copies or substantial portions of the Software.

*

* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,

* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF

* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND

* NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT

* HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,

* WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER

* DEALINGS IN THE SOFTWARE.

*/

#define NOSTRSAFE

#include "credprov.h"

#include <tchar.h>

#include <shellapi.h>

#include <htmlhelp.h>

#include <strsafe.h>

#include <assert.h>

static ATOM message_window_class = 0;

static HWND notifier_window = NULL;

static volatile BOOL notification_icon_added = FALSE;

#define TOKEN_ICON_ID 1

#define TOKEN_MESSAGE_ID WM_USER

static khm_int32

get_default_notifier_action(void)

{

khm_int32 cmd = KHUI_ACTION_OPEN_APP;

khc_read_int32(NULL, L"CredWindow\\NotificationAction", &cmd);

return cmd;

}

static void

prepare_context_menu(HMENU hmenu)

{

khm_int32 cmd;

wchar_t caption[128];

cmd = get_default_notifier_action();

if (cmd == KHUI_ACTION_NEW_CRED)

LoadString(hResModule, IDS_ACT_NEW, caption, ARRAYLENGTH(caption));

else

LoadString(hResModule, IDS_ACT_OPEN, caption, ARRAYLENGTH(caption));

ModifyMenu(hmenu, ID_DEFAULT, MF_STRING|MF_BYCOMMAND, ID_DEFAULT, caption);

SetMenuDefaultItem(hmenu, ID_DEFAULT, FALSE);

}

static void

handle_context_menu(void)

{

POINT pt;

HMENU hMenu;

HMENU hMenuBar;

GetCursorPos(&pt);

hMenuBar = LoadMenu(hResModule, MAKEINTRESOURCE(IDR_CTXMENU));

hMenu = GetSubMenu(hMenuBar, 0);

if (hMenu) {

prepare_context_menu(hMenu);

TrackPopupMenu(hMenu, TPM_NONOTIFY, pt.x, pt.y, 0, notifier_window, NULL);

}

{

NOTIFYICONDATA idata;

ZeroMemory(&idata, sizeof(idata));

Shell_NotifyIcon(NIM_SETFOCUS, &idata);

}

}

static LRESULT CALLBACK

notifier_wnd_proc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam)

{

if (uMsg == TOKEN_MESSAGE_ID) {

switch (lParam) {

case NIN_SELECT:

case NIN_KEYSELECT:

{

NOTIFYICONDATA idata;

khm_int32 cmd = KHUI_ACTION_OPEN_APP;

khc_read_int32(NULL, L"CredWindow\\NotificationAction", &cmd);

khui_action_trigger(cmd, NULL);

ZeroMemory(&idata, sizeof(idata));

Shell_NotifyIcon(NIM_SETFOCUS, &idata);

}

return 0;

case WM_CONTEXTMENU:

handle_context_menu();

return TRUE;

default:

return 0;

}

}

else if (uMsg == WM_COMMAND) {

switch (LOWORD(wParam)) {

case ID_DEFAULT:

{

khm_int32 cmd;

cmd = get_default_notifier_action();

khui_action_trigger(cmd, NULL);

}

return TRUE;

case ID_SHOWHELP:

{

wchar_t helploc[MAX_PATH + MAX_PATH];

get_help_file(helploc, sizeof(helploc));

StringCbCat(helploc, sizeof(helploc), L"::index.html");

HtmlHelp(notifier_window, helploc, HH_DISPLAY_TOPIC, 0);

}

return TRUE;

}

}

return DefWindowProc(hwnd, uMsg, wParam, lParam);

}

#pragma warning(push)

#pragma warning(disable: 4204)

static void

initialize_if_necessary(void)

{

if (message_window_class == 0) {

WNDCLASSEX c = {

sizeof(WNDCLASSEX), /* cbSize */

0, /* style */

notifier_wnd_proc, /* lpfnWndProc */

0, /* cbClsExtra */

0, /* cbWndExtra */

hInstance, /* hinstance */

NULL, /* hIcon */

NULL, /* hCursor */

NULL, /* hbrBackground */

NULL, /* lpszMenuName */

L"KCACredStateIconNotifier", /* lpszClassName */

NULL, /* hIconSm */

};

message_window_class = RegisterClassEx(&c);

}

if (notifier_window == NULL && message_window_class != 0) {

notifier_window = CreateWindow(MAKEINTATOM(message_window_class),

L"KCACredStateIconNotifierWindow",

0, 0, 0, 0, 0,

HWND_MESSAGE,

NULL,

hInstance,

NULL);

}

assert(notifier_window != NULL);

if (!notification_icon_added && notifier_window != NULL) {

NOTIFYICONDATA idata;

ZeroMemory(&idata, sizeof(idata));

idata.cbSize = sizeof(idata);

idata.hWnd = notifier_window;

idata.uID = TOKEN_ICON_ID;

idata.uFlags = NIF_ICON | NIF_MESSAGE;

idata.uCallbackMessage = TOKEN_MESSAGE_ID;

idata.hIcon = (HICON) LoadImage(hResModule, MAKEINTRESOURCE(IDI_CRED_NONE),

IMAGE_ICON, 0, 0,

LR_DEFAULTSIZE | LR_DEFAULTCOLOR | LR_SHARED);

notification_icon_added = Shell_NotifyIcon(NIM_ADD, &idata);

idata.cbSize = sizeof(idata);

idata.uVersion = NOTIFYICON_VERSION;

Shell_NotifyIcon(NIM_SETVERSION, &idata);

assert(notification_icon_added);

}

}

#pragma warning(pop)

void

kca_remove_icon(void)

{

NOTIFYICONDATA idata;

ZeroMemory(&idata, sizeof(idata));

idata.cbSize = sizeof(idata);

idata.hWnd = notifier_window;

idata.uID = TOKEN_ICON_ID;

Shell_NotifyIcon(NIM_DELETE, &idata);

notification_icon_added = FALSE;

}

static void

set_tooltip_and_icon(UINT tooltip_text, const wchar_t * postfix, UINT icon_id)

{

NOTIFYICONDATA idata;

wchar_t buf[ARRAYLENGTH(idata.szTip)];

ZeroMemory(&idata, sizeof(idata));

idata.cbSize = sizeof(idata);

idata.hWnd = notifier_window;

idata.uID = TOKEN_ICON_ID;

idata.uFlags = NIF_ICON | NIF_TIP;

idata.hIcon = (HICON) LoadImage(hResModule, MAKEINTRESOURCE(icon_id),

IMAGE_ICON, 0, 0,

LR_DEFAULTCOLOR | LR_DEFAULTSIZE | LR_SHARED);

if (tooltip_text != 0) {

LoadString(hResModule, tooltip_text, buf, ARRAYLENGTH(buf));

}

StringCbPrintf(idata.szTip, sizeof(idata.szTip),

L"%s%s",

(tooltip_text != 0)? buf : L"",

(postfix != NULL)? postfix : L"");

Shell_NotifyIcon(NIM_MODIFY, &idata);

}

struct state_data {

khm_handle credset;

};

#define COLLECT_STR_LEN 256

static khm_int32 KHMAPI

collect_kca_cert_names(khm_handle cred, void * rock)

{

wchar_t *str = (wchar_t *) rock;

wchar_t cert_name[KCDB_MAXCCH_NAME] = L"";

FILETIME ft_now;

FILETIME ft_expire;

khm_size cb;

cb = sizeof(ft_expire);

if (KHM_FAILED(kcdb_cred_get_attr(cred, KCDB_ATTR_EXPIRE, NULL, &ft_expire, &cb)))

return KHM_ERROR_SUCCESS;

GetSystemTimeAsFileTime(&ft_now);

if (CompareFileTime(&ft_now, &ft_expire) >= 0)

return KHM_ERROR_SUCCESS;

cb = sizeof(cert_name);

if (KHM_SUCCEEDED(kcdb_cred_get_attr(cred, attr_id_auth_realm, NULL, cert_name, &cb)) &&

cert_name[0] != L'\0') {

StringCchCat(str, COLLECT_STR_LEN, cert_name);

StringCchCat(str, COLLECT_STR_LEN, L"\n");

} else if ((cb = sizeof(cert_name)) &&

KHM_SUCCEEDED(kcdb_cred_get_attr(cred, attr_id_subj_email, NULL, cert_name, &cb))

&& cert_name[0] != L'\0') {

StringCchCat(str, COLLECT_STR_LEN, cert_name);

StringCchCat(str, COLLECT_STR_LEN, L"\n");

}

return KHM_ERROR_SUCCESS;

}

static khm_int32 KHMAPI

set_state_from_ui_thread(HWND hwnd_main, void * stuff)

{

struct state_data * d = (struct state_data *) stuff;

wchar_t certs[COLLECT_STR_LEN] = L"";

initialize_if_necessary();

if (d->credset)

kcdb_credset_apply(d->credset, collect_kca_cert_names, certs);

if (certs[0] == L'\0') {

set_tooltip_and_icon(IDS_CRED_TT_NONE, NULL, IDI_CRED_NONE);

return KHM_ERROR_SUCCESS;

}

set_tooltip_and_icon(IDS_CRED_TT_GOOD, certs, IDI_CRED_GOOD);

return KHM_ERROR_SUCCESS;

}

void

kca_icon_set_state(khm_handle credset_with_tokens)

{

struct state_data d;

#if KH_VERSION_API < 7

if (pkhui_request_UI_callback == NULL)

return;

#endif

d.credset = credset_with_tokens;

if (notification_icon_added) {

set_state_from_ui_thread(NULL, &d);

} else {

khui_request_UI_callback(set_state_from_ui_thread, &d);

}

}