Introduced wrapper macros snew(), snewn() and sresize() for the
[u/mdw/putty] / pageant.c
index 2cfcdd0..41ae837 100644 (file)
--- a/pageant.c
+++ b/pageant.c
@@ -9,10 +9,14 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <ctype.h>
+#include <assert.h>
 #include <tchar.h>
 
+#include "putty.h"
 #include "ssh.h"
+#include "misc.h"
 #include "tree234.h"
+#include "winstuff.h"
 
 #define IDI_MAINICON 200
 #define IDI_TRAYICON 201
 #define IDM_CLOSE    0x0010
 #define IDM_VIEWKEYS 0x0020
 #define IDM_ADDKEY   0x0030
-#define IDM_ABOUT    0x0040
+#define IDM_HELP     0x0040
+#define IDM_ABOUT    0x0050
 
 #define APPNAME "Pageant"
 
 extern char ver[];
 
 static HINSTANCE instance;
-static HWND hwnd;
+static HWND main_hwnd;
 static HWND keylist;
 static HWND aboutbox;
-static HMENU systray_menu;
+static HMENU systray_menu, session_menu;
 static int already_running;
+static int requested_help;
+
+static char *help_path;
+static char *putty_path;
+
+#define IDM_PUTTY         0x0060
+#define IDM_SESSIONS_BASE 0x1000
+#define IDM_SESSIONS_MAX  0x2000
+#define PUTTY_REGKEY      "Software\\SimonTatham\\PuTTY\\Sessions"
+#define PUTTY_DEFAULT     "Default%20Settings"
+static int initial_menuitems_count;
+
+/*
+ * Print a modal (Really Bad) message box and perform a fatal exit.
+ */
+void modalfatalbox(char *fmt, ...)
+{
+    va_list ap;
+    char *buf;
+
+    va_start(ap, fmt);
+    buf = dupvprintf(fmt, ap);
+    va_end(ap);
+    MessageBox(main_hwnd, buf, "Pageant Fatal Error",
+              MB_SYSTEMMODAL | MB_ICONERROR | MB_OK);
+    sfree(buf);
+    exit(1);
+}
+
+/* Un-munge session names out of the registry. */
+static void unmungestr(char *in, char *out, int outlen)
+{
+    while (*in) {
+       if (*in == '%' && in[1] && in[2]) {
+           int i, j;
+
+           i = in[1] - '0';
+           i -= (i > 9 ? 7 : 0);
+           j = in[2] - '0';
+           j -= (j > 9 ? 7 : 0);
+
+           *out++ = (i << 4) + j;
+           if (!--outlen)
+               return;
+           in += 3;
+       } else {
+           *out++ = *in++;
+           if (!--outlen)
+               return;
+       }
+    }
+    *out = '\0';
+    return;
+}
 
 static tree234 *rsakeys, *ssh2keys;
 
@@ -61,22 +120,12 @@ void agent_query(void *in, int inlen, void **out, int *outlen);
 int agent_exists(void);
 
 /*
- * We need this to link with the RSA code, because rsaencrypt()
- * pads its data with random bytes. Since we only use rsadecrypt()
- * and the signing functions, which are deterministic, this should
- * never be called.
- *
- * If it _is_ called, there is a _serious_ problem, because it
- * won't generate true random numbers. So we must scream, panic,
- * and exit immediately if that should happen.
+ * Forward references
  */
-int random_byte(void)
-{
-    MessageBox(hwnd, "Internal Error", APPNAME, MB_OK | MB_ICONERROR);
-    exit(0);
-    /* this line can't be reached but it placates MSVC's warnings :-) */
-    return 0;
-}
+static void *make_keylist1(int *length);
+static void *make_keylist2(int *length);
+static void *get_keylist1(void);
+static void *get_keylist2(void);
 
 /*
  * Blob structure for passing to the asymmetric SSH2 key compare
@@ -88,14 +137,6 @@ struct blob {
 };
 static int cmpkeys_ssh2_asymm(void *av, void *bv);
 
-/*
- * This function is needed to link with the DES code. We need not
- * have it do anything at all.
- */
-void logevent(char *msg)
-{
-}
-
 #define GET_32BIT(cp) \
     (((unsigned long)(unsigned char)(cp)[0] << 24) | \
     ((unsigned long)(unsigned char)(cp)[1] << 16) | \
@@ -123,7 +164,6 @@ static tree234 *passphrases = NULL;
  */
 static void forget_passphrases(void)
 {
-    int i;
     while (count234(passphrases) > 0) {
        char *pp = index234(passphrases, 0);
        memset(pp, 0, strlen(pp));
@@ -173,7 +213,7 @@ static int CALLBACK AboutProc(HWND hwnd, UINT msg,
            return 0;
          case 101:
            EnableWindow(hwnd, 0);
-           DialogBox(instance, MAKEINTRESOURCE(214), NULL, LicenceProc);
+           DialogBox(instance, MAKEINTRESOURCE(214), hwnd, LicenceProc);
            EnableWindow(hwnd, 1);
            SetActiveWindow(hwnd);
            return 0;
@@ -333,7 +373,7 @@ static void keylist_update(void)
 /*
  * This function loads a key from a file and adds it.
  */
-static void add_keyfile(char *filename)
+static void add_keyfile(Filename filename)
 {
     char passphrase[PASSPHRASE_MAXLEN];
     struct RSAKey *rkey = NULL;
@@ -343,23 +383,83 @@ static void add_keyfile(char *filename)
     int attempts;
     char *comment;
     struct PassphraseProcStruct pps;
-    int ver;
+    int type;
     int original_pass;
        
-    ver = keyfile_version(filename);
-    if (ver == 0) {
-       MessageBox(NULL, "Couldn't load private key.", APPNAME,
-                  MB_OK | MB_ICONERROR);
+    type = key_type(&filename);
+    if (type != SSH_KEYTYPE_SSH1 && type != SSH_KEYTYPE_SSH2) {
+       char msg[256];
+       sprintf(msg, "Couldn't load this key (%s)", key_type_to_str(type));
+       MessageBox(NULL, msg, APPNAME, MB_OK | MB_ICONERROR);
        return;
     }
 
-    if (ver == 1)
-       needs_pass = rsakey_encrypted(filename, &comment);
+    /*
+     * See if the key is already loaded (in the primary Pageant,
+     * which may or may not be us).
+     */
+    {
+       void *blob;
+       unsigned char *keylist, *p;
+       int i, nkeys, bloblen;
+
+       if (type == SSH_KEYTYPE_SSH1) {
+           if (!rsakey_pubblob(&filename, &blob, &bloblen)) {
+               MessageBox(NULL, "Couldn't load private key.", APPNAME,
+                          MB_OK | MB_ICONERROR);
+               return;
+           }
+           keylist = get_keylist1();
+       } else {
+           unsigned char *blob2;
+           blob = ssh2_userkey_loadpub(&filename, NULL, &bloblen);
+           if (!blob) {
+               MessageBox(NULL, "Couldn't load private key.", APPNAME,
+                          MB_OK | MB_ICONERROR);
+               return;
+           }
+           /* For our purposes we want the blob prefixed with its length */
+           blob2 = snewn(bloblen+4, unsigned char);
+           PUT_32BIT(blob2, bloblen);
+           memcpy(blob2 + 4, blob, bloblen);
+           sfree(blob);
+           blob = blob2;
+
+           keylist = get_keylist2();
+       }
+       if (keylist) {
+           nkeys = GET_32BIT(keylist);
+           p = keylist + 4;
+
+           for (i = 0; i < nkeys; i++) {
+               if (!memcmp(blob, p, bloblen)) {
+                   /* Key is already present; we can now leave. */
+                   sfree(keylist);
+                   sfree(blob);
+                   return;
+               }
+               /* Now skip over public blob */
+               if (type == SSH_KEYTYPE_SSH1)
+                   p += rsa_public_blob_len(p);
+               else
+                   p += 4 + GET_32BIT(p);
+               /* Now skip over comment field */
+               p += 4 + GET_32BIT(p);
+           }
+
+           sfree(keylist);
+       }
+
+       sfree(blob);
+    }
+
+    if (type == SSH_KEYTYPE_SSH1)
+       needs_pass = rsakey_encrypted(&filename, &comment);
     else
-       needs_pass = ssh2_userkey_encrypted(filename, &comment);
+       needs_pass = ssh2_userkey_encrypted(&filename, &comment);
     attempts = 0;
-    if (ver == 1)
-       rkey = smalloc(sizeof(*rkey));
+    if (type == SSH_KEYTYPE_SSH1)
+       rkey = snew(struct RSAKey);
     pps.passphrase = passphrase;
     pps.comment = comment;
     original_pass = 0;
@@ -378,17 +478,17 @@ static void add_keyfile(char *filename)
                if (!dlgret) {
                    if (comment)
                        sfree(comment);
-                   if (ver == 1)
+                   if (type == SSH_KEYTYPE_SSH1)
                        sfree(rkey);
                    return;                    /* operation cancelled */
                }
            }
        } else
            *passphrase = '\0';
-       if (ver == 1)
-           ret = loadrsakey(filename, rkey, passphrase);
+       if (type == SSH_KEYTYPE_SSH1)
+           ret = loadrsakey(&filename, rkey, passphrase);
        else {
-           skey = ssh2_load_userkey(filename, passphrase);
+           skey = ssh2_load_userkey(&filename, passphrase);
            if (skey == SSH2_WRONG_PASSPHRASE)
                ret = -1;
            else if (!skey)
@@ -410,11 +510,11 @@ static void add_keyfile(char *filename)
     if (ret == 0) {
        MessageBox(NULL, "Couldn't load private key.", APPNAME,
                   MB_OK | MB_ICONERROR);
-       if (ver == 1)
+       if (type == SSH_KEYTYPE_SSH1)
            sfree(rkey);
        return;
     }
-    if (ver == 1) {
+    if (type == SSH_KEYTYPE_SSH1) {
        if (already_running) {
            unsigned char *request, *response;
            void *vresponse;
@@ -432,7 +532,7 @@ static void add_keyfile(char *filename)
                ssh1_bignum_length(rkey->q) + 4 + clen  /* comment */
                ;
 
-           request = smalloc(reqlen);
+           request = snewn(reqlen, unsigned char);
 
            request[4] = SSH1_AGENTC_ADD_RSA_IDENTITY;
            reqlen = 5;
@@ -457,6 +557,9 @@ static void add_keyfile(char *filename)
                MessageBox(NULL, "The already running Pageant "
                           "refused to add the key.", APPNAME,
                           MB_OK | MB_ICONERROR);
+
+           sfree(request);
+           sfree(response);
        } else {
            if (add234(rsakeys, rkey) != rkey)
                sfree(rkey);           /* already present, don't waste RAM */
@@ -477,7 +580,7 @@ static void add_keyfile(char *filename)
                4 + clen               /* comment */
                ;
 
-           request = smalloc(reqlen);
+           request = snewn(reqlen, unsigned char);
 
            request[4] = SSH2_AGENTC_ADD_IDENTITY;
            reqlen = 5;
@@ -496,9 +599,12 @@ static void add_keyfile(char *filename)
            agent_query(request, reqlen, &vresponse, &resplen);
            response = vresponse;
            if (resplen < 5 || response[4] != SSH_AGENT_SUCCESS)
-               MessageBox(NULL, "The already running Pageant"
+               MessageBox(NULL, "The already running Pageant "
                           "refused to add the key.", APPNAME,
                           MB_OK | MB_ICONERROR);
+
+           sfree(request);
+           sfree(response);
        } else {
            if (add234(ssh2keys, skey) != skey) {
                skey->alg->freekey(skey->data);
@@ -509,6 +615,162 @@ static void add_keyfile(char *filename)
 }
 
 /*
+ * Create an SSH1 key list in a malloc'ed buffer; return its
+ * length.
+ */
+static void *make_keylist1(int *length)
+{
+    int i, nkeys, len;
+    struct RSAKey *key;
+    unsigned char *blob, *p, *ret;
+    int bloblen;
+
+    /*
+     * Count up the number and length of keys we hold.
+     */
+    len = 4;
+    nkeys = 0;
+    for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
+       nkeys++;
+       blob = rsa_public_blob(key, &bloblen);
+       len += bloblen;
+       sfree(blob);
+       len += 4 + strlen(key->comment);
+    }
+
+    /* Allocate the buffer. */
+    p = ret = snewn(len, unsigned char);
+    if (length) *length = len;
+
+    PUT_32BIT(p, nkeys);
+    p += 4;
+    for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
+       blob = rsa_public_blob(key, &bloblen);
+       memcpy(p, blob, bloblen);
+       p += bloblen;
+       sfree(blob);
+       PUT_32BIT(p, strlen(key->comment));
+       memcpy(p + 4, key->comment, strlen(key->comment));
+       p += 4 + strlen(key->comment);
+    }
+
+    assert(p - ret == len);
+    return ret;
+}
+
+/*
+ * Create an SSH2 key list in a malloc'ed buffer; return its
+ * length.
+ */
+static void *make_keylist2(int *length)
+{
+    struct ssh2_userkey *key;
+    int i, len, nkeys;
+    unsigned char *blob, *p, *ret;
+    int bloblen;
+
+    /*
+     * Count up the number and length of keys we hold.
+     */
+    len = 4;
+    nkeys = 0;
+    for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
+       nkeys++;
+       len += 4;              /* length field */
+       blob = key->alg->public_blob(key->data, &bloblen);
+       len += bloblen;
+       sfree(blob);
+       len += 4 + strlen(key->comment);
+    }
+
+    /* Allocate the buffer. */
+    p = ret = snewn(len, unsigned char);
+    if (length) *length = len;
+
+    /*
+     * Packet header is the obvious five bytes, plus four
+     * bytes for the key count.
+     */
+    PUT_32BIT(p, nkeys);
+    p += 4;
+    for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
+       blob = key->alg->public_blob(key->data, &bloblen);
+       PUT_32BIT(p, bloblen);
+       p += 4;
+       memcpy(p, blob, bloblen);
+       p += bloblen;
+       sfree(blob);
+       PUT_32BIT(p, strlen(key->comment));
+       memcpy(p + 4, key->comment, strlen(key->comment));
+       p += 4 + strlen(key->comment);
+    }
+
+    assert(p - ret == len);
+    return ret;
+}
+
+/*
+ * Acquire a keylist1 from the primary Pageant; this means either
+ * calling make_keylist1 (if that's us) or sending a message to the
+ * primary Pageant (if it's not).
+ */
+static void *get_keylist1(void)
+{
+    void *ret;
+
+    if (already_running) {
+       unsigned char request[5], *response;
+       void *vresponse;
+       int resplen;
+       request[4] = SSH1_AGENTC_REQUEST_RSA_IDENTITIES;
+       PUT_32BIT(request, 4);
+
+       agent_query(request, 5, &vresponse, &resplen);
+       response = vresponse;
+       if (resplen < 5 || response[4] != SSH1_AGENT_RSA_IDENTITIES_ANSWER)
+           return NULL;
+
+       ret = snewn(resplen-5, unsigned char);
+       memcpy(ret, response+5, resplen-5);
+       sfree(response);
+    } else {
+       ret = make_keylist1(NULL);
+    }
+    return ret;
+}
+
+/*
+ * Acquire a keylist2 from the primary Pageant; this means either
+ * calling make_keylist2 (if that's us) or sending a message to the
+ * primary Pageant (if it's not).
+ */
+static void *get_keylist2(void)
+{
+    void *ret;
+
+    if (already_running) {
+       unsigned char request[5], *response;
+       void *vresponse;
+       int resplen;
+
+       request[4] = SSH2_AGENTC_REQUEST_IDENTITIES;
+       PUT_32BIT(request, 4);
+
+       agent_query(request, 5, &vresponse, &resplen);
+       response = vresponse;
+       if (resplen < 5 || response[4] != SSH2_AGENT_IDENTITIES_ANSWER)
+           return NULL;
+
+       ret = snewn(resplen-5, unsigned char);
+       memcpy(ret, response+5, resplen-5);
+       sfree(response);
+    } else {
+       ret = make_keylist2(NULL);
+    }
+    return ret;
+}
+
+/*
  * This is the main agent function that answers messages.
  */
 static void answer_msg(void *msg)
@@ -529,42 +791,18 @@ static void answer_msg(void *msg)
         * Reply with SSH1_AGENT_RSA_IDENTITIES_ANSWER.
         */
        {
-           struct RSAKey *key;
-           int len, nkeys;
-           int i;
-
-           /*
-            * Count up the number and length of keys we hold.
-            */
-           len = nkeys = 0;
-           for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
-               nkeys++;
-               len += 4;              /* length field */
-               len += ssh1_bignum_length(key->exponent);
-               len += ssh1_bignum_length(key->modulus);
-               len += 4 + strlen(key->comment);
-           }
+           int len;
+           void *keylist;
 
-           /*
-            * Packet header is the obvious five bytes, plus four
-            * bytes for the key count.
-            */
-           len += 5 + 4;
-           if (len > AGENT_MAX_MSGLEN)
-               goto failure;          /* aaargh! too much stuff! */
-           PUT_32BIT(ret, len - 4);
            ret[4] = SSH1_AGENT_RSA_IDENTITIES_ANSWER;
-           PUT_32BIT(ret + 5, nkeys);
-           p = ret + 5 + 4;
-           for (i = 0; NULL != (key = index234(rsakeys, i)); i++) {
-               PUT_32BIT(p, bignum_bitcount(key->modulus));
-               p += 4;
-               p += ssh1_write_bignum(p, key->exponent);
-               p += ssh1_write_bignum(p, key->modulus);
-               PUT_32BIT(p, strlen(key->comment));
-               memcpy(p + 4, key->comment, strlen(key->comment));
-               p += 4 + strlen(key->comment);
+           keylist = make_keylist1(&len);
+           if (len + 5 > AGENT_MAX_MSGLEN) {
+               sfree(keylist);
+               goto failure;
            }
+           PUT_32BIT(ret, len + 1);
+           memcpy(ret + 5, keylist, len);
+           sfree(keylist);
        }
        break;
       case SSH2_AGENTC_REQUEST_IDENTITIES:
@@ -572,47 +810,18 @@ static void answer_msg(void *msg)
         * Reply with SSH2_AGENT_IDENTITIES_ANSWER.
         */
        {
-           struct ssh2_userkey *key;
-           int len, nkeys;
-           unsigned char *blob;
-           int bloblen;
-           int i;
-
-           /*
-            * Count up the number and length of keys we hold.
-            */
-           len = nkeys = 0;
-           for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
-               nkeys++;
-               len += 4;              /* length field */
-               blob = key->alg->public_blob(key->data, &bloblen);
-               len += bloblen;
-               sfree(blob);
-               len += 4 + strlen(key->comment);
-           }
+           int len;
+           void *keylist;
 
-           /*
-            * Packet header is the obvious five bytes, plus four
-            * bytes for the key count.
-            */
-           len += 5 + 4;
-           if (len > AGENT_MAX_MSGLEN)
-               goto failure;          /* aaargh! too much stuff! */
-           PUT_32BIT(ret, len - 4);
            ret[4] = SSH2_AGENT_IDENTITIES_ANSWER;
-           PUT_32BIT(ret + 5, nkeys);
-           p = ret + 5 + 4;
-           for (i = 0; NULL != (key = index234(ssh2keys, i)); i++) {
-               blob = key->alg->public_blob(key->data, &bloblen);
-               PUT_32BIT(p, bloblen);
-               p += 4;
-               memcpy(p, blob, bloblen);
-               p += bloblen;
-               sfree(blob);
-               PUT_32BIT(p, strlen(key->comment));
-               memcpy(p + 4, key->comment, strlen(key->comment));
-               p += 4 + strlen(key->comment);
+           keylist = make_keylist2(&len);
+           if (len + 5 > AGENT_MAX_MSGLEN) {
+               sfree(keylist);
+               goto failure;
            }
+           PUT_32BIT(ret, len + 1);
+           memcpy(ret + 5, keylist, len);
+           sfree(keylist);
        }
        break;
       case SSH1_AGENTC_RSA_CHALLENGE:
@@ -704,7 +913,7 @@ static void answer_msg(void *msg)
            struct RSAKey *key;
            char *comment;
             int commentlen;
-           key = smalloc(sizeof(struct RSAKey));
+           key = snew(struct RSAKey);
            memset(key, 0, sizeof(struct RSAKey));
            p += makekey(p, key, NULL, 1);
            p += makeprivate(p, key);
@@ -712,7 +921,7 @@ static void answer_msg(void *msg)
            p += ssh1_read_bignum(p, &key->p);  /* p */
            p += ssh1_read_bignum(p, &key->q);  /* q */
             commentlen = GET_32BIT(p);
-           comment = smalloc(commentlen+1);
+           comment = snewn(commentlen+1, char);
            if (comment) {
                memcpy(comment, p + 4, commentlen);
                 comment[commentlen] = '\0';
@@ -740,7 +949,7 @@ static void answer_msg(void *msg)
            int alglen, commlen;
            int bloblen;
 
-           key = smalloc(sizeof(struct ssh2_userkey));
+           key = snew(struct ssh2_userkey);
 
            alglen = GET_32BIT(p);
            p += 4;
@@ -768,7 +977,7 @@ static void answer_msg(void *msg)
            commlen = GET_32BIT(p);
            p += 4;
 
-           comment = smalloc(commlen + 1);
+           comment = snewn(commlen + 1, char);
            if (comment) {
                memcpy(comment, p, commlen);
                comment[commlen] = '\0';
@@ -1013,7 +1222,7 @@ static void prompt_add_keyfile(void)
 {
     OPENFILENAME of;
     char filename[FILENAME_MAX];
-    char *filelist = smalloc(8192);
+    char *filelist = snewn(8192, char);
     char *filewalker;
     int n, dirlen;
        
@@ -1023,8 +1232,9 @@ static void prompt_add_keyfile(void)
 #else
     of.lStructSize = sizeof(of);
 #endif
-    of.hwndOwner = hwnd;
-    of.lpstrFilter = "All Files\0*\0\0\0";
+    of.hwndOwner = main_hwnd;
+    of.lpstrFilter = "PuTTY Private Key Files (*.ppk)\0*.ppk\0"
+       "All Files (*.*)\0*\0\0\0";
     of.lpstrCustomFilter = NULL;
     of.nFilterIndex = 1;
     of.lpstrFile = filelist;
@@ -1037,7 +1247,7 @@ static void prompt_add_keyfile(void)
     if (GetOpenFileName(&of)) {
        if(strlen(filelist) > of.nFileOffset)
            /* Only one filename returned? */
-           add_keyfile(filelist);
+           add_keyfile(filename_from_str(filelist));
        else {
            /* we are returned a bunch of strings, end to
             * end. first string is the directory, the
@@ -1065,7 +1275,7 @@ static void prompt_add_keyfile(void)
                memcpy(filename + dirlen, filewalker, n);
                filewalker += n;
 
-               add_keyfile(filename);
+               add_keyfile(filename_from_str(filename));
            }
        }
 
@@ -1101,6 +1311,16 @@ static int CALLBACK KeyListProc(HWND hwnd, UINT msg,
                           rd.right - rd.left, rd.bottom - rd.top, TRUE);
        }
 
+        if (help_path)
+            SetWindowLong(hwnd, GWL_EXSTYLE,
+                          GetWindowLong(hwnd, GWL_EXSTYLE) | WS_EX_CONTEXTHELP);
+        else {
+            HWND item = GetDlgItem(hwnd, 103);   /* the Help button */
+            if (item)
+                DestroyWindow(item);
+        }
+        requested_help = FALSE;
+
        keylist = hwnd;
        {
            static int tabs[] = { 35, 60, 210 };
@@ -1149,7 +1369,7 @@ static int CALLBACK KeyListProc(HWND hwnd, UINT msg,
                }
 
                /* get item indices in an array */
-               selectedArray = smalloc(numSelected * sizeof(int));
+               selectedArray = snewn(numSelected, int);
                SendDlgItemMessage(hwnd, 100, LB_GETSELITEMS,
                                numSelected, (WPARAM)selectedArray);
                
@@ -1189,8 +1409,35 @@ static int CALLBACK KeyListProc(HWND hwnd, UINT msg,
                keylist_update();
            }
            return 0;
+         case 103:                    /* help */
+            if (HIWORD(wParam) == BN_CLICKED ||
+                HIWORD(wParam) == BN_DOUBLECLICKED) {
+                if (help_path) {
+                    WinHelp(main_hwnd, help_path, HELP_COMMAND,
+                            (DWORD)"JI(`',`pageant.general')");
+                    requested_help = TRUE;
+                }
+            }
+           return 0;
        }
        return 0;
+      case WM_HELP:
+        if (help_path) {
+            int id = ((LPHELPINFO)lParam)->iCtrlId;
+            char *cmd = NULL;
+            switch (id) {
+              case 100: cmd = "JI(`',`pageant.keylist')"; break;
+              case 101: cmd = "JI(`',`pageant.addkey')"; break;
+              case 102: cmd = "JI(`',`pageant.remkey')"; break;
+            }
+            if (cmd) {
+                WinHelp(main_hwnd, help_path, HELP_COMMAND, (DWORD)cmd);
+                requested_help = TRUE;
+            } else {
+                MessageBeep(0);
+            }
+        }
+        break;
       case WM_CLOSE:
        keylist = NULL;
        DestroyWindow(hwnd);
@@ -1226,6 +1473,59 @@ static BOOL AddTrayIcon(HWND hwnd)
     return res;
 }
 
+/* Update the saved-sessions menu. */
+static void update_sessions(void)
+{
+    int num_entries;
+    HKEY hkey;
+    TCHAR buf[MAX_PATH + 1];
+    MENUITEMINFO mii;
+
+    int index_key, index_menu;
+
+    if (!putty_path)
+       return;
+
+    if(ERROR_SUCCESS != RegOpenKey(HKEY_CURRENT_USER, PUTTY_REGKEY, &hkey))
+       return;
+
+    for(num_entries = GetMenuItemCount(session_menu);
+       num_entries > initial_menuitems_count;
+       num_entries--)
+       RemoveMenu(session_menu, 0, MF_BYPOSITION);
+
+    index_key = 0;
+    index_menu = 0;
+
+    while(ERROR_SUCCESS == RegEnumKey(hkey, index_key, buf, MAX_PATH)) {
+       TCHAR session_name[MAX_PATH + 1];
+       unmungestr(buf, session_name, MAX_PATH);
+       if(strcmp(buf, PUTTY_DEFAULT) != 0) {
+           memset(&mii, 0, sizeof(mii));
+           mii.cbSize = sizeof(mii);
+           mii.fMask = MIIM_TYPE | MIIM_STATE | MIIM_ID;
+           mii.fType = MFT_STRING;
+           mii.fState = MFS_ENABLED;
+           mii.wID = (index_menu * 16) + IDM_SESSIONS_BASE;
+           mii.dwTypeData = session_name;
+           InsertMenuItem(session_menu, index_menu, TRUE, &mii);
+           index_menu++;
+       }
+       index_key++;
+    }
+
+    RegCloseKey(hkey);
+
+    if(index_menu == 0) {
+       mii.cbSize = sizeof(mii);
+       mii.fMask = MIIM_TYPE | MIIM_STATE;
+       mii.fType = MFT_STRING;
+       mii.fState = MFS_GRAYED;
+       mii.dwTypeData = _T("(No sessions)");
+       InsertMenuItem(session_menu, index_menu, TRUE, &mii);
+    }
+}
+
 static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
                                WPARAM wParam, LPARAM lParam)
 {
@@ -1260,6 +1560,7 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
       case WM_SYSTRAY2:
        if (!menuinprogress) {
            menuinprogress = 1;
+           update_sessions();
            SetForegroundWindow(hwnd);
            ret = TrackPopupMenu(systray_menu,
                                 TPM_RIGHTALIGN | TPM_BOTTOMALIGN |
@@ -1271,6 +1572,13 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
       case WM_COMMAND:
       case WM_SYSCOMMAND:
        switch (wParam & ~0xF) {       /* low 4 bits reserved to Windows */
+         case IDM_PUTTY:
+           if((int)ShellExecute(hwnd, NULL, putty_path, _T(""), _T(""),
+                                SW_SHOW) <= 32) {
+               MessageBox(NULL, "Unable to execute PuTTY!",
+                          "Error", MB_OK | MB_ICONERROR);
+           }
+           break;
          case IDM_CLOSE:
            if (passphrase_box)
                SendMessage(passphrase_box, WM_CLOSE, 0, 0);
@@ -1281,14 +1589,17 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
                keylist = CreateDialog(instance, MAKEINTRESOURCE(211),
                                       NULL, KeyListProc);
                ShowWindow(keylist, SW_SHOWNORMAL);
-               /* 
-                * Sometimes the window comes up minimised / hidden
-                * for no obvious reason. Prevent this.
-                */
-               SetForegroundWindow(keylist);
-               SetWindowPos(keylist, HWND_TOP, 0, 0, 0, 0,
-                            SWP_NOMOVE | SWP_NOSIZE | SWP_SHOWWINDOW);
            }
+           /* 
+            * Sometimes the window comes up minimised / hidden for
+            * no obvious reason. Prevent this. This also brings it
+            * to the front if it's already present (the user
+            * selected View Keys because they wanted to _see_ the
+            * thing).
+            */
+           SetForegroundWindow(keylist);
+           SetWindowPos(keylist, HWND_TOP, 0, 0, 0, 0,
+                        SWP_NOMOVE | SWP_NOSIZE | SWP_SHOWWINDOW);
            break;
          case IDM_ADDKEY:
            if (passphrase_box) {
@@ -1312,9 +1623,42 @@ static LRESULT CALLBACK WndProc(HWND hwnd, UINT message,
                             SWP_NOMOVE | SWP_NOSIZE | SWP_SHOWWINDOW);
            }
            break;
+         case IDM_HELP:
+            if (help_path) {
+                WinHelp(main_hwnd, help_path, HELP_COMMAND,
+                        (DWORD)"JI(`',`pageant.general')");
+                requested_help = TRUE;
+            }
+           break;
+         default:
+           {
+               if(wParam >= IDM_SESSIONS_BASE && wParam <= IDM_SESSIONS_MAX) {
+                   MENUITEMINFO mii;
+                   TCHAR buf[MAX_PATH + 1];
+                   TCHAR param[MAX_PATH + 1];
+                   memset(&mii, 0, sizeof(mii));
+                   mii.cbSize = sizeof(mii);
+                   mii.fMask = MIIM_TYPE;
+                   mii.cch = MAX_PATH;
+                   mii.dwTypeData = buf;
+                   GetMenuItemInfo(session_menu, wParam, FALSE, &mii);
+                   strcpy(param, "@");
+                   strcat(param, mii.dwTypeData);
+                   if((int)ShellExecute(hwnd, NULL, putty_path, param,
+                                        _T(""), SW_SHOW) <= 32) {
+                       MessageBox(NULL, "Unable to execute PuTTY!", "Error",
+                                  MB_OK | MB_ICONERROR);
+                   }
+               }
+           }
+           break;
        }
        break;
       case WM_DESTROY:
+        if (requested_help) {
+            WinHelp(main_hwnd, help_path, HELP_QUIT, 0);
+            requested_help = FALSE;
+        }
        PostQuitMessage(0);
        return 0;
       case WM_COPYDATA:
@@ -1422,13 +1766,16 @@ void spawn_cmd(char *cmdline, char * args, int show)
 {
     if (ShellExecute(NULL, _T("open"), cmdline,
                     args, NULL, show) <= (HINSTANCE) 32) {
-       TCHAR sMsg[140];
-       sprintf(sMsg, _T("Failed to run \"%.100s\", Error: %d"), cmdline,
-               (int)GetLastError());
-       MessageBox(NULL, sMsg, APPNAME, MB_OK | MB_ICONEXCLAMATION);
+       char *msg;
+       msg = dupprintf("Failed to run \"%.100s\", Error: %d", cmdline,
+                       (int)GetLastError());
+       MessageBox(NULL, msg, APPNAME, MB_OK | MB_ICONEXCLAMATION);
+       sfree(msg);
     }
 }
 
+void cleanup_exit(int code) { exit(code); }
+
 int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
 {
     WNDCLASS wndclass;
@@ -1437,6 +1784,8 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
     HMODULE advapi;
     char *command = NULL;
     int added_keys = 0;
+    int argc, i;
+    char **argv, **argstart;
 
     /*
      * Determine whether we're an NT system (should have security
@@ -1477,6 +1826,47 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
     instance = inst;
 
     /*
+     * See if we can find our Help file.
+     */
+    {
+        char b[2048], *p, *q, *r;
+        FILE *fp;
+        GetModuleFileName(NULL, b, sizeof(b) - 1);
+        r = b;
+        p = strrchr(b, '\\');
+        if (p && p >= r) r = p+1;
+        q = strrchr(b, ':');
+        if (q && q >= r) r = q+1;
+        strcpy(r, "putty.hlp");
+        if ( (fp = fopen(b, "r")) != NULL) {
+            help_path = dupstr(b);
+            fclose(fp);
+        } else
+            help_path = NULL;
+    }
+
+    /*
+     * Look for the PuTTY binary (we will enable the saved session
+     * submenu if we find it).
+     */
+    {
+        char b[2048], *p, *q, *r;
+        FILE *fp;
+        GetModuleFileName(NULL, b, sizeof(b) - 1);
+        r = b;
+        p = strrchr(b, '\\');
+        if (p && p >= r) r = p+1;
+        q = strrchr(b, ':');
+        if (q && q >= r) r = q+1;
+        strcpy(r, "putty.exe");
+        if ( (fp = fopen(b, "r")) != NULL) {
+            putty_path = dupstr(b);
+            fclose(fp);
+        } else
+            putty_path = NULL;
+    }
+
+    /*
      * Find out if Pageant is already running.
      */
     already_running = FALSE;
@@ -1499,25 +1889,37 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
            RegisterClass(&wndclass);
        }
 
-       hwnd = keylist = NULL;
+       main_hwnd = keylist = NULL;
 
-       hwnd = CreateWindow(APPNAME, APPNAME,
-                           WS_OVERLAPPEDWINDOW | WS_VSCROLL,
-                           CW_USEDEFAULT, CW_USEDEFAULT,
-                           100, 100, NULL, NULL, inst, NULL);
+       main_hwnd = CreateWindow(APPNAME, APPNAME,
+                                 WS_OVERLAPPEDWINDOW | WS_VSCROLL,
+                                 CW_USEDEFAULT, CW_USEDEFAULT,
+                                 100, 100, NULL, NULL, inst, NULL);
 
        /* Set up a system tray icon */
-       AddTrayIcon(hwnd);
+       AddTrayIcon(main_hwnd);
 
+        /* Accelerators used: nsvkxa */
         systray_menu = CreatePopupMenu();
-        /* accelerators used: vkxa */
+       if (putty_path) {
+           session_menu = CreateMenu();
+           AppendMenu(systray_menu, MF_ENABLED, IDM_PUTTY, "&New Session");
+           AppendMenu(systray_menu, MF_POPUP | MF_ENABLED,
+                      (UINT) session_menu, "&Saved Sessions");
+           AppendMenu(systray_menu, MF_SEPARATOR, 0, 0);
+       }
         AppendMenu(systray_menu, MF_ENABLED, IDM_VIEWKEYS,
                "&View Keys");
         AppendMenu(systray_menu, MF_ENABLED, IDM_ADDKEY, "Add &Key");
+       AppendMenu(systray_menu, MF_SEPARATOR, 0, 0);
+        if (help_path)
+            AppendMenu(systray_menu, MF_ENABLED, IDM_HELP, "&Help");
         AppendMenu(systray_menu, MF_ENABLED, IDM_ABOUT, "&About");
+       AppendMenu(systray_menu, MF_SEPARATOR, 0, 0);
         AppendMenu(systray_menu, MF_ENABLED, IDM_CLOSE, "E&xit");
+       initial_menuitems_count = GetMenuItemCount(session_menu);
 
-       ShowWindow(hwnd, SW_HIDE);
+       ShowWindow(main_hwnd, SW_HIDE);
 
        /*
         * Initialise storage for RSA keys.
@@ -1528,52 +1930,34 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
     }
 
     /*
+     * Initialise the random number generator.
+     */
+    random_init();
+
+    /*
      * Initialise storage for short-term passphrase cache.
      */
     passphrases = newtree234(NULL);
 
     /*
      * Process the command line and add keys as listed on it.
-     * If we already determined that we need to spawn a program from above we
-     * need to ignore the first two arguments. [DBW]
      */
-    {
-       char *p;
-       int inquotes = 0;
-       p = cmdline;
-       while (*p) {
-           while (*p && isspace(*p))
-               p++;
-               if (*p && !isspace(*p)) {
-               char *q = p, *pp = p;
-               while (*p && (inquotes || !isspace(*p))) {
-                       if (*p == '"') {
-                       inquotes = !inquotes;
-                       p++;
-                       continue;
-                   }
-                   *pp++ = *p++;
-               }
-               if (*pp) {
-                   if (*p)
-                       p++;
-                   *pp++ = '\0';
-               }
-               if (!strcmp(q, "-c")) {
-                   /*
-                    * If we see `-c', then the rest of the
-                    * command line should be treated as a
-                    * command to be spawned.
-                    */
-                   while (*p && isspace(*p))
-                       p++;
-                   command = p;
-                   break;
-               } else {
-                   add_keyfile(q);
-                   added_keys = TRUE;
-               }
-           }
+    split_into_argv(cmdline, &argc, &argv, &argstart);
+    for (i = 0; i < argc; i++) {
+       if (!strcmp(argv[i], "-c")) {
+           /*
+            * If we see `-c', then the rest of the
+            * command line should be treated as a
+            * command to be spawned.
+            */
+           if (i < argc-1)
+               command = argstart[i+1];
+           else
+               command = "";
+           break;
+       } else {
+           add_keyfile(filename_from_str(argv[i]));
+           added_keys = TRUE;
        }
     }
 
@@ -1615,8 +1999,11 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
      * Main message loop.
      */
     while (GetMessage(&msg, NULL, 0, 0) == 1) {
-       TranslateMessage(&msg);
-       DispatchMessage(&msg);
+       if (!(IsWindow(keylist) && IsDialogMessage(keylist, &msg)) &&
+           !(IsWindow(aboutbox) && IsDialogMessage(aboutbox, &msg))) {
+           TranslateMessage(&msg);
+           DispatchMessage(&msg);
+       }
     }
 
     /* Clean up the system tray icon */
@@ -1624,7 +2011,7 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
        NOTIFYICONDATA tnid;
 
        tnid.cbSize = sizeof(NOTIFYICONDATA);
-       tnid.hWnd = hwnd;
+       tnid.hWnd = main_hwnd;
        tnid.uID = 1;
 
        Shell_NotifyIcon(NIM_DELETE, &tnid);
@@ -1634,5 +2021,5 @@ int WINAPI WinMain(HINSTANCE inst, HINSTANCE prev, LPSTR cmdline, int show)
 
     if (advapi)
        FreeLibrary(advapi);
-    exit(msg.wParam);
+    return msg.wParam;
 }