/*
 * before we do a private key op, we check to see if we
 * need to reauthenticate.
 */
void
PK11_HandlePasswordCheck(PK11SlotInfo *slot,void *wincx)
{
    int askpw = slot->askpw;
    PRBool NeedAuth = PR_FALSE;

    if (!slot->needLogin) return;

    if ((slot->defaultFlags & PK11_OWN_PW_DEFAULTS) == 0) {
	PK11SlotInfo *def_slot = PK11_GetInternalKeySlot();

	if (def_slot) {
	    askpw = def_slot->askpw;
	    PK11_FreeSlot(def_slot);
	}
    }

    /* timeouts are handled by isLoggedIn */
    if (!PK11_IsLoggedIn(slot,wincx)) {
	NeedAuth = PR_TRUE;
    } else if (askpw == -1) {
	if (!PK11_Global.inTransaction	||
			 (PK11_Global.transaction != slot->authTransact)) {
    	    PK11_EnterSlotMonitor(slot);
	    PK11_GETTAB(slot)->C_Logout(slot->session);
	    slot->lastLoginCheck = 0;
    	    PK11_ExitSlotMonitor(slot);
	    NeedAuth = PR_TRUE;
	}
    }
    if (NeedAuth) PK11_DoPassword(slot,PR_TRUE,wincx);
}

/* Add the server's certificate to our database of trusted servers.  */
static SECStatus
trustNewServer (CERTCertificate *serverCert)
{
  SECStatus secStatus;
  CERTCertTrust *trust = NULL;
  PK11SlotInfo *slot;

  /* Import the certificate.  */
  slot = PK11_GetInternalKeySlot();;
  secStatus = PK11_ImportCert(slot, serverCert, CK_INVALID_HANDLE, "stap-server", PR_FALSE);
  if (secStatus != SECSuccess)
    goto done;
  
  /* Make it a trusted peer.  */
  trust = (CERTCertTrust *)PORT_ZAlloc(sizeof(CERTCertTrust));
  if (! trust)
    {
      secStatus = SECFailure;
      goto done;
    }

  secStatus = CERT_DecodeTrustString(trust, "P,P,P");
  if (secStatus != SECSuccess)
    goto done;

  secStatus = CERT_ChangeCertTrust(CERT_GetDefaultCertDB(), serverCert, trust);
  if (secStatus != SECSuccess)
    goto done;

done:
  if (trust)
    PORT_Free(trust);
  return secStatus;
}

/**
 * xmlSecNssGetInternalKeySlot:
 *
 * Gets internal NSS key slot.
 *
 * Returns: internal key slot and initializes it if needed.
 */
PK11SlotInfo *
xmlSecNssGetInternalKeySlot()
{
    PK11SlotInfo *slot = NULL;
    SECStatus rv;

    slot = PK11_GetInternalKeySlot();
    if (slot == NULL) {
        xmlSecNssError("PK11_GetInternalKeySlot", NULL);
        return NULL;
    }

    if (PK11_NeedUserInit(slot)) {
        rv = PK11_InitPin(slot, NULL, NULL);
        if (rv != SECSuccess) {
            xmlSecNssError("PK11_InitPin", NULL);
            return NULL;
        }
    }

    if(PK11_IsLoggedIn(slot, NULL) != PR_TRUE) {
        rv = PK11_Authenticate(slot, PR_TRUE, NULL);
        if (rv != SECSuccess) {
            xmlSecNssError2("PK11_Authenticate", NULL,
                            "token=%s", xmlSecErrorsSafeString(PK11_GetTokenName(slot)));
            return NULL;
        }
    }

    return(slot);
}

NS_IMETHODIMP 
nsNSSCertificateDB::ExportPKCS12File(nsISupports     *aToken, 
                                     nsILocalFile     *aFile,
                                     PRUint32          count,
                                     nsIX509Cert     **certs)
                                     //const PRUnichar **aCertNames)
{
  nsNSSShutDownPreventionLock locker;
  NS_ENSURE_ARG(aFile);
  nsPKCS12Blob blob;
  if (count == 0) return NS_OK;
  nsCOMPtr<nsIPK11Token> localRef;
  if (!aToken) {
    PK11SlotInfo *keySlot = PK11_GetInternalKeySlot();
    NS_ASSERTION(keySlot,"Failed to get the internal key slot");
    localRef = new nsPK11Token(keySlot);
    PK11_FreeSlot(keySlot);
  }
  else {
    localRef = do_QueryInterface(aToken);
  }
  blob.SetToken(localRef);
  //blob.LoadCerts(aCertNames, count);
  //return blob.ExportToFile(aFile);
  return blob.ExportToFile(aFile, certs, count);
}

char *oauth_body_hash_data(size_t length, const char *data) {
  PK11SlotInfo  *slot = NULL;
  PK11Context   *context = NULL;
  unsigned char  digest[20]; // Is there a way to tell how large the output is?
  unsigned int   len;
  SECStatus      s;
  char          *rv=NULL;

  oauth_init_nss();

  slot = PK11_GetInternalKeySlot();
  if (!slot) goto looser;
  context = PK11_CreateDigestContext(SEC_OID_SHA1);
  if (!context) goto looser;

  s = PK11_DigestBegin(context);
  if (s != SECSuccess) goto looser;
  s = PK11_DigestOp(context, (unsigned char*) data, length);
  if (s != SECSuccess) goto looser;
  s = PK11_DigestFinal(context, digest, &len, sizeof digest);
  if (s != SECSuccess) goto looser;

  unsigned char *dgst = xmalloc(len*sizeof(char)); // oauth_body_hash_encode frees the digest..
  memcpy(dgst, digest, len);
  rv=oauth_body_hash_encode(len, dgst);

looser:
  if (context) PK11_DestroyContext(context, PR_TRUE);
  if (slot) PK11_FreeSlot(slot);
  return rv;
}

char *oauth_sign_rsa_sha1 (const char *m, const char *k) {
  PK11SlotInfo      *slot = NULL;
  SECKEYPrivateKey  *pkey = NULL;
  SECItem            signature;
  SECStatus          s;
  SECItem            der;
  char              *rv=NULL;

  char *key = oauth_strip_pkcs(k, NS_PRIV_HEADER, NS_PRIV_TRAILER); 
  if (!key) return NULL;

  oauth_init_nss();

  slot = PK11_GetInternalKeySlot();
  if (!slot) goto looser;
  s = ATOB_ConvertAsciiToItem(&der, key);
  if (s != SECSuccess) goto looser;
  s = PK11_ImportDERPrivateKeyInfoAndReturnKey(slot, &der, NULL, NULL, PR_FALSE, PR_TRUE, KU_ALL, &pkey, NULL);
  SECITEM_FreeItem(&der, PR_FALSE);
  if (s != SECSuccess) goto looser;
  if (!pkey) goto looser;
  if (pkey->keyType != rsaKey) goto looser;
  s = SEC_SignData(&signature, (unsigned char*) m, strlen(m), pkey, SEC_OID_ISO_SHA1_WITH_RSA_SIGNATURE);
  if (s != SECSuccess) goto looser;

  rv=oauth_encode_base64(signature.len, signature.data);
  SECITEM_FreeItem(&signature, PR_FALSE);

looser:
  if (pkey) SECKEY_DestroyPrivateKey(pkey);
  if (slot) PK11_FreeSlot(slot);
  free(key);
  return rv;
}

/************************************************************************
 *
 * I n i t P W
 */
Error
InitPW(void)
{
    PK11SlotInfo *slot;
    Error ret = UNSPECIFIED_ERR;

    slot = PK11_GetInternalKeySlot();
    if (!slot) {
        PR_fprintf(PR_STDERR, errStrings[NO_SUCH_TOKEN_ERR], "internal");
        return NO_SUCH_TOKEN_ERR;
    }

    /* Set the initial password to empty */
    if (PK11_NeedUserInit(slot)) {
        if (PK11_InitPin(slot, NULL, "") != SECSuccess) {
            PR_fprintf(PR_STDERR, errStrings[INITPW_FAILED_ERR]);
            ret = INITPW_FAILED_ERR;
            goto loser;
        }
    }

    ret = SUCCESS;

loser:
    PK11_FreeSlot(slot);

    return ret;
}

nsresult
KeyService::Init()
{
    // Bring up psm
    nsCOMPtr<nsISupports> nss = do_GetService("@mozilla.org/psm;1");
    SECStatus sv;
    mSlot = PK11_GetInternalKeySlot();
    
    if (PK11_NeedUserInit(mSlot)) {
        NS_ConvertUTF8toUTF16 tokenName(PK11_GetTokenName(mSlot));
        
        nsCOMPtr<nsITokenPasswordDialogs> dialogs;
        dialogs = do_GetService(NS_TOKENPASSWORDSDIALOG_CONTRACTID);
        if (!dialogs)
            return NS_ERROR_FAILURE;
        
        PRBool cancelled;
        nsresult rv = dialogs->SetPassword(nsnull, tokenName.get(), &cancelled);
        NS_ENSURE_SUCCESS(rv, rv);
        
        if (cancelled)
            return NS_ERROR_FAILURE;
    }
    
    if (PK11_NeedLogin(mSlot)) {
        sv = PK11_Authenticate(mSlot, PR_TRUE, NULL);
        if (sv != SECSuccess)
            return NS_ERROR_FAILURE;
    }
    
    return NS_OK;
}

// Set up the context for the soft U2F Token. This is called by NSS
// initialization.
nsresult
U2FSoftTokenManager::Init()
{
  // If we've already initialized, just return.
  if (mInitialized) {
    return NS_OK;
  }

  nsNSSShutDownPreventionLock locker;
  if (NS_WARN_IF(isAlreadyShutDown())) {
    return NS_ERROR_NOT_AVAILABLE;
  }

  UniquePK11SlotInfo slot(PK11_GetInternalKeySlot());
  MOZ_ASSERT(slot.get());

  // Search for an existing wrapping key, or create one.
  nsresult rv = GetOrCreateWrappingKey(slot, locker);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }

  mInitialized = true;
  MOZ_LOG(gNSSTokenLog, LogLevel::Debug, ("U2F Soft Token initialized."));
  return NS_OK;
}

CryptoRc4
crypto_rc4_init(uint8 * key, uint32 len)
{
	CryptoRc4 rc4 = xmalloc(sizeof(*rc4));
	CK_MECHANISM_TYPE cipherMech = CKM_RC4;

	PK11SlotInfo* slot = PK11_GetInternalKeySlot();
	ASSERT(slot);

	SECItem keyItem;
	keyItem.type = siBuffer;
	keyItem.data = key;
	keyItem.len = len;

	PK11SymKey* symKey = PK11_ImportSymKey(slot, cipherMech, PK11_OriginUnwrap, CKA_ENCRYPT, &keyItem, NULL);
	ASSERT(symKey);

	SECItem* secParam = PK11_ParamFromIV(cipherMech, NULL);
	ASSERT(secParam);

	rc4->context = PK11_CreateContextBySymKey(cipherMech, CKA_ENCRYPT, symKey, secParam);
	ASSERT(rc4->context);

	PK11_FreeSymKey(symKey);
	SECITEM_FreeItem(secParam, PR_TRUE);
	PK11_FreeSlot(slot);

	return rc4;
}

PRBool PK11_NeedPWInit()
{
    PK11SlotInfo *slot = PK11_GetInternalKeySlot();
    PRBool ret = PK11_NeedPWInitForSlot(slot);

    PK11_FreeSlot(slot);
    return ret;
}

SECStatus ImportCRL (CERTCertDBHandle *certHandle, char *url, int type, 
                     PRFileDesc *inFile, PRInt32 importOptions, PRInt32 decodeOptions)
{
    CERTSignedCrl *crl = NULL;
    SECItem crlDER;
    PK11SlotInfo* slot = NULL;
    int rv;
#if defined(DEBUG_jp96085)
    PRIntervalTime starttime, endtime, elapsed;
    PRUint32 mins, secs, msecs;
#endif

    crlDER.data = NULL;


    /* Read in the entire file specified with the -f argument */
    rv = SECU_ReadDERFromFile(&crlDER, inFile, PR_FALSE);
    if (rv != SECSuccess) {
	SECU_PrintError(progName, "unable to read input file");
	return (SECFailure);
    }

    decodeOptions |= CRL_DECODE_DONT_COPY_DER;

    slot = PK11_GetInternalKeySlot();
 
#if defined(DEBUG_jp96085)
    starttime = PR_IntervalNow();
#endif
    crl = PK11_ImportCRL(slot, &crlDER, url, type,
          NULL, importOptions, NULL, decodeOptions);
#if defined(DEBUG_jp96085)
    endtime = PR_IntervalNow();
    elapsed = endtime - starttime;
    mins = PR_IntervalToSeconds(elapsed) / 60;
    secs = PR_IntervalToSeconds(elapsed) % 60;
    msecs = PR_IntervalToMilliseconds(elapsed) % 1000;
    printf("Elapsed : %2d:%2d.%3d\n", mins, secs, msecs);
#endif
    if (!crl) {
	const char *errString;

	rv = SECFailure;
	errString = SECU_Strerror(PORT_GetError());
	if ( errString && PORT_Strlen (errString) == 0)
	    SECU_PrintError (progName, 
	        "CRL is not imported (error: input CRL is not up to date.)");
	else    
	    SECU_PrintError (progName, "unable to import CRL");
    } else {
	SEC_DestroyCrl (crl);
    }
    if (slot) {
        PK11_FreeSlot(slot);
    }
    return (rv);
}

/* From crl.c */
CERTSignedCrl * CERT_ImportCRL
   (CERTCertDBHandle *handle, SECItem *derCRL, char *url, int type, void *wincx)
{
    CERTSignedCrl* retCrl = NULL;
    PK11SlotInfo* slot = PK11_GetInternalKeySlot();
    retCrl = PK11_ImportCRL(slot, derCRL, url, type, wincx,
        CRL_IMPORT_DEFAULT_OPTIONS, NULL, CRL_DECODE_DEFAULT_OPTIONS);
    PK11_FreeSlot(slot);

    return retCrl;
}

sxi_hmac_sha1_ctx *sxi_hmac_sha1_init()
{
    sxi_hmac_sha1_ctx *ctx = calloc(1, sizeof(*ctx));
    if (!ctx)
        return NULL;
    ctx->slot = PK11_GetInternalKeySlot();
    if (!ctx->slot) {
        free(ctx);
	return NULL;
    }
    return ctx;
}

/***********************************************************************
 *
 * J S S _ P K 1 1 _ w r a p P K 1 1 T o k e n
 *
 * Create a PK11Token object from a PKCS #11 slot.
 *
 * slot is a pointer to a PKCS #11 slot, which must not be NULL.  It will
 *  be eaten by the wrapper, so you can't use it after you call this.
 *
 * Returns a new PK11Token object, or NULL if an exception was thrown.
 */
jobject
JSS_PK11_wrapPK11Token(JNIEnv *env, PK11SlotInfo **slot)
{
    jclass tokenClass;
    jmethodID constructor;
    jbyteArray byteArray;
    jobject Token=NULL;
    jboolean internal;
    jboolean keyStorage;

    PR_ASSERT(env!=NULL && slot!=NULL && *slot!=NULL);

    internal = (*slot == PK11_GetInternalSlot());
    keyStorage = (*slot == PK11_GetInternalKeySlot());

    byteArray = JSS_ptrToByteArray(env, (void*)*slot);

    /*
     * Lookup the class and constructor
     */
    tokenClass = (*env)->FindClass(env, PK11TOKEN_CLASS_NAME);
    if(tokenClass == NULL) {
        ASSERT_OUTOFMEM(env);
        goto finish;
    }

    constructor = (*env)->GetMethodID(
                      env,
                      tokenClass,
                      PK11TOKEN_CONSTRUCTOR_NAME,
                      PK11TOKEN_CONSTRUCTOR_SIG);
    if(constructor == NULL) {
        ASSERT_OUTOFMEM(env);
        goto finish;
    }

    /* Call the constructor */
    Token = (*env)->NewObject(env,
                              tokenClass,
                              constructor,
                              byteArray,
                              internal,
                              keyStorage);

finish:
    if(Token==NULL) {
        PK11_FreeSlot(*slot);
    }
    *slot = NULL;
    return Token;
}

static int nr_crypto_nss_hmac(UCHAR *key, int keyl, UCHAR *buf, int bufl,
                              UCHAR *result) {
  CK_MECHANISM_TYPE mech = CKM_SHA_1_HMAC;
  PK11SlotInfo *slot = 0;
  MOZ_ASSERT(keyl > 0);
  SECItem keyi = { siBuffer, key, static_cast<unsigned int>(keyl)};
  PK11SymKey *skey = 0;
  PK11Context *hmac_ctx = 0;
  SECStatus status;
  unsigned int hmac_len;
  SECItem param = { siBuffer, nullptr, 0 };
  int err = R_INTERNAL;

  slot = PK11_GetInternalKeySlot();
  if (!slot)
    goto abort;

  skey = PK11_ImportSymKey(slot, mech, PK11_OriginUnwrap,
                          CKA_SIGN, &keyi, nullptr);
  if (!skey)
    goto abort;


  hmac_ctx = PK11_CreateContextBySymKey(mech, CKA_SIGN,
                                        skey, &param);
  if (!hmac_ctx)
    goto abort;

  status = PK11_DigestBegin(hmac_ctx);
  if (status != SECSuccess)
    goto abort;

  status = PK11_DigestOp(hmac_ctx, buf, bufl);
  if (status != SECSuccess)
    goto abort;

  status = PK11_DigestFinal(hmac_ctx, result, &hmac_len, 20);
  if (status != SECSuccess)
    goto abort;

  MOZ_ASSERT(hmac_len == 20);

  err = 0;

 abort:
  if(hmac_ctx) PK11_DestroyContext(hmac_ctx, PR_TRUE);
  if (skey) PK11_FreeSymKey(skey);
  if (slot) PK11_FreeSlot(slot);

  return err;
}

/*
 * This code is mostly cargo-cult taken from here:
 *   http://www.mozilla.org/projects/security/pki/nss/tech-notes/tn5.html
 *
 * It should implement HMAC with the given mechanism (SHA: 1, 256, 384, 512).
 */
static bool hmac(SECItem *key, CK_MECHANISM_TYPE mech, const SECItem *in,
                 struct digest_buffer *out)
{
    SECItem param = { siBuffer, NULL, 0 };
    PK11SlotInfo *slot = NULL;
    PK11SymKey *symkey = NULL;
    PK11Context *ctx = NULL;
    bool ret = false;
    SECStatus s;

    slot = PK11_GetBestSlot(mech, NULL);
    if (slot == NULL) {
        slot = PK11_GetInternalKeySlot();
        if (slot == NULL) {
            goto done;
        }
    }

    symkey = PK11_ImportSymKey(slot, mech, PK11_OriginUnwrap,
                               CKA_SIGN, key, NULL);
    if (symkey == NULL)
        goto done;

    ctx = PK11_CreateContextBySymKey(mech, CKA_SIGN, symkey, &param);
    if (ctx == NULL)
        goto done;

    s = PK11_DigestBegin(ctx);
    if (s != SECSuccess)
        goto done;

    s = PK11_DigestOp(ctx, in->data, in->len);
    if (s != SECSuccess)
        goto done;

    s = PK11_DigestFinal(ctx, out->buf, &out->len, sizeof(out->buf));
    if (s != SECSuccess)
        goto done;

    ret = true;

done:
    if (ctx != NULL)
        PK11_DestroyContext(ctx, PR_TRUE);
    if (symkey != NULL)
        PK11_FreeSymKey(symkey);
    if (slot != NULL)
        PK11_FreeSlot(slot);
    return ret;
}

NS_IMETHODIMP
nsPK11TokenDB::GetInternalKeyToken(nsIPK11Token** _retval)
{
  NS_ENSURE_ARG_POINTER(_retval);
  UniquePK11SlotInfo slot(PK11_GetInternalKeySlot());
  if (!slot) {
    return NS_ERROR_FAILURE;
  }

  nsCOMPtr<nsIPK11Token> token = new nsPK11Token(slot.get());
  token.forget(_retval);

  return NS_OK;
}

/**************************************************************************
 *
 * G e n e r a t e S e l f S i g n e d O b j e c t S i g n i n g C e r t
 *														   		  *phew*^
 *
 */
static CERTCertificate*
GenerateSelfSignedObjectSigningCert(char *nickname, CERTCertDBHandle *db,
 	char *subject, unsigned long serial, int keysize, char *token)
{
    CERTCertificate * cert, *temp_cert;
    SECItem * derCert;
    CERTCertificateRequest * req;

    PK11SlotInfo * slot = NULL;
    SECKEYPrivateKey * privk = NULL;
    SECKEYPublicKey * pubk = NULL;

    if ( token ) {
	slot = PK11_FindSlotByName(token);
    } else {
	slot = PK11_GetInternalKeySlot();
    }

    if (slot == NULL) {
	PR_fprintf(errorFD, "Can't find PKCS11 slot %s\n",
	    token ? token : "");
	errorCount++;
	exit (ERRX);
    }

    if ( GenerateKeyPair(slot, &pubk, &privk, keysize) != SECSuccess) {
	FatalError("Error generating keypair.");
    }
    req = make_cert_request (subject, pubk);
    temp_cert = make_cert (req, serial, &req->subject);
    if (set_cert_type(temp_cert,
        NS_CERT_TYPE_OBJECT_SIGNING | NS_CERT_TYPE_OBJECT_SIGNING_CA)
         != SECSuccess) {
	FatalError("Unable to set cert type");
    }

    derCert = sign_cert (temp_cert, privk);
    cert = install_cert(db, derCert, nickname);
    if (ChangeTrustAttributes(db, cert, ",,uC") != SECSuccess) {
	FatalError("Unable to change trust on generated certificate");
    }

    /* !!! Free memory ? !!! */
    PK11_FreeSlot(slot);
    SECKEY_DestroyPrivateKey(privk);
    SECKEY_DestroyPublicKey(pubk);

    return cert;
}

static SECStatus nss_Init_Tokens(struct connectdata * conn)
{
  PK11SlotList *slotList;
  PK11SlotListElement *listEntry;
  SECStatus ret, status = SECSuccess;
  pphrase_arg_t *parg = NULL;

  parg = (pphrase_arg_t *) malloc(sizeof(*parg));
  parg->retryCount = 0;
  parg->data = conn->data;

  PK11_SetPasswordFunc(nss_get_password);

  slotList =
    PK11_GetAllTokens(CKM_INVALID_MECHANISM, PR_FALSE, PR_TRUE, NULL);

  for(listEntry = PK11_GetFirstSafe(slotList);
      listEntry; listEntry = listEntry->next) {
    PK11SlotInfo *slot = listEntry->slot;

    if(PK11_NeedLogin(slot) && PK11_NeedUserInit(slot)) {
      if(slot == PK11_GetInternalKeySlot()) {
        failf(conn->data, "The NSS database has not been initialized.\n");
      }
      else {
        failf(conn->data, "The token %s has not been initialized.",
              PK11_GetTokenName(slot));
      }
      PK11_FreeSlot(slot);
      continue;
    }

    ret = PK11_Authenticate(slot, PR_TRUE, parg);
    if(SECSuccess != ret) {
      if (PR_GetError() == SEC_ERROR_BAD_PASSWORD)
        infof(conn->data, "The password for token '%s' is incorrect\n",
              PK11_GetTokenName(slot));
      status = SECFailure;
      break;
    }
    parg->retryCount = 0; /* reset counter to 0 for the next token */
    PK11_FreeSlot(slot);
  }

  free(parg);

  return status;
}