/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "signtool.h"
#include "zip.h" 
#include "prmem.h"
#include "blapi.h"
#include "sechash.h"	/* for HASH_GetHashObject() */

static int	create_pk7 (char *dir, char *keyName, int *keyType);
static int	jar_find_key_type (CERTCertificate *cert);
static int	manifesto (char *dirname, char *install_script, PRBool recurse);
static int	manifesto_fn(char *relpath, char *basedir, char *reldir,
			char *filename, void *arg);
static int	manifesto_xpi_fn(char *relpath, char *basedir, char *reldir,
			char *filename, void *arg);
static int	sign_all_arc_fn(char *relpath, char *basedir, char *reldir,
			char *filename, void *arg);
static int	add_meta (FILE *fp, char *name);
static int	SignFile (FILE *outFile, FILE *inFile, CERTCertificate *cert);
static int	generate_SF_file (char *manifile, char *who);
static int	calculate_MD5_range (FILE *fp, long r1, long r2, 
			JAR_Digest *dig);
static void	SignOut (void *arg, const char *buf, unsigned long len);

static char	*metafile = NULL;
static int	optimize = 0;
static FILE *mf;
static ZIPfile *zipfile = NULL;

/* 
 *  S i g n A r c h i v e
 *
 *  Sign an individual archive tree. A directory 
 *  called META-INF is created underneath this.
 *
 */
int
SignArchive(char *tree, char *keyName, char *zip_file, int javascript,
	char *meta_file, char *install_script, int _optimize, PRBool recurse)
{
    int	status;
    char	tempfn [FNSIZE], fullfn [FNSIZE];
    int	keyType = rsaKey;

    metafile = meta_file;
    optimize = _optimize;

    /* To create XPI compatible Archive manifesto() must be run before 
     * the zipfile is opened. This is so the signed files are not added
     * the archive before the crucial rsa/dsa file*/
    if (xpi_arc) {
	manifesto (tree, install_script, recurse);
    }

    if (zip_file) {
	zipfile = JzipOpen(zip_file, NULL /*no comment*/);
    }

    /*Sign and add files to the archive normally with manifesto()*/
    if (!xpi_arc) {
	manifesto (tree, install_script, recurse);
    }

    if (keyName) {
	status = create_pk7 (tree, keyName, &keyType);
	if (status < 0) {
	    PR_fprintf(errorFD, "the tree \"%s\" was NOT SUCCESSFULLY SIGNED\n",
	         tree);
	    errorCount++;
	    exit (ERRX);
	}
    }

    /* Add the rsa/dsa file as the first file in the archive. This is crucial
     * for a XPInstall compatible archive */
    if (xpi_arc) {
	if (verbosity >= 0) {
	    PR_fprintf(outputFD, "%s \n", XPI_TEXT);
	}

	/* rsa/dsa to zip */
	sprintf (tempfn, "META-INF/%s.%s", base, (keyType == dsaKey ?
	    "dsa" : "rsa"));
	sprintf (fullfn, "%s/%s", tree, tempfn);
	JzipAdd(fullfn, tempfn, zipfile, compression_level);

	/* Loop through all files & subdirectories, add to archive */
	foreach (tree, "", manifesto_xpi_fn, recurse, PR_FALSE /*include dirs */,
	     		(void * )NULL);
    }
    /* mf to zip */
    strcpy (tempfn, "META-INF/manifest.mf");
    sprintf (fullfn, "%s/%s", tree, tempfn);
    JzipAdd(fullfn, tempfn, zipfile, compression_level);

    /* sf to zip */
    sprintf (tempfn, "META-INF/%s.sf", base);
    sprintf (fullfn, "%s/%s", tree, tempfn);
    JzipAdd(fullfn, tempfn, zipfile, compression_level);

    /* Add the rsa/dsa file to the zip archive normally */
    if (!xpi_arc) {
	/* rsa/dsa to zip */
	sprintf (tempfn, "META-INF/%s.%s", base, (keyType == dsaKey ?
	    "dsa" : "rsa"));
	sprintf (fullfn, "%s/%s", tree, tempfn);
	JzipAdd(fullfn, tempfn, zipfile, compression_level);
    }

    JzipClose(zipfile);

    if (verbosity >= 0) {
	if (javascript) {
	    PR_fprintf(outputFD, "jarfile \"%s\" signed successfully\n",
	         				zip_file);
	} else {
	    PR_fprintf(outputFD, "tree \"%s\" signed successfully\n",
	         tree);
	}
    }

    return 0;
}


typedef struct {
    char	*keyName;
    int	javascript;
    char	*metafile;
    char	*install_script;
    int	optimize;
} SignArcInfo;

/* 
 *  S i g n A l l A r c
 *
 *  Javascript may generate multiple .arc directories, one
 *  for each jar archive needed. Sign them all.
 *
 */
int
SignAllArc(char *jartree, char *keyName, int javascript, char *metafile,
char *install_script, int optimize, PRBool recurse)
{
    SignArcInfo info;

    info.keyName = keyName;
    info.javascript = javascript;
    info.metafile = metafile;
    info.install_script = install_script;
    info.optimize = optimize;

    return foreach(jartree, "", sign_all_arc_fn, recurse,
        PR_TRUE /*include dirs*/, (void * )&info);
}


static int	
sign_all_arc_fn(char *relpath, char *basedir, char *reldir, char *filename,
 	void *arg)
{
    char	*zipfile = NULL;
    char	*arc = NULL, *archive = NULL;
    int	retval = 0;
    SignArcInfo * infop = (SignArcInfo * )arg;

    /* Make sure there is one and only one ".arc" in the relative path, 
     * and that it is at the end of the path (don't sign .arcs within .arcs) */
    if ( (PL_strcaserstr(relpath, ".arc") == relpath + strlen(relpath) -
        4) && 
        (PL_strcasestr(relpath, ".arc") == relpath + strlen(relpath) - 4) ) {

	if (!infop) {
	    PR_fprintf(errorFD, "%s: Internal failure\n", PROGRAM_NAME);
	    errorCount++;
	    retval = -1;
	    goto finish;
	}
	archive = PR_smprintf("%s/%s", basedir, relpath);

	zipfile = PL_strdup(archive);
	arc = PORT_Strrchr (zipfile, '.');

	if (arc == NULL) {
	    PR_fprintf(errorFD, "%s: Internal failure\n", PROGRAM_NAME);
	    errorCount++;
	    retval = -1;
	    goto finish;
	}

	PL_strcpy (arc, ".jar");

	if (verbosity >= 0) {
	    PR_fprintf(outputFD, "\nsigning: %s\n", zipfile);
	}
	retval = SignArchive(archive, infop->keyName, zipfile,
	    infop->javascript, infop->metafile, infop->install_script,
	     			infop->optimize, PR_TRUE /* recurse */);
    }
finish:
    if (archive) 
	PR_Free(archive);
    if (zipfile) 
	PR_Free(zipfile);

    return retval;
}


/*********************************************************************
 *
 * c r e a t e _ p k 7
 */
static int	
create_pk7 (char *dir, char *keyName, int *keyType)
{
    int	status = 0;
    char	*file_ext;

    CERTCertificate * cert;
    CERTCertDBHandle * db;

    FILE * in, *out;

    char	sf_file [FNSIZE];
    char	pk7_file [FNSIZE];

    /* open cert database */
    db = CERT_GetDefaultCertDB();

    if (db == NULL)
	return - 1;

    /* find cert */
    /*cert = CERT_FindCertByNicknameOrEmailAddr(db, keyName);*/
    cert = PK11_FindCertFromNickname(keyName, &pwdata);

    if (cert == NULL) {
	SECU_PrintError ( PROGRAM_NAME,
	    "Cannot find the cert \"%s\"", keyName);
	return -1;
    }


    /* determine the key type, which sets the extension for pkcs7 object */

    *keyType = jar_find_key_type (cert);
    file_ext = (*keyType == dsaKey) ? "dsa" : "rsa";

    sprintf (sf_file, "%s/META-INF/%s.sf", dir, base);
    sprintf (pk7_file, "%s/META-INF/%s.%s", dir, base, file_ext);

    if ((in = fopen (sf_file, "rb")) == NULL) {
	PR_fprintf(errorFD, "%s: Can't open %s for reading\n", PROGRAM_NAME,
	     sf_file);
	errorCount++;
	exit (ERRX);
    }

    if ((out = fopen (pk7_file, "wb")) == NULL) {
	PR_fprintf(errorFD, "%s: Can't open %s for writing\n", PROGRAM_NAME,
	     sf_file);
	errorCount++;
	exit (ERRX);
    }

    status = SignFile (out, in, cert);

    CERT_DestroyCertificate (cert);
    fclose (in);
    fclose (out);

    if (status) {
	PR_fprintf(errorFD, "%s: PROBLEM signing data (%s)\n",
	    PROGRAM_NAME, SECU_Strerror(PORT_GetError()));
	errorCount++;
	return - 1;
    }

    return 0;
}


/*
 *  j a r _ f i n d _ k e y _ t y p e
 * 
 *  Determine the key type for a given cert, which 
 * should be rsaKey or dsaKey. Any error return 0.
 *
 */
static int	
jar_find_key_type (CERTCertificate *cert)
{
    SECKEYPrivateKey * privk = NULL;
    KeyType keyType;

    /* determine its type */
    privk = PK11_FindKeyByAnyCert (cert, &pwdata);
    if (privk == NULL) {
	PR_fprintf(errorFD, "warning - can't find private key for this cert\n");
	warningCount++;
	return 0;
    }

    keyType = privk->keyType;
    SECKEY_DestroyPrivateKey (privk);
    return keyType;
}


/*
 *  m a n i f e s t o
 *
 *  Run once for every subdirectory in which a 
 *  manifest is to be created -- usually exactly once.
 *
 */
static int	
manifesto (char *dirname, char *install_script, PRBool recurse)
{
    char	metadir [FNSIZE], sfname [FNSIZE];

    /* Create the META-INF directory to hold signing info */

    if (PR_Access (dirname, PR_ACCESS_READ_OK)) {
	PR_fprintf(errorFD, "%s: unable to read your directory: %s\n",
	     PROGRAM_NAME, dirname);
	errorCount++;
	perror (dirname);
	exit (ERRX);
    }

    if (PR_Access (dirname, PR_ACCESS_WRITE_OK)) {
	PR_fprintf(errorFD, "%s: unable to write to your directory: %s\n",
	     			PROGRAM_NAME, dirname);
	errorCount++;
	perror(dirname);
	exit(ERRX);
    }

    sprintf (metadir, "%s/META-INF", dirname);

    strcpy (sfname, metadir);

    PR_MkDir (metadir, 0777);

    strcat (metadir, "/");
    strcat (metadir, MANIFEST);

    if ((mf = fopen (metadir, "wb")) == NULL) {
	perror (MANIFEST);
	PR_fprintf(errorFD, "%s: Probably, the directory you are trying to"
			    " sign has\n", PROGRAM_NAME);
	PR_fprintf(errorFD, "%s: permissions problems or may not exist.\n",
	     		PROGRAM_NAME);
	errorCount++;
	exit (ERRX);
    }

    if (verbosity >= 0) {
	PR_fprintf(outputFD, "Generating %s file..\n", metadir);
    }

    fprintf(mf, "Manifest-Version: 1.0\n");
    fprintf (mf, "Created-By: %s\n", CREATOR);
    fprintf (mf, "Comments: %s\n", BREAKAGE);

    if (scriptdir) {
	fprintf (mf, "Comments: --\n");
	fprintf (mf, "Comments: --\n");
	fprintf (mf, "Comments: -- This archive signs Javascripts which may not necessarily\n");
	fprintf (mf, "Comments: -- be included in the physical jar file.\n");
	fprintf (mf, "Comments: --\n");
	fprintf (mf, "Comments: --\n");
    }

    if (install_script)
	fprintf (mf, "Install-Script: %s\n", install_script);

    if (metafile)
	add_meta (mf, "+");

    /* Loop through all files & subdirectories */
    foreach (dirname, "", manifesto_fn, recurse, PR_FALSE /*include dirs */,
         		(void * )NULL);

    fclose (mf);

    strcat (sfname, "/");
    strcat (sfname, base);
    strcat (sfname, ".sf");

    if (verbosity >= 0) {
	PR_fprintf(outputFD, "Generating %s.sf file..\n", base);
    }
    generate_SF_file (metadir, sfname);

    return 0;
}


/*
 *  m a n i f e s t o _ x p i _ f n
 *
 *  Called by pointer from SignArchive(), once for
 *  each file within the directory. This function
 *  is only used for adding to XPI compatible archive
 *
 */
static int	manifesto_xpi_fn 
(char *relpath, char *basedir, char *reldir, char *filename, void *arg)
{
    char	fullname [FNSIZE];

    if (verbosity >= 0) {
	PR_fprintf(outputFD, "--> %s\n", relpath);
    }

    /* extension matching */
    if (extensionsGiven) {
	char	*ext = PL_strrchr(relpath, '.');
	if (!ext) 
	    return 0;
	if (!PL_HashTableLookup(extensions, ext)) 
	    return 0;
    }
    sprintf (fullname, "%s/%s", basedir, relpath);
    JzipAdd(fullname, relpath, zipfile, compression_level);

    return 0;
}


/*
 *  m a n i f e s t o _ f n
 *
 *  Called by pointer from manifesto(), once for
 *  each file within the directory.
 *
 */
static int	manifesto_fn 
(char *relpath, char *basedir, char *reldir, char *filename, void *arg)
{
    int	use_js;

    JAR_Digest dig;
    char	fullname [FNSIZE];

    if (verbosity >= 0) {
	PR_fprintf(outputFD, "--> %s\n", relpath);
    }

    /* extension matching */
    if (extensionsGiven) {
	char	*ext = PL_strrchr(relpath, '.');
	if (!ext) 
	    return 0;
	if (!PL_HashTableLookup(extensions, ext)) 
	    return 0;
    }

    sprintf (fullname, "%s/%s", basedir, relpath);

    fprintf (mf, "\n");

    use_js = 0;

    if (scriptdir && !PORT_Strcmp (scriptdir, reldir))
	use_js++;

    /* sign non-.js files inside .arc directories using the javascript magic */

    if ( (PL_strcaserstr(filename, ".js") != filename + strlen(filename) - 3)
         && (PL_strcaserstr(reldir, ".arc") == reldir + strlen(filename) - 4))
	use_js++;

    if (use_js) {
	fprintf (mf, "Name: %s\n", filename);
	fprintf (mf, "Magic: javascript\n");

	if (optimize == 0)
	    fprintf (mf, "javascript.id: %s\n", filename);

	if (metafile)
	    add_meta (mf, filename);
    } else {
	fprintf (mf, "Name: %s\n", relpath);
	if (metafile)
	    add_meta (mf, relpath);
    }

    JAR_digest_file (fullname, &dig);


    if (optimize == 0) {
	fprintf (mf, "Digest-Algorithms: MD5 SHA1\n");
	fprintf (mf, "MD5-Digest: %s\n", BTOA_DataToAscii (dig.md5,
	     MD5_LENGTH));
    }

    fprintf (mf, "SHA1-Digest: %s\n", BTOA_DataToAscii (dig.sha1, SHA1_LENGTH));

    if (!use_js) {
	JzipAdd(fullname, relpath, zipfile, compression_level);
    }

    return 0;
}


/*
 *  a d d _ m e t a
 *
 *  Parse the metainfo file, and add any details
 *  necessary to the manifest file. In most cases you
 *  should be using the -i option (ie, for SmartUpdate).
 *
 */
static int	add_meta (FILE *fp, char *name)
{
    FILE * met;
    char	buf [BUFSIZ];

    int	place;
    char	*pattern, *meta;

    int	num = 0;

    if ((met = fopen (metafile, "r")) != NULL) {
	while (fgets (buf, BUFSIZ, met)) {
	    char	*s;

	    for (s = buf; *s && *s != '\n' && *s != '\r'; s++)
		;
	    *s = 0;

	    if (*buf == 0)
		continue;

	    pattern = buf;

	    /* skip to whitespace */
	    for (s = buf; *s && *s != ' ' && *s != '\t'; s++)
		;

	    /* terminate pattern */
	    if (*s == ' ' || *s == '\t') 
		*s++ = 0;

	    /* eat through whitespace */
	    while (*s == ' ' || *s == '\t') 
		s++;

	    meta = s;

	    /* this will eventually be regexp matching */

	    place = 0;
	    if (!PORT_Strcmp (pattern, name))
		place = 1;

	    if (place) {
		num++;
		if (verbosity >= 0) {
		    PR_fprintf(outputFD, "[%s] %s\n", name, meta);
		}
		fprintf (fp, "%s\n", meta);
	    }
	}
	fclose (met);
    } else {
	PR_fprintf(errorFD, "%s: can't open metafile: %s\n", PROGRAM_NAME,
	     metafile);
	errorCount++;
	exit (ERRX);
    }

    return num;
}


/**********************************************************************
 *
 * S i g n F i l e
 */
static int	
SignFile (FILE *outFile, FILE *inFile, CERTCertificate *cert)
{
    int	nb;
    char	ibuf[4096], digestdata[32];
    const SECHashObject *hashObj;
    void	*hashcx;
    unsigned int	len;

    SECItem digest;
    SEC_PKCS7ContentInfo * cinfo;
    SECStatus rv;

    if (outFile == NULL || inFile == NULL || cert == NULL)
	return - 1;

    /* XXX probably want to extend interface to allow other hash algorithms */
    hashObj = HASH_GetHashObject(HASH_AlgSHA1);

    hashcx = (*hashObj->create)();
    if (hashcx == NULL)
	return - 1;

    (*hashObj->begin)(hashcx);

    for (; ; ) {
	if (feof(inFile)) 
	    break;
	nb = fread(ibuf, 1, sizeof(ibuf), inFile);
	if (nb == 0) {
	    if (ferror(inFile)) {
		PORT_SetError(SEC_ERROR_IO);
		(*hashObj->destroy)(hashcx, PR_TRUE);
		return - 1;
	    }
	    /* eof */
	    break;
	}
	(*hashObj->update)(hashcx, (unsigned char *) ibuf, nb);
    }

    (*hashObj->end)(hashcx, (unsigned char *) digestdata, &len, 32);
    (*hashObj->destroy)(hashcx, PR_TRUE);

    digest.data = (unsigned char *) digestdata;
    digest.len = len;

    cinfo = SEC_PKCS7CreateSignedData 
        (cert, certUsageObjectSigner, NULL, 
        SEC_OID_SHA1, &digest, NULL, NULL);

    if (cinfo == NULL)
	return - 1;

    rv = SEC_PKCS7IncludeCertChain (cinfo, NULL);
    if (rv != SECSuccess) {
	SEC_PKCS7DestroyContentInfo (cinfo);
	return - 1;
    }

    if (no_time == 0) {
	rv = SEC_PKCS7AddSigningTime (cinfo);
	if (rv != SECSuccess) {
	    /* don't check error */
	}
    }

    rv = SEC_PKCS7Encode(cinfo, SignOut, outFile, NULL, NULL, &pwdata);

    SEC_PKCS7DestroyContentInfo (cinfo);

    if (rv != SECSuccess)
	return - 1;

    return 0;
}


/*
 *  g e n e r a t e _ S F _ f i l e 
 *
 *  From the supplied manifest file, calculates
 *  digests on the various sections, creating a .SF
 *  file in the process.
 * 
 */
static int	generate_SF_file (char *manifile, char *who)
{
    FILE * sf;
    FILE * mf;
    long	r1, r2, r3;
    char	whofile [FNSIZE];
    char	*buf, *name = NULL;
    JAR_Digest dig;
    int	line = 0;

    strcpy (whofile, who);

    if ((mf = fopen (manifile, "rb")) == NULL) {
	perror (manifile);
	exit (ERRX);
    }

    if ((sf = fopen (whofile, "wb")) == NULL) {
	perror (who);
	exit (ERRX);
    }

    buf = (char *) PORT_ZAlloc (BUFSIZ);

    if (buf)
	name = (char *) PORT_ZAlloc (BUFSIZ);

    if (buf == NULL || name == NULL)
	out_of_memory();

    fprintf (sf, "Signature-Version: 1.0\n");
    fprintf (sf, "Created-By: %s\n", CREATOR);
    fprintf (sf, "Comments: %s\n", BREAKAGE);

    if (fgets (buf, BUFSIZ, mf) == NULL) {
	PR_fprintf(errorFD, "%s: empty manifest file!\n", PROGRAM_NAME);
	errorCount++;
	exit (ERRX);
    }

    if (strncmp (buf, "Manifest-Version:", 17)) {
	PR_fprintf(errorFD, "%s: not a manifest file!\n", PROGRAM_NAME);
	errorCount++;
	exit (ERRX);
    }

    fseek (mf, 0L, SEEK_SET);

    /* Process blocks of headers, and calculate their hashen */

    while (1) {
	/* Beginning range */
	r1 = ftell (mf);

	if (fgets (name, BUFSIZ, mf) == NULL)
	    break;

	line++;

	if (r1 != 0 && strncmp (name, "Name:", 5)) {
	    PR_fprintf(errorFD, 
	    "warning: unexpected input in manifest file \"%s\" at line %d:\n",
	         manifile, line);
	    PR_fprintf(errorFD, "%s\n", name);
	    warningCount++;
	}

	r2 = r1;
	while (fgets (buf, BUFSIZ, mf)) {
	    if (*buf == 0 || *buf == '\n' || *buf == '\r')
		break;

	    line++;

	    /* Ending range for hashing */
	    r2 = ftell (mf);
	}

	r3 = ftell (mf);

	if (r1) {
	    fprintf (sf, "\n");
	    fprintf (sf, "%s", name);
	}

	calculate_MD5_range (mf, r1, r2, &dig);

	if (optimize == 0) {
	    fprintf (sf, "Digest-Algorithms: MD5 SHA1\n");
	    fprintf (sf, "MD5-Digest: %s\n", 
	        BTOA_DataToAscii (dig.md5, MD5_LENGTH));
	}

	fprintf (sf, "SHA1-Digest: %s\n", 
	    BTOA_DataToAscii (dig.sha1, SHA1_LENGTH));

	/* restore normalcy after changing offset position */
	fseek (mf, r3, SEEK_SET);
    }

    PORT_Free (buf);
    PORT_Free (name);

    fclose (sf);
    fclose (mf);

    return 0;
}


/*
 *  c a l c u l a t e _ M D 5 _ r a n g e
 *
 *  Calculate the MD5 digest on a range of bytes in
 *  the specified fopen'd file. Returns base64.
 *
 */
static int	
calculate_MD5_range (FILE *fp, long r1, long r2, JAR_Digest *dig)
{
    int	num;
    int	range;
    unsigned char	*buf;
    SECStatus rv;

    range = r2 - r1;

    /* position to the beginning of range */
    fseek (fp, r1, SEEK_SET);

    buf = (unsigned char *) PORT_ZAlloc (range);
    if (buf == NULL)
	out_of_memory();

    if ((num = fread (buf, 1, range, fp)) != range) {
	PR_fprintf(errorFD, "%s: expected %d bytes, got %d\n", PROGRAM_NAME,
	     		range, num);
	errorCount++;
	exit (ERRX);
    }

    rv = PK11_HashBuf(SEC_OID_MD5, dig->md5, buf, range);
    if (rv == SECSuccess) {
	rv =PK11_HashBuf(SEC_OID_SHA1, dig->sha1, buf, range);
    }
    if (rv != SECSuccess) {
	PR_fprintf(errorFD, "%s: can't generate digest context\n",
	     PROGRAM_NAME);
	errorCount++;
	exit (ERRX);
    }

    PORT_Free (buf);

    return 0;
}


static void	SignOut (void *arg, const char *buf, unsigned long len)
{
    fwrite (buf, len, 1, (FILE * ) arg);
}