/*
 * The Cryptonit security software suite is developped by IDEALX
 * Cryptonit Team (http://IDEALX.org/ and http://cryptonit.org).
 *
 * Copyright 2003-2006 IDEALX
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License version
 * 2 as published by the Free Software Foundation.
 *
 * 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., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301, USA.
 */

#define WINVER        0x0500
#define _WIN32_WINNT  0x0500

#include "aes.h"
#include "aes.c"
#define CIPHER_BLOCK_SIZE AES_BLOCK_SIZE
#include "sha2.h"
#include "sha2.c"

#include <windows.h>
#include <commctrl.h>

/* BUGFIX: it seems mingw32 (at least in Debian packaging) forgot to
   define off64_t */
typedef long long off64_t; 
#include <stdio.h>

#include "types.h"
#include "resource.h"

#ifdef DEBUG
void dump(unsigned char *val, int size)
{
  int i;
  for(i=0;i<size;i++) {
    if(i%32 == 0) {
      printf("\n");
    }
    printf("%02x", val[i]);
  }
  printf("\n");
}
#define DUMP(x,y,l) \
  printf("##x ="); dump(y,l); 
#else
#define DUMP(x,y,l)
#endif

#define BLOCK_SIZE 1024

unsigned char password[256];
char szEFileName[MAX_PATH];
char szFileName[MAX_PATH];
#ifdef HARD_WAY
HWND hEdit;
HWND hOk;
HWND hCancel;
#endif

#include "pkcs5.c"

int decrypt(HWND hwnd)
{
  unsigned int size = 0, index = 0, next = 0, delete = 0;
  unsigned int authentication_failed = 0; /* mark wether hmac successfull or not  */
  DWORD foo = 0;
  char *msg = NULL;
  hmac_ctx hmac;
  aes_ctx aes;

  uint64 salt;
  uint8 buffer[BLOCK_SIZE];
  uint8 key[SHA256_DIGEST_LENGTH];
  uint8 hash[SHA256_DIGEST_LENGTH];
  HANDLE hFin = INVALID_HANDLE_VALUE;
  HANDLE hFout = INVALID_HANDLE_VALUE;

  /* Open ourselves for reading */
  if((hFin = CreateFile(szEFileName, GENERIC_READ, FILE_SHARE_READ, NULL,
			OPEN_EXISTING, 0, NULL)) == INVALID_HANDLE_VALUE) {
    msg = "Could not read input file.";
    goto end;
  }


  /* Read original file size stored in last four bytes */
  SetFilePointer(hFin, -(sizeof(unsigned int) + sizeof(uint64)), 0, FILE_END);
  if(ReadFile(hFin, &size, sizeof(unsigned int), &foo, NULL) == 0) {
    msg = "Could not read input file size.";
    goto end;
  }



  /* Read original salt stored in last four bytes */
  if(ReadFile(hFin, &salt, sizeof(uint64), &foo, NULL) == 0) {
    msg = "Could not read salt.";
    goto end;
  }

  

  /* Derive password into key */
  derive_key(password, strlen(password), 
	     (char *)&salt, 8, key);

  DUMP(key, key, SHA256_DIGEST_LENGTH);

  /* Rewinding to start of encrypted file : 
   * = 4 bytes : size of original file
   * + SHA256_DIGEST_LENGTH bytes : hash
   * + size bytes : file size
   * + salt size
   * + padding to CIPHER_BLOCK_SIZE
   */

  index = 4 + SHA256_DIGEST_LENGTH + size + sizeof(uint64);
  if((size % CIPHER_BLOCK_SIZE)) {
    index += CIPHER_BLOCK_SIZE - (size % CIPHER_BLOCK_SIZE);
  }
  SetFilePointer (hFin, -index, 0, FILE_END);


  /* Initializing HMAC using the drived key */
  hmac_init( key, SHA256_DIGEST_LENGTH, &hmac );

  for(index = 0; index < size; index += next) {
    int extra = 0;
    next =  size - index;
    if(next >= BLOCK_SIZE) {
      next = BLOCK_SIZE;
    } else {
      /* Was there any padding ? */
      if(next % CIPHER_BLOCK_SIZE) {
		extra = CIPHER_BLOCK_SIZE - (next % CIPHER_BLOCK_SIZE);
      }
    }
    /* Read encrypted block */
    if(ReadFile(hFin, buffer, (next + extra), &foo, NULL) == 0) {
      msg = "Failed while reading from input file.";
	  delete = 1;
      goto end;
    }
    /* Update the hmac with the encrypted block */
    hmac_update(buffer, (next + extra), &hmac);
  }
  
  /* Finish HMAC */
  hmac_final(hash, &hmac); 
  
  /* Read encrypted hash */
  if(ReadFile(hFin, buffer, SHA256_DIGEST_LENGTH, &foo, NULL) == 0) {
    msg = "Failed while reading HMAC from input file, the file is probably corrupted.";
    /*     delete = 1; */
    authentication_failed = 1;
    goto end;
  }
  
  /* Compare hmac */
  if((memcmp(buffer, hash, SHA256_DIGEST_LENGTH)) != 0) {
    msg = "Your password is invalid, or the file was corrupted (HMAC is invalid).";
    /*     delete = 1; */
    authentication_failed = 1;
    goto end;
  }
  
  /* 
   * decrypt file content and write it to output file, now that
   * we are sure that it can be decrypted
   */
  
  /* Opening output file for writing */
  if((hFout = CreateFile
      (szFileName, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS,
       FILE_ATTRIBUTE_NORMAL, NULL)) == INVALID_HANDLE_VALUE) {
    msg = "Could not create output file.";
    goto end;
  }

  /* Get at the right offset in the file to decypher it */
  index = 4 + SHA256_DIGEST_LENGTH + size + sizeof(uint64);
  if((size % CIPHER_BLOCK_SIZE)) {
    index += CIPHER_BLOCK_SIZE - (size % CIPHER_BLOCK_SIZE);
  }
  SetFilePointer (hFin, -index, 0, FILE_END);

  aes_set_key(&aes, key, 256, 0);
  
  for(index = 0; index < size; index += next) {
    int extra = 0;
    next =  size - index;
    if(next >= BLOCK_SIZE) {
      next = BLOCK_SIZE;
    } else {
      /* Was there any padding ? */
      if(next % CIPHER_BLOCK_SIZE) {
		extra = CIPHER_BLOCK_SIZE - (next % CIPHER_BLOCK_SIZE);
      }
    }

    /* Read encrypted block */
    if(ReadFile(hFin, buffer, (next + extra), &foo, NULL) == 0) {
      msg = "Failed while reading from input file.";
	  delete = 1;
      goto end;
    }

    /* Decrypt block */
    aes_cbc_decrypt(&aes, key, buffer, (next + extra));
    
    /* Write block */
    if(WriteFile(hFout, buffer, next, &foo, NULL) == 0) {
      msg = "Failed while writing to output file.";
      delete = 1;
      goto end;
    }
  }

 end:
  if(hFin != INVALID_HANDLE_VALUE)
    CloseHandle(hFin);
  if(hFout != INVALID_HANDLE_VALUE)
    CloseHandle(hFout);

  if(delete) {
	DeleteFile(szFileName);
  } else 
    if (!authentication_failed) {
      ShellExecute(NULL, "open", szFileName, NULL, NULL, SW_MAXIMIZE);
    }

  if(msg) {
	MessageBox(hwnd, msg, "Warning", MB_OK);
  }

  EndDialog(hwnd, 0);
}

void DoFileSave(HWND hwnd)
{
  OPENFILENAME ofn;

  // When somebody clicks the Ok button, first we get the
  // password they entered
  int len = GetWindowTextLength(GetDlgItem(hwnd, IDC_EDIT));
  if(len > 0 && len < sizeof(password)) {
    GetDlgItemText(hwnd, IDC_EDIT, password, len + 1);

 
    memcpy(szFileName, szEFileName, MAX_PATH);
    len = strlen(szFileName);
    if(szFileName[len-4]=='.' &&
       (szFileName[len-3]=='E' || szFileName[len-3]=='e') &&
       (szFileName[len-2]=='X' || szFileName[len-2]=='x') &&
       (szFileName[len-1]=='E' || szFileName[len-1]=='e')) {
      szFileName[len-4]='\0';
    } else {
      szFileName[0]='\0';
    }
    
#ifdef OPENFILENAME_SIZE_VERSION_400
    ZeroMemory(&ofn, OPENFILENAME_SIZE_VERSION_400);
    ofn.lStructSize = sizeof(OPENFILENAME_SIZE_VERSION_400);
#else
    ZeroMemory(&ofn, sizeof(ofn));
    ofn.lStructSize = sizeof(ofn);
#endif	
    
    ofn.hwndOwner = hwnd;
    ofn.lpstrFilter = "All Files (*.*)\0*.*\0";
    ofn.lpstrFile = szFileName;
    ofn.nMaxFile = MAX_PATH;
    ofn.Flags = OFN_EXPLORER | OFN_HIDEREADONLY | OFN_PATHMUSTEXIST
      | OFN_OVERWRITEPROMPT;
    
    if(GetSaveFileName(&ofn)) {
      decrypt(hwnd);
      PostMessage(hwnd, WM_CLOSE, 0, 0);
    }
  }
}

#ifndef HARD_WAY
BOOL CALLBACK DlgProc(HWND hwnd, UINT Message, WPARAM wParam, LPARAM lParam)
#else
LRESULT CALLBACK WndProc(HWND hwnd, UINT Message, WPARAM wParam, LPARAM lParam)
#endif
{
  static HBITMAP hBitmap;
  BITMAP bitmap;
  static int cxSource, cySource;
  static int cxClient, cyClient;
  HINSTANCE      hInstance ;
  HDC            hdc, hdcMem ;
  int            x, y ;
  PAINTSTRUCT    ps ;

  switch(Message)
    {
#ifndef HARD_WAY
    case WM_INITDIALOG:
      // This is where we set up the dialog box, and initialise any
      // default values
      SetDlgItemText(hwnd, IDC_EDIT, "");
#else

    case WM_CREATE:
      {
//	HFONT hfDefault;

	hEdit = CreateWindowEx(WS_EX_CLIENTEDGE, "EDIT", "", WS_CHILD | WS_VISIBLE | ES_PASSWORD |
			       ES_AUTOHSCROLL | ES_NOHIDESEL,
			       2, 2, 100, 24, hwnd, (HMENU)IDC_EDIT, GetModuleHandle(NULL), NULL);

	hOk = CreateWindowEx(WS_EX_CLIENTEDGE, "BUTTON", "Ok", WS_CHILD | WS_VISIBLE | BS_DEFPUSHBUTTON,
			       102, 2, 60, 24, hwnd, (HMENU)IDC_OK, GetModuleHandle(NULL), NULL);
	hCancel = CreateWindowEx(WS_EX_CLIENTEDGE, "BUTTON", "Cancel", WS_CHILD | WS_VISIBLE,
				 164, 2, 60, 24, hwnd, (HMENU)IDC_CANCEL, GetModuleHandle(NULL), NULL);
//	if(hEdit == NULL)
//	MessageBox(hwnd, "Could not create edit box.", "Error", MB_OK | MB_ICONERROR);
// 	hfDefault = GetStockObject(DEFAULT_GUI_FONT);
// 	SendMessage(hEdit, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
// 	SendMessage(hOk, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));
// 	SendMessage(hCancel, WM_SETFONT, (WPARAM)hfDefault, MAKELPARAM(FALSE, 0));

	SetFocus(hEdit);

	hInstance = ((LPCREATESTRUCT) lParam)->hInstance ;
	hBitmap = LoadBitmap (hInstance, MAKEINTRESOURCE(IDB_BITMAP1)) ;
	GetObject (hBitmap, sizeof (BITMAP), &bitmap) ;
	cxSource = bitmap.bmWidth ;
	cySource = bitmap.bmHeight ;
	
      }
#endif
      return 0;

    case WM_SIZE:
      cxClient = LOWORD (lParam) ;
      cyClient = HIWORD (lParam) ;
      return 0;

    case WM_PAINT:
      hdc = BeginPaint (hwnd, &ps) ;

      hdcMem = CreateCompatibleDC (hdc) ;
      SelectObject (hdcMem, hBitmap) ;
      x = y = 0;
      BitBlt (hdc, 0, 28, cxSource, cySource, hdcMem, 0, 0, SRCCOPY) ;

      DeleteDC (hdcMem) ;
      EndPaint (hwnd, &ps) ;
      return 0;

    case WM_DESTROY:
      DeleteObject (hBitmap) ;
      PostQuitMessage (0) ;
      return 0;

    case WM_COMMAND:
      switch(LOWORD(wParam)) {
      case IDC_OK: 
	DoFileSave(hwnd);
	break;
      case IDC_CANCEL:
	PostMessage(hwnd, WM_CLOSE, 0, 0);
	break;
      }
      break;

    case WM_CLOSE:
      DestroyWindow(hwnd);
      // EndDialog(hwnd, 0);
      return 0;
    }
  return DefWindowProc (hwnd, Message, wParam, lParam) ;
}

const char g_szClassName[] = "myWindowClass";

int STDCALL WinMain (HINSTANCE hInstance, HINSTANCE hPrev,
		     LPSTR lpCmd, int nShow)
{
#ifndef HARD_WAY
    GetModuleFileName(0, szEFileName, MAX_PATH);
    return DialogBox(hInstance, MAKEINTRESOURCE(IDD_DIALOG1),
		     NULL, DlgProc);
#else
  WNDCLASSEX wc;
  HWND hwnd;
  MSG Msg;
  
  GetModuleFileName(0, szEFileName, MAX_PATH);
  wc.cbSize	   = sizeof(WNDCLASSEX);
  wc.style	   = CS_HREDRAW | CS_VREDRAW;
  wc.lpfnWndProc   = WndProc;
  wc.cbClsExtra	   = 0;
  wc.cbWndExtra	   = 0;
  wc.hInstance	   = hInstance;
  wc.hIcon	   = LoadIcon(hInstance, MAKEINTRESOURCE(IDI_SELFEXTRACT));
  wc.hCursor	   = LoadCursor(NULL, IDC_ARROW);
  wc.hbrBackground = (HBRUSH) GetStockObject (WHITE_BRUSH) ;
  wc.lpszClassName = g_szClassName;
  wc.hIconSm	   = LoadIcon(hInstance, MAKEINTRESOURCE(IDI_SELFEXTRACT));

  if(!RegisterClassEx(&wc)) {
    MessageBox(NULL, "Window Registration Failed!", "Error!",
	       MB_ICONEXCLAMATION | MB_OK);
    return 0;
  }
  
  // WS_EX_APPWINDOW |
  // WS_EX_TOPMOST
  // DS_SYSMODAL |  | DS_MODALFRAME
  // DS_ABSALIGN | DS_CENTER | WS_CAPTION | WS_SYSMENU | WS_POPUP


/*    hwnd = CreateWindowEx (0, g_szClassName, "Enter Password:",  */
/*  			 DS_SYSMODAL | DS_CENTER | WS_CAPTION | WS_SYSMENU, CW_USEDEFAULT, CW_USEDEFAULT, 226+4, 24+27, NULL, NULL,  */
/*  			 hInstance, NULL);  */

  hwnd = CreateWindowEx (0, g_szClassName, "Enter Password:",
			 WS_OVERLAPPEDWINDOW | DS_CENTER, CW_USEDEFAULT, CW_USEDEFAULT, 226+4, 137, NULL, NULL,
			 hInstance, NULL);

  
  if(hwnd == NULL) {
    MessageBox(NULL, "Window Creation Failed!", "Error!",
	       MB_ICONEXCLAMATION | MB_OK);
    return 0;
  }
  
  ShowWindow(hwnd, nShow);
  UpdateWindow(hwnd);
  
  while(GetMessage(&Msg, NULL, 0, 0) > 0) {
    TranslateMessage(&Msg);
    /* This is a nasty Hack */
    if((Msg.message == WM_KEYUP) || (Msg.message == WM_KEYDOWN)) {
      if(Msg.wParam == VK_RETURN)
	DoFileSave(hwnd);
    }
    DispatchMessage(&Msg);
  }
  return Msg.wParam;
#endif
}
