/*
 *  Off-the-Record Messaging Toolkit
 *  Copyright (C) 2004-2005  Nikita Borisov and Ian Goldberg
 *                           <otr@cypherpunks.ca>
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of version 2 of the GNU General Public License 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

/* system headers */
#include <stdio.h>
#include <stdlib.h>

/* libotr headers */
#include "b64.h"

/* toolkit headers */
#include "sha1hmac.h"
#include "parse.h"

/* Dump an unsigned int to a FILE * */
void dump_int(FILE *stream, const char *title, unsigned int val)
{
    fprintf(stream, "%s: %u\n", title, val);
}

/* Dump an mpi to a FILE * */
void dump_mpi(FILE *stream, const char *title, gcry_mpi_t val)
{
    size_t plen;
    char *d;
    
    gcry_mpi_print(GCRYMPI_FMT_USG, NULL, 0, &plen, val);
    d = malloc(plen);
    gcry_mpi_print(GCRYMPI_FMT_USG, d, plen, NULL, val);
    dump_data(stream, title, d, plen);
    free(d);
}

/* Dump data to a FILE * */
void dump_data(FILE *stream, const char *title, const unsigned char *data,
	size_t datalen)
{
    size_t i;
    fprintf(stream, "%s: ", title);
    for(i=0;i<datalen;++i) {
	fprintf(stream, "%02x", data[i]);
    }
    fprintf(stream, "\n");
}

/* base64 decode the message, and put the resulting size into *lenp */
static unsigned char *decode(const char *msg, size_t *lenp)
{
    const char *header, *footer;
    unsigned char *raw;
	
    /* Find the header */
    header = strstr(msg, "?OTR:");
    if (!header) return NULL;
    /* Skip the header */
    header += 5;

    /* Find the trailing '.' */
    footer = strchr(header, '.');
    if (!footer) footer = header + strlen(header);

    raw = malloc((footer-header) / 4 * 3);
    if (raw == NULL && (footer-header >= 4)) return NULL;
    *lenp = otrl_base64_decode(raw, header, footer-header);

    return raw;
}

#define require_len(l) do { if (lenp < (l)) goto inv; } while(0)
#define read_int(x) do { \
	require_len(4); \
	(x) = (bufp[0] << 24) | (bufp[1] << 16) | (bufp[2] << 8 ) | bufp[3]; \
	bufp += 4; lenp -= 4; \
    } while(0)
#define read_mpi(x) do { \
	size_t mpilen; \
	read_int(mpilen); \
	require_len(mpilen); \
	gcry_mpi_scan(&(x), GCRYMPI_FMT_USG, bufp, mpilen, NULL); \
	bufp += mpilen; lenp -= mpilen; \
    } while(0)
#define read_raw(b, l) do { \
	require_len(l); \
	memmove((b), bufp, (l)); \
	bufp += (l); lenp -= (l); \
    } while(0)
#define write_int(x) do { \
	bufp[0] = ((x) >> 24) & 0xff; \
	bufp[1] = ((x) >> 16) & 0xff; \
	bufp[2] = ((x) >> 8) & 0xff; \
	bufp[3] = (x) & 0xff; \
	bufp += 4; lenp -= 4; \
    } while(0)
#define write_mpi(x,l) do { \
	write_int(l); \
	gcry_mpi_print(GCRYMPI_FMT_USG, bufp, lenp, NULL, (x)); \
	bufp += (l); lenp -= (l); \
    } while(0)
#define write_raw(x,l) do { \
	memmove(bufp, (x), (l)); \
	bufp += (l); lenp -= (l); \
    } while(0)

/* Parse a Key Exchange Message into a newly-allocated KeyExchMsg structure */
KeyExchMsg parse_keyexch(const char *msg)
{
    KeyExchMsg kem = NULL;
    size_t lenp;
    unsigned char *raw = decode(msg, &lenp);
    unsigned char *bufp = raw;
    if (!raw) goto inv;

    kem = calloc(1, sizeof(struct s_KeyExchMsg));
    if (!kem) {
	free(raw);
	goto inv;
    }

    kem->raw = raw;
    kem->sigstart = bufp;

    require_len(3);
    if (memcmp(bufp, "\x00\x01\x0a", 3)) goto inv;
    bufp += 3; lenp -= 3;

    require_len(1);
    kem->reply = *bufp;
    bufp += 1; lenp -= 1;

    read_mpi(kem->p);
    read_mpi(kem->q);
    read_mpi(kem->g);
    read_mpi(kem->e);

    read_int(kem->keyid);

    read_mpi(kem->y);

    kem->sigend = bufp;

    require_len(40);
    gcry_mpi_scan(&kem->r, GCRYMPI_FMT_USG, bufp, 20, NULL);
    gcry_mpi_scan(&kem->s, GCRYMPI_FMT_USG, bufp+20, 20, NULL);
    bufp += 40; lenp -= 40;

    if (lenp != 0) goto inv;

    return kem;
inv:
    free_keyexch(kem);
    return NULL;
}

/* Deallocate a KeyExchMsg and all of the data it points to */
void free_keyexch(KeyExchMsg keyexch)
{
    if (!keyexch) return;
    free(keyexch->raw);
    gcry_mpi_release(keyexch->p);
    gcry_mpi_release(keyexch->q);
    gcry_mpi_release(keyexch->g);
    gcry_mpi_release(keyexch->e);
    gcry_mpi_release(keyexch->y);
    gcry_mpi_release(keyexch->r);
    gcry_mpi_release(keyexch->s);
    free(keyexch);
}

/* Parse a Data Message into a newly-allocated DataMsg structure */
DataMsg parse_datamsg(const char *msg)
{
    DataMsg datam = NULL;
    size_t lenp;
    unsigned char *raw = decode(msg, &lenp);
    unsigned char *bufp = raw;
    if (!raw) goto inv;

    datam = calloc(1, sizeof(struct s_DataMsg));
    if (!datam) {
	free(raw);
	goto inv;
    }

    datam->raw = raw;
    datam->rawlen = lenp;
    datam->macstart = bufp;

    require_len(3);
    if (memcmp(bufp, "\x00\x01\x03", 3)) goto inv;
    bufp += 3; lenp -= 3;

    read_int(datam->sender_keyid);
    read_int(datam->rcpt_keyid);
    read_mpi(datam->y);
    read_raw(datam->ctr, 8);
    read_int(datam->encmsglen);
    datam->encmsg = malloc(datam->encmsglen);
    if (!datam->encmsg && datam->encmsglen > 0) goto inv;
    read_raw(datam->encmsg, datam->encmsglen);
    datam->macend = bufp;
    read_raw(datam->mac, 20);
    read_int(datam->mackeyslen);
    datam->mackeys = malloc(datam->mackeyslen);
    if (!datam->mackeys && datam->mackeyslen > 0) goto inv;
    read_raw(datam->mackeys, datam->mackeyslen);

    if (lenp != 0) goto inv;

    return datam;
inv:
    free_datamsg(datam);
    return NULL;
}

/* Recalculate the MAC on the message, base64-encode the resulting MAC'd
 * message, and put on the appropriate header and footer.  Return a
 * newly-allocated pointer to the result, which the caller will have to
 * free(). */
char *remac_datamsg(DataMsg datamsg, unsigned char mackey[20])
{
    size_t rawlen, lenp;
    size_t ylen;
    size_t base64len;
    char *outmsg;
    unsigned char *raw, *bufp;
    
    /* Calculate the size of the message that will result */
    gcry_mpi_print(GCRYMPI_FMT_USG, NULL, 0, &ylen, datamsg->y);
    rawlen = 3 + 4 + 4 + 4 + ylen + 8 + 4 + datamsg->encmsglen + 20 +
	4 + datamsg->mackeyslen;

    /* Construct the new raw message (note that some of the pieces may
     * have been altered, so we construct it from scratch). */
    raw = malloc(rawlen);
    if (!raw) {
	fprintf(stderr, "Out of memory!\n");
	exit(1);
    }
    bufp = raw;
    lenp = rawlen;
    datamsg->macstart = raw;
    datamsg->macend = NULL;
    free(datamsg->raw);
    datamsg->raw = raw;
    datamsg->rawlen = rawlen;

    memmove(bufp, "\x00\x01\x03", 3);
    bufp += 3; lenp -= 3;
    write_int(datamsg->sender_keyid);
    write_int(datamsg->rcpt_keyid);
    write_mpi(datamsg->y, ylen);
    write_raw(datamsg->ctr, 8);
    write_int(datamsg->encmsglen);
    write_raw(datamsg->encmsg, datamsg->encmsglen);
    datamsg->macend = bufp;

    /* Recalculate the MAC */
    sha1hmac(datamsg->mac, mackey, datamsg->macstart,
	    datamsg->macend - datamsg->macstart);

    write_raw(datamsg->mac, 20);
    write_int(datamsg->mackeyslen);
    write_raw(datamsg->mackeys, datamsg->mackeyslen);
    
    if (lenp != 0) {
	fprintf(stderr, "Error creating OTR Data Message.\n");
	exit(1);
    }

    base64len = 5 + ((datamsg->rawlen + 2) / 3) * 4 + 1 + 1;
    outmsg = malloc(base64len);
    if (!outmsg) return NULL;

    memmove(outmsg, "?OTR:", 5);
    otrl_base64_encode(outmsg + 5, datamsg->raw, datamsg->rawlen);
    strcpy(outmsg + base64len - 2, ".");
    return outmsg;
}

/* Assemble a new Data Message from its pieces.  Return a
 * newly-allocated string containing the base64 representation. */
char *assemble_datamsg(unsigned char mackey[20], unsigned int sender_keyid,
	unsigned int rcpt_keyid, gcry_mpi_t y, unsigned char ctr[8],
	unsigned char *encmsg, size_t encmsglen, unsigned char *mackeys,
	size_t mackeyslen)
{
    DataMsg datam = calloc(1, sizeof(struct s_DataMsg));
    char *newmsg = NULL;
    if (!datam) goto inv;
    datam->sender_keyid = sender_keyid;
    datam->rcpt_keyid = rcpt_keyid;
    datam->y = gcry_mpi_copy(y);
    memmove(datam->ctr, ctr, 8);
    datam->encmsg = malloc(encmsglen);
    if (!datam->encmsg && encmsglen > 0) goto inv;
    memmove(datam->encmsg, encmsg, encmsglen);
    datam->encmsglen = encmsglen;
    datam->mackeys = malloc(mackeyslen);
    if (!datam->mackeys && mackeyslen > 0) goto inv;
    memmove(datam->mackeys, mackeys, mackeyslen);
    datam->mackeyslen = mackeyslen;

    /* Recalculate the MAC and base64-encode the result */
    newmsg = remac_datamsg(datam, mackey);
    free_datamsg(datam);
    return newmsg;
inv:
    free_datamsg(datam);
    return NULL;
}

/* Deallocate a DataMsg and all of the data it points to */
void free_datamsg(DataMsg datamsg)
{
    if (!datamsg) return;
    free(datamsg->raw);
    gcry_mpi_release(datamsg->y);
    free(datamsg->encmsg);
    free(datamsg->mackeys);
    free(datamsg);
}

static int ctoh(char c)
{
    if (c >= '0' && c <= '9') return (c-'0');
    if (c >= 'a' && c <= 'f') return (c-'a'+10);
    if (c >= 'A' && c <= 'F') return (c-'A'+10);
    return -1;
}

/* Convert a string of hex chars to a buffer of unsigned chars. */
void argv_to_buf(unsigned char **bufp, size_t *lenp, char *arg)
{
    unsigned char *buf;
    size_t len, i;

    *bufp = NULL;
    *lenp = 0;

    len = strlen(arg);
    if (len % 2) {
	fprintf(stderr, "Argument ``%s'' must have even length.\n", arg);
	return;
    }
    buf = malloc(len/2);
    if (buf == NULL && len > 0) {
	fprintf(stderr, "Out of memory!\n");
	return;
    }

    for(i=0;i<len/2;++i) {
	int hi = ctoh(arg[2*i]);
	int lo = ctoh(arg[2*i+1]);
	if (hi < 0 || lo < 0) {
	    free(buf);
	    fprintf(stderr, "Illegal hex char in argument ``%s''.\n", arg);
	    return;
	}
	buf[i] = (hi << 4) + lo;
    }
    *bufp = buf;
    *lenp = len/2;
}