/*
 * PuTTY key generation front end (Windows).
 */

#include <time.h>
#include <stdio.h>
#include <stdlib.h>

#define PUTTY_DO_GLOBALS

#include "putty.h"
#include "ssh.h"

#include <commctrl.h>

#ifdef MSVC4
#define ICON_BIG        1
#endif

#define WM_DONEKEY (WM_XUSER + 1)

#define DEFAULT_KEYSIZE 1024

static int requested_help;

static char *cmdline_keyfile = NULL;

/*
 * Print a modal (Really Bad) message box and perform a fatal exit.
 */
void modalfatalbox(char *fmt, ...)
{
    va_list ap;
    char *stuff;

    va_start(ap, fmt);
    stuff = dupvprintf(fmt, ap);
    va_end(ap);
    MessageBox(NULL, stuff, "PuTTYgen Fatal Error",
	       MB_SYSTEMMODAL | MB_ICONERROR | MB_OK);
    sfree(stuff);
    exit(1);
}

/* ----------------------------------------------------------------------
 * Progress report code. This is really horrible :-)
 */
#define PROGRESSRANGE 65535
#define MAXPHASE 5
struct progress {
    int nphases;
    struct {
	int exponential;
	unsigned startpoint, total;
	unsigned param, current, n;    /* if exponential */
	unsigned mult;		       /* if linear */
    } phases[MAXPHASE];
    unsigned total, divisor, range;
    HWND progbar;
};

static void progress_update(void *param, int action, int phase, int iprogress)
{
    struct progress *p = (struct progress *) param;
    unsigned progress = iprogress;
    int position;

    if (action < PROGFN_READY && p->nphases < phase)
	p->nphases = phase;
    switch (action) {
      case PROGFN_INITIALISE:
	p->nphases = 0;
	break;
      case PROGFN_LIN_PHASE:
	p->phases[phase-1].exponential = 0;
	p->phases[phase-1].mult = p->phases[phase].total / progress;
	break;
      case PROGFN_EXP_PHASE:
	p->phases[phase-1].exponential = 1;
	p->phases[phase-1].param = 0x10000 + progress;
	p->phases[phase-1].current = p->phases[phase-1].total;
	p->phases[phase-1].n = 0;
	break;
      case PROGFN_PHASE_EXTENT:
	p->phases[phase-1].total = progress;
	break;
      case PROGFN_READY:
	{
	    unsigned total = 0;
	    int i;
	    for (i = 0; i < p->nphases; i++) {
		p->phases[i].startpoint = total;
		total += p->phases[i].total;
	    }
	    p->total = total;
	    p->divisor = ((p->total + PROGRESSRANGE - 1) / PROGRESSRANGE);
	    p->range = p->total / p->divisor;
	    SendMessage(p->progbar, PBM_SETRANGE, 0, MAKELPARAM(0, p->range));
	}
	break;
      case PROGFN_PROGRESS:
	if (p->phases[phase-1].exponential) {
	    while (p->phases[phase-1].n < progress) {
		p->phases[phase-1].n++;
		p->phases[phase-1].current *= p->phases[phase-1].param;
		p->phases[phase-1].current /= 0x10000;
	    }
	    position = (p->phases[phase-1].startpoint +
			p->phases[phase-1].total - p->phases[phase-1].current);
	} else {
	    position = (p->phases[phase-1].startpoint +
			progress * p->phases[phase-1].mult);
	}
	SendMessage(p->progbar, PBM_SETPOS, position / p->divisor, 0);
	break;
    }
}

extern char ver[];

#define PASSPHRASE_MAXLEN 512

struct PassphraseProcStruct {
    char *passphrase;
    char *comment;
};

/*
 * Dialog-box function for the passphrase box.
 */
static int CALLBACK PassphraseProc(HWND hwnd, UINT msg,
				   WPARAM wParam, LPARAM lParam)
{
    static char *passphrase = NULL;
    struct PassphraseProcStruct *p;

    switch (msg) {
      case WM_INITDIALOG:
	SetForegroundWindow(hwnd);
	SetWindowPos(hwnd, HWND_TOP, 0, 0, 0, 0,
		     SWP_NOMOVE | SWP_NOSIZE | SWP_SHOWWINDOW);

	/*
	 * Centre the window.
	 */
	{			       /* centre the window */
	    RECT rs, rd;
	    HWND hw;

	    hw = GetDesktopWindow();
	    if (GetWindowRect(hw, &rs) && GetWindowRect(hwnd, &rd))
		MoveWindow(hwnd,
			   (rs.right + rs.left + rd.left - rd.right) / 2,
			   (rs.bottom + rs.top + rd.top - rd.bottom) / 2,
			   rd.right - rd.left, rd.bottom - rd.top, TRUE);
	}

	p = (struct PassphraseProcStruct *) lParam;
	passphrase = p->passphrase;
	if (p->comment)
	    SetDlgItemText(hwnd, 101, p->comment);
	*passphrase = 0;
	SetDlgItemText(hwnd, 102, passphrase);
	return 0;
      case WM_COMMAND:
	switch (LOWORD(wParam)) {
	  case IDOK:
	    if (*passphrase)
		EndDialog(hwnd, 1);
	    else
		MessageBeep(0);
	    return 0;
	  case IDCANCEL:
	    EndDialog(hwnd, 0);
	    return 0;
	  case 102:		       /* edit box */
	    if ((HIWORD(wParam) == EN_CHANGE) && passphrase) {
		GetDlgItemText(hwnd, 102, passphrase,
			       PASSPHRASE_MAXLEN - 1);
		passphrase[PASSPHRASE_MAXLEN - 1] = '\0';
	    }
	    return 0;
	}
	return 0;
      case WM_CLOSE:
	EndDialog(hwnd, 0);
	return 0;
    }
    return 0;
}

/*
 * Prompt for a key file. Assumes the filename buffer is of size
 * FILENAME_MAX.
 */
static int prompt_keyfile(HWND hwnd, char *dlgtitle,
			  char *filename, int save, int ppk)
{
    OPENFILENAME of;
    memset(&of, 0, sizeof(of));
#ifdef OPENFILENAME_SIZE_VERSION_400
    of.lStructSize = OPENFILENAME_SIZE_VERSION_400;
#else
    of.lStructSize = sizeof(of);
#endif
    of.hwndOwner = hwnd;
    if (ppk) {
	of.lpstrFilter = "PuTTY Private Key Files (*.ppk)\0*.ppk\0"
	    "All Files (*.*)\0*\0\0\0";
	of.lpstrDefExt = ".ppk";
    } else {
	of.lpstrFilter = "All Files (*.*)\0*\0\0\0";
    }
    of.lpstrCustomFilter = NULL;
    of.nFilterIndex = 1;
    of.lpstrFile = filename;
    *filename = '\0';
    of.nMaxFile = FILENAME_MAX;
    of.lpstrFileTitle = NULL;
    of.lpstrInitialDir = NULL;
    of.lpstrTitle = dlgtitle;
    of.Flags = 0;
    if (save)
	return GetSaveFileName(&of);
    else
	return GetOpenFileName(&of);
}

/*
 * Dialog-box function for the Licence box.
 */
static int CALLBACK LicenceProc(HWND hwnd, UINT msg,
				WPARAM wParam, LPARAM lParam)
{
    switch (msg) {
      case WM_INITDIALOG:
	/*
	 * Centre the window.
	 */
	{			       /* centre the window */
	    RECT rs, rd;
	    HWND hw;

	    hw = GetDesktopWindow();
	    if (GetWindowRect(hw, &rs) && GetWindowRect(hwnd, &rd))
		MoveWindow(hwnd,
			   (rs.right + rs.left + rd.left - rd.right) / 2,
			   (rs.bottom + rs.top + rd.top - rd.bottom) / 2,
			   rd.right - rd.left, rd.bottom - rd.top, TRUE);
	}

	return 1;
      case WM_COMMAND:
	switch (LOWORD(wParam)) {
	  case IDOK:
	    EndDialog(hwnd, 1);
	    return 0;
	}
	return 0;
      case WM_CLOSE:
	EndDialog(hwnd, 1);
	return 0;
    }
    return 0;
}

/*
 * Dialog-box function for the About box.
 */
static int CALLBACK AboutProc(HWND hwnd, UINT msg,
			      WPARAM wParam, LPARAM lParam)
{
    switch (msg) {
      case WM_INITDIALOG:
	/*
	 * Centre the window.
	 */
	{			       /* centre the window */
	    RECT rs, rd;
	    HWND hw;

	    hw = GetDesktopWindow();
	    if (GetWindowRect(hw, &rs) && GetWindowRect(hwnd, &rd))
		MoveWindow(hwnd,
			   (rs.right + rs.left + rd.left - rd.right) / 2,
			   (rs.bottom + rs.top + rd.top - rd.bottom) / 2,
			   rd.right - rd.left, rd.bottom - rd.top, TRUE);
	}

	SetDlgItemText(hwnd, 100, ver);
	return 1;
      case WM_COMMAND:
	switch (LOWORD(wParam)) {
	  case IDOK:
	    EndDialog(hwnd, 1);
	    return 0;
	  case 101:
	    EnableWindow(hwnd, 0);
	    DialogBox(hinst, MAKEINTRESOURCE(214), hwnd, LicenceProc);
	    EnableWindow(hwnd, 1);
	    SetActiveWindow(hwnd);
	    return 0;
	}
	return 0;
      case WM_CLOSE:
	EndDialog(hwnd, 1);
	return 0;
    }
    return 0;
}

/*
 * Thread to generate a key.
 */
struct rsa_key_thread_params {
    HWND progressbar;		       /* notify this with progress */
    HWND dialog;		       /* notify this on completion */
    int keysize;		       /* bits in key */
    int is_dsa;
    struct RSAKey *key;
    struct dss_key *dsskey;
};
static DWORD WINAPI generate_rsa_key_thread(void *param)
{
    struct rsa_key_thread_params *params =
	(struct rsa_key_thread_params *) param;
    struct progress prog;
    prog.progbar = params->progressbar;

    progress_update(&prog, PROGFN_INITIALISE, 0, 0);

    if (params->is_dsa)
	dsa_generate(params->dsskey, params->keysize, progress_update, &prog);
    else
	rsa_generate(params->key, params->keysize, progress_update, &prog);

    PostMessage(params->dialog, WM_DONEKEY, 0, 0);

    sfree(params);
    return 0;
}

struct MainDlgState {
    int collecting_entropy;
    int generation_thread_exists;
    int key_exists;
    int entropy_got, entropy_required, entropy_size;
    int keysize;
    int ssh2, is_dsa;
    char **commentptr;		       /* points to key.comment or ssh2key.comment */
    struct ssh2_userkey ssh2key;
    unsigned *entropy;
    struct RSAKey key;
    struct dss_key dsskey;
    HMENU filemenu, keymenu, cvtmenu;
};

static void hidemany(HWND hwnd, const int *ids, int hideit)
{
    while (*ids) {
	ShowWindow(GetDlgItem(hwnd, *ids++), (hideit ? SW_HIDE : SW_SHOW));
    }
}

static void setupbigedit1(HWND hwnd, int id, int idstatic, struct RSAKey *key)
{
    char *buffer;
    char *dec1, *dec2;

    dec1 = bignum_decimal(key->exponent);
    dec2 = bignum_decimal(key->modulus);
    buffer = dupprintf("%d %s %s %s", bignum_bitcount(key->modulus),
		       dec1, dec2, key->comment);
    SetDlgItemText(hwnd, id, buffer);
    SetDlgItemText(hwnd, idstatic,
		   "&Public key for pasting into authorized_keys file:");
    sfree(dec1);
    sfree(dec2);
    sfree(buffer);
}

static void setupbigedit2(HWND hwnd, int id, int idstatic,
			  struct ssh2_userkey *key)
{
    unsigned char *pub_blob;
    char *buffer, *p;
    int pub_len;
    int i;

    pub_blob = key->alg->public_blob(key->data, &pub_len);
    buffer = snewn(strlen(key->alg->name) + 4 * ((pub_len + 2) / 3) +
		   strlen(key->comment) + 3, char);
    strcpy(buffer, key->alg->name);
    p = buffer + strlen(buffer);
    *p++ = ' ';
    i = 0;
    while (i < pub_len) {
	int n = (pub_len - i < 3 ? pub_len - i : 3);
	base64_encode_atom(pub_blob + i, n, p);
	i += n;
	p += 4;
    }
    *p++ = ' ';
    strcpy(p, key->comment);
    SetDlgItemText(hwnd, id, buffer);
    SetDlgItemText(hwnd, idstatic, "&Public key for pasting into "
		   "OpenSSH authorized_keys file:");
    sfree(pub_blob);
    sfree(buffer);
}

static int save_ssh1_pubkey(char *filename, struct RSAKey *key)
{
    char *dec1, *dec2;
    FILE *fp;

    dec1 = bignum_decimal(key->exponent);
    dec2 = bignum_decimal(key->modulus);
    fp = fopen(filename, "wb");
    if (!fp)
	return 0;
    fprintf(fp, "%d %s %s %s\n",
	    bignum_bitcount(key->modulus), dec1, dec2, key->comment);
    fclose(fp);
    sfree(dec1);
    sfree(dec2);
    return 1;
}

/*
 * Warn about the obsolescent key file format.
 */
void old_keyfile_warning(void)
{
    static const char mbtitle[] = "PuTTY Key File Warning";
    static const char message[] =
	"You are loading an SSH 2 private key which has an\n"
	"old version of the file format. This means your key\n"
	"file is not fully tamperproof. Future versions of\n"
	"PuTTY may stop supporting this private key format,\n"
	"so we recommend you convert your key to the new\n"
	"format.\n"
	"\n"
	"Once the key is loaded into PuTTYgen, you can perform\n"
	"this conversion simply by saving it again.";

    MessageBox(NULL, message, mbtitle, MB_OK);
}

static int save_ssh2_pubkey(char *filename, struct ssh2_userkey *key)
{
    unsigned char *pub_blob;
    char *p;
    int pub_len;
    int i, column;
    FILE *fp;

    pub_blob = key->alg->public_blob(key->data, &pub_len);

    fp = fopen(filename, "wb");
    if (!fp)
	return 0;

    fprintf(fp, "---- BEGIN SSH2 PUBLIC KEY ----\n");

    fprintf(fp, "Comment: \"");
    for (p = key->comment; *p; p++) {
	if (*p == '\\' || *p == '\"')
	    fputc('\\', fp);
	fputc(*p, fp);
    }
    fprintf(fp, "\"\n");

    i = 0;
    column = 0;
    while (i < pub_len) {
	char buf[5];
	int n = (pub_len - i < 3 ? pub_len - i : 3);
	base64_encode_atom(pub_blob + i, n, buf);
	i += n;
	buf[4] = '\0';
	fputs(buf, fp);
	if (++column >= 16) {
	    fputc('\n', fp);