// TortoiseSVN - a Windows shell extension for easy version control

// Copyright (C) 2003 - Tim Kemp and Stefan Kueng

// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 2
// of the License, or (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

#include "LoginDialog.h"
#include "TortoisePlinkRes.h"
#include <string>

#include <wincrypt.h>
#pragma comment(lib, "Crypt32.lib")
#define MY_ENCODING_TYPE  (PKCS_7_ASN_ENCODING | X509_ASN_ENCODING)


HINSTANCE g_hmodThisDll;
HWND g_hwndMain;

class LoginDialog
{
public:
   LoginDialog(const std::string& prompt, bool is_pw);
   
   static bool DoLoginDialog(std::string& password, const std::string& prompt, bool& save_pw, bool is_pw);

private:
   bool myOK;
   bool savePassword;
   HWND _hdlg;

   std::string  myPassword;
   std::string  myPrompt;
   bool         myIsPW;
   
   void CreateModule(void);
   void RetrieveValues();
   
   std::string GetPassword();

   friend BOOL CALLBACK LoginDialogProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam);
};


BOOL DoLoginDialog(char* password, int maxlen, const char* prompt, BOOL is_pw)
{
   BOOL res = FALSE;
   bool save_pw = false;
   LPWSTR pDescrOut = NULL;
   DATA_BLOB encrypted, decrypted;
   //CRYPTPROTECT_PROMPTSTRUCT PromptStruct;
   std::string passwordstr;

   ZeroMemory(&decrypted, sizeof(decrypted));
   /*ZeroMemory(&PromptStruct, sizeof(PromptStruct));
   PromptStruct.cbSize = sizeof(PromptStruct);
   PromptStruct.dwPromptFlags = CRYPTPROTECT_PROMPT_ON_PROTECT | CRYPTPROTECT_PROMPT_ON_UNPROTECT;
   PromptStruct.szPrompt = L"Tortoise PLINK SSH Password";*/

   // open tortoise registry key
   HKEY hKey = NULL;
   RegOpenKeyEx(HKEY_CURRENT_USER, "SOFTWARE\\TortoiseSVN", 0, KEY_SET_VALUE | KEY_QUERY_VALUE, &hKey);

   DWORD regType = REG_BINARY, dataLen = 0;
   // the item is not a password, or cannot open registry key or there is no password on registry, show dialog
   if (!is_pw || !hKey || RegQueryValueEx(hKey, "PLINK_Password", NULL, &regType, NULL, &dataLen))
   {
	   g_hmodThisDll = GetModuleHandle(0);
	   g_hwndMain = GetParentHwnd();
	   res = LoginDialog::DoLoginDialog(passwordstr, prompt, save_pw, is_pw ? true : false);
	   if (res)
		   strncpy(password, passwordstr.c_str(), maxlen);
   }
   else
   {
	    // gets encrypted password from registry
	    encrypted.pbData = (BYTE*)LocalAlloc(LMEM_ZEROINIT, dataLen);
		encrypted.cbData = dataLen;
		RegQueryValueEx(hKey, "PLINK_Password", NULL, &regType, encrypted.pbData, &dataLen);

		// decrypt the password
		res = FALSE;
		if (CryptUnprotectData(
			&encrypted,
			NULL,
			NULL,                 // Optional entropy
			NULL,                 // Reserved
			NULL,        // Optional PromptStruct
			0,
			&decrypted))
		{
			strncpy(password, (char*)decrypted.pbData, maxlen);
			LocalFree(encrypted.pbData);
			LocalFree(decrypted.pbData);
			RegCloseKey(hKey);
			return TRUE;
		}
		else
		{
			// failed decripting ? Show dialog asking password
			g_hmodThisDll = GetModuleHandle(0);
			g_hwndMain = GetParentHwnd();
			std::string passwordstr;
			res = LoginDialog::DoLoginDialog(passwordstr, prompt, save_pw, is_pw ? true : false);
		}
   }

   // succedded and is a password and user wants save password in registry ? Do it
   if (res && is_pw && save_pw)
   {
	   strncpy(password, passwordstr.c_str(), maxlen);
	   decrypted.pbData = (BYTE*)passwordstr.c_str();
	   decrypted.cbData = passwordstr.length() + 1;

	   if(CryptProtectData(
		   &decrypted,
		   L"Tortoise PLINK SSH password.",     // A description string. 
		   NULL,                               // Optional entropy not used.
		   NULL,                               // Reserved.
		   NULL,                      // Pass a PromptStruct.
		   0,
		   &encrypted))
	   {
		   RegSetValueEx(hKey, "PLINK_Password", NULL, REG_BINARY, encrypted.pbData, encrypted.cbData);
	   }

	   decrypted.pbData = NULL; // not allocated by LocalAlloc !
   }

   LocalFree(encrypted.pbData);
   LocalFree(decrypted.pbData);
   RegCloseKey(hKey);
   return res;
}


BOOL CALLBACK LoginDialogProc(HWND hwndDlg, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
   if (uMsg == WM_INITDIALOG)
   {
      LoginDialog* pDlg = (LoginDialog*) lParam;
      pDlg->_hdlg = hwndDlg;
      SetWindowLong(hwndDlg, GWL_USERDATA, lParam);
      // Set prompt text
      SendDlgItemMessage(hwndDlg, IDC_LOGIN_PROMPT, WM_SETTEXT,
                         pDlg->myPrompt.length(), (LPARAM) pDlg->myPrompt.c_str());
      // Make sure edit control has the focus
      //SendDlgItemMessage(hwndDlg, IDC_LOGIN_PASSWORD, WM_SETFOCUS, 0, 0);
      if (GetDlgCtrlID((HWND) wParam) != IDC_LOGIN_PASSWORD)
      { 
         SetFocus(GetDlgItem(hwndDlg, IDC_LOGIN_PASSWORD));
         return FALSE; 
      } 
      return TRUE; 
   }
   else if (uMsg == WM_COMMAND && LOWORD(wParam) == IDCANCEL && HIWORD(wParam) == BN_CLICKED)
   {
      LoginDialog* pDlg = (LoginDialog*) GetWindowLong(hwndDlg, GWL_USERDATA);
      pDlg->myOK = false;
      EndDialog(hwndDlg, IDCANCEL);
      return 1;
   }
   else if (uMsg == WM_COMMAND && LOWORD(wParam) == IDOK && HIWORD(wParam) == BN_CLICKED)
   {
      LoginDialog* pDlg = (LoginDialog*) GetWindowLong(hwndDlg, GWL_USERDATA);
      pDlg->myOK = true;
      pDlg->RetrieveValues();
      EndDialog(hwndDlg, IDOK);
      return 1;
   }
   return 0;
}

LoginDialog::LoginDialog(const std::string& prompt, bool is_pw)
{
   myPrompt = prompt;
   myIsPW = is_pw;
   savePassword = false;
}

void LoginDialog::CreateModule(void)
{
   DialogBoxParam(g_hmodThisDll, MAKEINTRESOURCE(IDD_LOGIN), g_hwndMain,
                  (DLGPROC)(LoginDialogProc), (long)this);
}


bool LoginDialog::DoLoginDialog(std::string& password, const std::string& prompt, bool& save_pw, bool is_pw)
{
   LoginDialog *pDlg = new LoginDialog(prompt, is_pw);

   pDlg->CreateModule();

   bool ret = pDlg->myOK;
   password = pDlg->myPassword;
   save_pw = pDlg->savePassword;
   
   delete pDlg;

   return ret;
}


std::string LoginDialog::GetPassword()
{
   char szTxt[256];
   SendDlgItemMessage(_hdlg, IDC_LOGIN_PASSWORD, WM_GETTEXT, sizeof(szTxt), (LPARAM)szTxt);
   std::string strText = szTxt;
   return strText;
}

void LoginDialog::RetrieveValues()
{
   myPassword = GetPassword();
   savePassword = IsDlgButtonChecked(_hdlg, IDC_REMEMBER_PASSWORD) == BST_CHECKED;
}


BOOL IsWinNT()
{
   OSVERSIONINFO vi;
   vi.dwOSVersionInfoSize = sizeof(vi);
   if (GetVersionEx(&vi))
   {
      if (vi.dwPlatformId == VER_PLATFORM_WIN32_NT)
      {
         return TRUE;
      }
   }
   return FALSE;
}

HWND GetParentHwnd()
{
   if (IsWinNT())
   {
      return GetDesktopWindow();
   }
   else
   {
      return GetForegroundWindow();
   }
}