package cn.jiguang.common.utils.sm2;

import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.ASN1Encoding;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.ASN1Primitive;
import org.bouncycastle.asn1.DERNull;
import org.bouncycastle.asn1.DEROctetString;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x9.X962Parameters;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.asn1.x9.X9ObjectIdentifiers;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.generators.ECKeyPairGenerator;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECKeyGenerationParameters;
import org.bouncycastle.crypto.params.ECKeyParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jcajce.provider.asymmetric.util.EC5Util;
import org.bouncycastle.jcajce.provider.asymmetric.util.ECUtil;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECNamedCurveSpec;
import org.bouncycastle.jce.spec.ECParameterSpec;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.pqc.math.linearalgebra.ByteUtils;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemReader;
import org.bouncycastle.util.io.pem.PemWriter;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * 这个工具类的方法，也适用于其他基于BC库的ECC算法
 */
public class BCECUtil {
    private static final String ALGO_NAME_EC = "EC";
    private static final String PEM_STRING_PUBLIC = "PUBLIC KEY";
    private static final String PEM_STRING_ECPRIVATEKEY = "EC PRIVATE KEY";

    /**
     * 生成ECC密钥对
     *
     * @return ECC密钥对
     */
    public static AsymmetricCipherKeyPair generateKeyPairParameter(ECDomainParameters domainParameters,
        SecureRandom random) {
        ECKeyGenerationParameters keyGenerationParams = new ECKeyGenerationParameters(domainParameters,
            random);
        ECKeyPairGenerator keyGen = new ECKeyPairGenerator();
        keyGen.init(keyGenerationParams);
        return keyGen.generateKeyPair();
    }

    public static KeyPair generateKeyPair(ECDomainParameters domainParameters, SecureRandom random)
        throws NoSuchProviderException, NoSuchAlgorithmException,
        InvalidAlgorithmParameterException {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance(ALGO_NAME_EC, BouncyCastleProvider.PROVIDER_NAME);
        ECParameterSpec parameterSpec = new ECParameterSpec(domainParameters.getCurve(), domainParameters.getG(),
            domainParameters.getN(), domainParameters.getH());
        kpg.initialize(parameterSpec, random);
        return kpg.generateKeyPair();
    }

    public static int getCurveLength(ECKeyParameters ecKey) {
        return getCurveLength(ecKey.getParameters());
    }

    public static int getCurveLength(ECDomainParameters domainParams) {
        return (domainParams.getCurve().getFieldSize() + 7) / 8;
    }

    public static byte[] fixToCurveLengthBytes(int curveLength, byte[] src) {
        if (src.length == curveLength) {
            return src;
        }

        byte[] result = new byte[curveLength];
        if (src.length > curveLength) {
            System.arraycopy(src, src.length - result.length, result, 0, result.length);
        } else {
            System.arraycopy(src, 0, result, result.length - src.length, src.length);
        }
        return result;
    }

    public static ECPrivateKeyParameters createECPrivateKeyParameters(BigInteger d,
        ECDomainParameters domainParameters) {
        return new ECPrivateKeyParameters(d, domainParameters);
    }

    public static ECPublicKeyParameters createECPublicKeyParameters(BigInteger x, BigInteger y,
        ECCurve curve, ECDomainParameters domainParameters) {
        return createECPublicKeyParameters(x.toByteArray(), y.toByteArray(), curve, domainParameters);
    }

    public static ECPublicKeyParameters createECPublicKeyParameters(String xHex, String yHex,
        ECCurve curve, ECDomainParameters domainParameters) {
        return createECPublicKeyParameters(ByteUtils.fromHexString(xHex), ByteUtils.fromHexString(yHex),
            curve, domainParameters);
    }

    public static ECPublicKeyParameters createECPublicKeyParameters(byte[] xBytes, byte[] yBytes,
        ECCurve curve, ECDomainParameters domainParameters) {
        final byte uncompressedFlag = 0x04;
        int curveLength = getCurveLength(domainParameters);
        xBytes = fixToCurveLengthBytes(curveLength, xBytes);
        yBytes = fixToCurveLengthBytes(curveLength, yBytes);
        byte[] encodedPubKey = new byte[1 + xBytes.length + yBytes.length];
        encodedPubKey[0] = uncompressedFlag;
        System.arraycopy(xBytes, 0, encodedPubKey, 1, xBytes.length);
        System.arraycopy(yBytes, 0, encodedPubKey, 1 + xBytes.length, yBytes.length);
        return new ECPublicKeyParameters(curve.decodePoint(encodedPubKey), domainParameters);
    }

    public static ECPrivateKeyParameters convertPrivateKeyToParameters(BCECPrivateKey ecPriKey) {
        ECParameterSpec parameterSpec = ecPriKey.getParameters();
        ECDomainParameters domainParameters = new ECDomainParameters(parameterSpec.getCurve(), parameterSpec.getG(),
            parameterSpec.getN(), parameterSpec.getH());
        return new ECPrivateKeyParameters(ecPriKey.getD(), domainParameters);
    }

    public static ECPublicKeyParameters convertPublicKeyToParameters(BCECPublicKey ecPubKey) {
        ECParameterSpec parameterSpec = ecPubKey.getParameters();
        ECDomainParameters domainParameters = new ECDomainParameters(parameterSpec.getCurve(), parameterSpec.getG(),
            parameterSpec.getN(), parameterSpec.getH());
        return new ECPublicKeyParameters(ecPubKey.getQ(), domainParameters);
    }

    public static BCECPublicKey createPublicKeyFromSubjectPublicKeyInfo(SubjectPublicKeyInfo subPubInfo) throws NoSuchProviderException,
        NoSuchAlgorithmException, InvalidKeySpecException, IOException {
        return BCECUtil.convertX509ToECPublicKey(subPubInfo.toASN1Primitive().getEncoded(ASN1Encoding.DER));
    }

    /**
     * 将ECC私钥转换为PKCS8标准的字节流
     *
     * @param priKey
     * @param pubKey 可以为空，但是如果为空的话得到的结果OpenSSL可能解析不了
     * @return
     */
    public static byte[] convertECPrivateKeyToPKCS8(ECPrivateKeyParameters priKey,
        ECPublicKeyParameters pubKey) {
        ECDomainParameters domainParams = priKey.getParameters();
        ECParameterSpec spec = new ECParameterSpec(domainParams.getCurve(), domainParams.getG(),
            domainParams.getN(), domainParams.getH());
        BCECPublicKey publicKey = null;
        if (pubKey != null) {
            publicKey = new BCECPublicKey(ALGO_NAME_EC, pubKey, spec,
                BouncyCastleProvider.CONFIGURATION);
        }
        BCECPrivateKey privateKey = new BCECPrivateKey(ALGO_NAME_EC, priKey, publicKey,
            spec, BouncyCastleProvider.CONFIGURATION);
        return privateKey.getEncoded();
    }

    /**
     * 将PKCS8标准的私钥字节流转换为私钥对象
     *
     * @param pkcs8Key
     * @return
     * @throws NoSuchAlgorithmException
     * @throws NoSuchProviderException
     * @throws InvalidKeySpecException
     */
    public static BCECPrivateKey convertPKCS8ToECPrivateKey(byte[] pkcs8Key)
        throws NoSuchAlgorithmException, NoSuchProviderException, InvalidKeySpecException {
        PKCS8EncodedKeySpec peks = new PKCS8EncodedKeySpec(pkcs8Key);
        KeyFactory kf = KeyFactory.getInstance(ALGO_NAME_EC, BouncyCastleProvider.PROVIDER_NAME);
        return (BCECPrivateKey) kf.generatePrivate(peks);
    }

    /**
     * 将PKCS8标准的私钥字节流转换为PEM
     *
     * @param encodedKey
     * @return
     * @throws IOException
     */
    public static String convertECPrivateKeyPKCS8ToPEM(byte[] encodedKey) throws IOException {
        return convertEncodedDataToPEM(PEM_STRING_ECPRIVATEKEY, encodedKey);
    }

    /**
     * 将PEM格式的私钥转换为PKCS8标准字节流
     *
     * @param pemString
     * @return
     * @throws IOException
     */
    public static byte[] convertECPrivateKeyPEMToPKCS8(String pemString) throws IOException {
        return convertPEMToEncodedData(pemString);
    }

    /**
     * 将ECC私钥转换为SEC1标准的字节流
     * openssl d2i_ECPrivateKey函数要求的DER编码的私钥也是SEC1标准的，
     * 这个工具函数的主要目的就是为了能生成一个openssl可以直接“识别”的ECC私钥.
     * 相对RSA私钥的PKCS1标准，ECC私钥的标准为SEC1
     *
     * @param priKey
     * @param pubKey
     * @return
     * @throws IOException
     */
    public static byte[] convertECPrivateKeyToSEC1(ECPrivateKeyParameters priKey,
        ECPublicKeyParameters pubKey) throws IOException {
        byte[] pkcs8Bytes = convertECPrivateKeyToPKCS8(priKey, pubKey);
        PrivateKeyInfo pki = PrivateKeyInfo.getInstance(pkcs8Bytes);
        ASN1Encodable encodable = pki.parsePrivateKey();
        ASN1Primitive primitive = encodable.toASN1Primitive();
        byte[] sec1Bytes = primitive.getEncoded();
        return sec1Bytes;
    }

    /**
     * 将SEC1标准的私钥字节流恢复为PKCS8标准的字节流
     *
     * @param sec1Key
     * @return
     * @throws IOException
     */
    public static byte[] convertECPrivateKeySEC1ToPKCS8(byte[] sec1Key) throws IOException {
        /**
         * 参考org.bouncycastle.asn1.pkcs.PrivateKeyInfo和
         * org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey，逆向拼装
         */
        X962Parameters params = getDomainParametersFromName(SM2Util.JDK_EC_SPEC, false);
        ASN1OctetString privKey = new DEROctetString(sec1Key);
        ASN1EncodableVector v = new ASN1EncodableVector();
        v.add(new ASN1Integer(0)); //版本号
        v.add(new AlgorithmIdentifier(X9ObjectIdentifiers.id_ecPublicKey, params)); //算法标识
        v.add(privKey);
        DERSequence ds = new DERSequence(v);
        return ds.getEncoded(ASN1Encoding.DER);
    }

    /**
     * 将SEC1标准的私钥字节流转为BCECPrivateKey对象
     *
     * @param sec1Key
     * @return
     * @throws NoSuchAlgorithmException
     * @throws NoSuchProviderException
     * @throws InvalidKeySpecException
     * @throws IOException
     */
    public static BCECPrivateKey convertSEC1ToBCECPrivateKey(byte[] sec1Key)
        throws NoSuchAlgorithmException, NoSuchProviderException, InvalidKeySpecException, IOException {
        PKCS8EncodedKeySpec peks = new PKCS8EncodedKeySpec(convertECPrivateKeySEC1ToPKCS8(sec1Key));
        KeyFactory kf = KeyFactory.getInstance(ALGO_NAME_EC, BouncyCastleProvider.PROVIDER_NAME);
        return (BCECPrivateKey) kf.generatePrivate(peks);
    }

    /**
     * 将SEC1标准的私钥字节流转为ECPrivateKeyParameters对象
     * openssl i2d_ECPrivateKey函数生成的DER编码的ecc私钥是：SEC1标准的、带有EC_GROUP、带有公钥的，
     * 这个工具函数的主要目的就是为了使Java程序能够“识别”openssl生成的ECC私钥
     *
     * @param sec1Key
     * @return
     * @throws NoSuchAlgorithmException
     * @throws NoSuchProviderException
     * @throws InvalidKeySpecException
     */
    public static ECPrivateKeyParameters convertSEC1ToECPrivateKey(byte[] sec1Key)
        throws NoSuchAlgorithmException, NoSuchProviderException, InvalidKeySpecException, IOException {
        BCECPrivateKey privateKey = convertSEC1ToBCECPrivateKey(sec1Key);
        return convertPrivateKeyToParameters(privateKey);
    }

    /**
     * 将ECC公钥对象转换为X509标准的字节流
     *
     * @param pubKey
     * @return
     */
    public static byte[] convertECPublicKeyToX509(ECPublicKeyParameters pubKey) {
        ECDomainParameters domainParams = pubKey.getParameters();
        ECParameterSpec spec = new ECParameterSpec(domainParams.getCurve(), domainParams.getG(),
            domainParams.getN(), domainParams.getH());
        BCECPublicKey publicKey = new BCECPublicKey(ALGO_NAME_EC, pubKey, spec,
            BouncyCastleProvider.CONFIGURATION);
        return publicKey.getEncoded();
    }

    /**
     * 将X509标准的公钥字节流转为公钥对象
     *
     * @param x509Bytes
     * @return
     * @throws NoSuchProviderException
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public static BCECPublicKey convertX509ToECPublicKey(byte[] x509Bytes) throws NoSuchProviderException,
        NoSuchAlgorithmException, InvalidKeySpecException {
        X509EncodedKeySpec eks = new X509EncodedKeySpec(x509Bytes);
        KeyFactory kf = KeyFactory.getInstance("EC", BouncyCastleProvider.PROVIDER_NAME);
        return (BCECPublicKey) kf.generatePublic(eks);
    }

    /**
     * 将X509标准的公钥字节流转为PEM
     *
     * @param encodedKey
     * @return
     * @throws IOException
     */
    public static String convertECPublicKeyX509ToPEM(byte[] encodedKey) throws IOException {
        return convertEncodedDataToPEM(PEM_STRING_PUBLIC, encodedKey);
    }

    /**
     * 将PEM格式的公钥转为X509标准的字节流
     *
     * @param pemString
     * @return
     * @throws IOException
     */
    public static byte[] convertECPublicKeyPEMToX509(String pemString) throws IOException {
        return convertPEMToEncodedData(pemString);
    }

    /**
     * copy from BC
     *
     * @param genSpec
     * @return
     */
    public static X9ECParameters getDomainParametersFromGenSpec(ECGenParameterSpec genSpec) {
        return getDomainParametersFromName(genSpec.getName());
    }

    /**
     * copy from BC
     *
     * @param curveName
     * @return
     */
    public static X9ECParameters getDomainParametersFromName(String curveName) {
        X9ECParameters domainParameters;
        try {
            if (curveName.charAt(0) >= '0' && curveName.charAt(0) <= '2') {
                ASN1ObjectIdentifier oidID = new ASN1ObjectIdentifier(curveName);
                domainParameters = ECUtil.getNamedCurveByOid(oidID);
            } else {
                if (curveName.indexOf(' ') > 0) {
                    curveName = curveName.substring(curveName.indexOf(' ') + 1);
                    domainParameters = ECUtil.getNamedCurveByName(curveName);
                } else {
                    domainParameters = ECUtil.getNamedCurveByName(curveName);
                }
            }
        } catch (IllegalArgumentException ex) {
            domainParameters = ECUtil.getNamedCurveByName(curveName);
        }
        return domainParameters;
    }

    /**
     * copy from BC
     *
     * @param ecSpec
     * @param withCompression
     * @return
     */
    public static X962Parameters getDomainParametersFromName(java.security.spec.ECParameterSpec ecSpec,
        boolean withCompression) {
        X962Parameters params;

        if (ecSpec instanceof ECNamedCurveSpec) {
            ASN1ObjectIdentifier curveOid = ECUtil.getNamedCurveOid(((ECNamedCurveSpec) ecSpec).getName());
            if (curveOid == null) {
                curveOid = new ASN1ObjectIdentifier(((ECNamedCurveSpec) ecSpec).getName());
            }
            params = new X962Parameters(curveOid);
        } else if (ecSpec == null) {
            params = new X962Parameters(DERNull.INSTANCE);
        } else {
            ECCurve curve = EC5Util.convertCurve(ecSpec.getCurve());

            X9ECParameters ecP = new X9ECParameters(
                curve,
                EC5Util.convertPoint(curve, ecSpec.getGenerator(), withCompression),
                ecSpec.getOrder(),
                BigInteger.valueOf(ecSpec.getCofactor()),
                ecSpec.getCurve().getSeed());

            params = new X962Parameters(ecP);
        }

        return params;
    }

    private static String convertEncodedDataToPEM(String type, byte[] encodedData) throws IOException {
        ByteArrayOutputStream bOut = new ByteArrayOutputStream();
        PemWriter pWrt = new PemWriter(new OutputStreamWriter(bOut));
        try {
            PemObject pemObj = new PemObject(type, encodedData);
            pWrt.writeObject(pemObj);
        } finally {
            pWrt.close();
        }
        return new String(bOut.toByteArray());
    }

    private static byte[] convertPEMToEncodedData(String pemString) throws IOException {
        ByteArrayInputStream bIn = new ByteArrayInputStream(pemString.getBytes());
        PemReader pRdr = new PemReader(new InputStreamReader(bIn));
        try {
            PemObject pemObject = pRdr.readPemObject();
            return pemObject.getContent();
        } finally {
            pRdr.close();
        }
    }
}


package cn.jiguang.common.utils.sm2;

import cn.jiguang.common.utils.Base64;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.math.ec.custom.gm.SM2P256V1Curve;

import java.math.BigInteger;
import java.security.*;
import java.security.spec.ECFieldFp;
import java.security.spec.EllipticCurve;
import java.util.concurrent.locks.ReentrantLock;

public class SM2Util extends GMBaseUtil {
    //////////////////////////////////////////////////////////////////////////////////////
    /*
     * 以下为SM2推荐曲线参数
     */
    public static final SM2P256V1Curve CURVE = new SM2P256V1Curve();
    public final static BigInteger SM2_ECC_P = CURVE.getQ();
    public final static BigInteger SM2_ECC_A = CURVE.getA().toBigInteger();
    public final static BigInteger SM2_ECC_B = CURVE.getB().toBigInteger();
    public final static BigInteger SM2_ECC_N = CURVE.getOrder();
    public final static BigInteger SM2_ECC_H = CURVE.getCofactor();
    public final static BigInteger SM2_ECC_GX = new BigInteger(
        "32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7", 16);
    public final static BigInteger SM2_ECC_GY = new BigInteger(
        "BC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0", 16);
    public static final ECPoint G_POINT = CURVE.createPoint(SM2_ECC_GX, SM2_ECC_GY);
    public static final ECDomainParameters DOMAIN_PARAMS = new ECDomainParameters(CURVE, G_POINT,
        SM2_ECC_N, SM2_ECC_H);
    public static final int CURVE_LEN = BCECUtil.getCurveLength(DOMAIN_PARAMS);
    //////////////////////////////////////////////////////////////////////////////////////

    public static final EllipticCurve JDK_CURVE = new EllipticCurve(new ECFieldFp(SM2_ECC_P), SM2_ECC_A, SM2_ECC_B);
    public static final java.security.spec.ECPoint JDK_G_POINT = new java.security.spec.ECPoint(
        G_POINT.getAffineXCoord().toBigInteger(), G_POINT.getAffineYCoord().toBigInteger());
    public static final java.security.spec.ECParameterSpec JDK_EC_SPEC = new java.security.spec.ECParameterSpec(
        JDK_CURVE, JDK_G_POINT, SM2_ECC_N, SM2_ECC_H.intValue());

    //////////////////////////////////////////////////////////////////////////////////////

    private static ECPublicKeyParameters RUNNING_PUBLIC_KEY;
    private static ReentrantLock LOCK = new ReentrantLock();

    /**
     * 加密数据
     * @param srcData
     * @param publicKey
     * @return
     * @throws Exception
     */
    public static byte[] encrypt(String srcData, String publicKey) throws Exception {
        if (RUNNING_PUBLIC_KEY == null) {
            try {
                LOCK.lock();
                if (RUNNING_PUBLIC_KEY == null) {
                    ECPoint ecPoint = CURVE.decodePoint(Base64.decode(publicKey.toCharArray()));
                    RUNNING_PUBLIC_KEY = new ECPublicKeyParameters(ecPoint, DOMAIN_PARAMS);
                }
            } catch (Exception e) {
                throw new RuntimeException("init public key error", e);
            } finally {
                LOCK.unlock();
            }
        }
        return encrypt(RUNNING_PUBLIC_KEY, srcData.getBytes());
    }

    /**
     * 生成ECC密钥对
     *
     * @return ECC密钥对
     */
    public static AsymmetricCipherKeyPair generateKeyPairParameter() {
        SecureRandom random = new SecureRandom();
        return BCECUtil.generateKeyPairParameter(DOMAIN_PARAMS, random);
    }

    public static KeyPair generateKeyPair() throws NoSuchProviderException, NoSuchAlgorithmException,
        InvalidAlgorithmParameterException {
        SecureRandom random = new SecureRandom();
        return BCECUtil.generateKeyPair(DOMAIN_PARAMS, random);
    }

    /**
     * 只获取私钥里的d，32字节
     *
     * @param privateKey
     * @return
     */
    public static byte[] getRawPrivateKey(BCECPrivateKey privateKey) {
        return fixToCurveLengthBytes(privateKey.getD().toByteArray());
    }

    /**
     * 只获取公钥里的XY分量，64字节
     *
     * @param publicKey
     * @return
     */
    public static byte[] getRawPublicKey(BCECPublicKey publicKey) {
        byte[] src65 = publicKey.getQ().getEncoded(false);
        byte[] rawXY = new byte[CURVE_LEN * 2];//SM2的话这里应该是64字节
        System.arraycopy(src65, 1, rawXY, 0, rawXY.length);
        return rawXY;
    }

    public static byte[] encrypt(BCECPublicKey pubKey, byte[] srcData) throws InvalidCipherTextException {
        ECPublicKeyParameters pubKeyParameters = BCECUtil.convertPublicKeyToParameters(pubKey);
        return encrypt(pubKeyParameters, srcData);
    }

    /**
     * ECC公钥加密
     *
     * @param pubKeyParameters ECC公钥
     * @param srcData          源数据
     * @return SM2密文，实际包含三部分：ECC公钥、真正的密文、公钥和原文的SM3-HASH值
     * @throws InvalidCipherTextException
     */
    public static byte[] encrypt(ECPublicKeyParameters pubKeyParameters, byte[] srcData)
        throws InvalidCipherTextException {
        SM2Engine engine = new SM2Engine();
        ParametersWithRandom pwr = new ParametersWithRandom(pubKeyParameters, new SecureRandom());
        engine.init(true, pwr);
        return engine.processBlock(srcData, 0, srcData.length);
    }

    public static byte[] decrypt(BCECPrivateKey priKey, byte[] sm2Cipher) throws InvalidCipherTextException {
        ECPrivateKeyParameters priKeyParameters = BCECUtil.convertPrivateKeyToParameters(priKey);
        return decrypt(priKeyParameters, sm2Cipher);
    }

    /**
     * ECC私钥解密
     *
     * @param priKeyParameters ECC私钥
     * @param sm2Cipher        SM2密文，实际包含三部分：ECC公钥、真正的密文、公钥和原文的SM3-HASH值
     * @return 原文
     * @throws InvalidCipherTextException
     */
    public static byte[] decrypt(ECPrivateKeyParameters priKeyParameters, byte[] sm2Cipher)
        throws InvalidCipherTextException {
        SM2Engine engine = new SM2Engine();
        engine.init(false, priKeyParameters);
        return engine.processBlock(sm2Cipher, 0, sm2Cipher.length);
    }

    private static byte[] fixToCurveLengthBytes(byte[] src) {
        if (src.length == CURVE_LEN) {
            return src;
        }

        byte[] result = new byte[CURVE_LEN];
        if (src.length > CURVE_LEN) {
            System.arraycopy(src, src.length - result.length, result, 0, result.length);
        } else {
            System.arraycopy(src, 0, result, result.length - src.length, src.length);
        }
        return result;
    }

}


/**
 * The MIT License (MIT)
 * Copyright (c) 2017 Microsoft Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package com.microsoft.azure.documentdb.bulkimport;

import java.time.Duration;

import com.google.common.base.Preconditions;

class InsertMetrics {

    final long numberOfDocumentsInserted;
    final Duration timeTaken;
    final double requestUnitsConsumed;
    final long numberOfThrottles;

    /**
     * Initializes a new instance of the InsertMetrics class (default constructor).
     */
    public InsertMetrics() {
        this(0, Duration.ZERO, 0, 0);
    }

    /**
     * Initializes a new instance of the InsertMetrics class (instance constructor).
     * @param numberOfDocumentsInserted Number of documents inserted.
     * @param timeTaken Amount of time taken to insert the documents.
     * @param requestUnitsConsumed The request units consumed to insert the documents.
     * @param numberOfThrottles The number of throttles encountered to insert the documents.
     */
    public InsertMetrics(long numberOfDocumentsInserted, Duration timeTaken, double requestUnitsConsumed, long numberOfThrottles) {
        Preconditions.checkArgument(numberOfDocumentsInserted >= 0, "numberOfDocumentsInserted must be non negative");
        Preconditions.checkArgument(requestUnitsConsumed >= 0, "requestUnitsConsumed must be non negative");
        Preconditions.checkArgument(numberOfThrottles >= 0, "numberOfThrottles must be non negative");

        this.numberOfDocumentsInserted = numberOfDocumentsInserted;
        this.timeTaken = timeTaken;
        this.requestUnitsConsumed = requestUnitsConsumed;
        this.numberOfThrottles = numberOfThrottles;
    }

    /**
     * Sums two {@link InsertMetrics} instances
     * @param m1
     * @param m2
     * @return the sum aggregate result
     */
    public static InsertMetrics sum(InsertMetrics m1, InsertMetrics m2) {
        long totalDocsInserted = m1.numberOfDocumentsInserted + m2.numberOfDocumentsInserted;
        Duration totalTimeTaken = m1.timeTaken.plus(m2.timeTaken);
        double totalRequestUnitsConsumed = m1.requestUnitsConsumed + m2.requestUnitsConsumed;
        long totalNumberOfThrottles = m1.numberOfThrottles + m2.numberOfThrottles;

        return new InsertMetrics(totalDocsInserted, totalTimeTaken, totalRequestUnitsConsumed, totalNumberOfThrottles);
    }
}


/**
 * The MIT License (MIT)
 * Copyright (c) 2017 Microsoft Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package com.microsoft.azure.documentdb.bulkimport;

import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Preconditions;
import com.google.common.base.Stopwatch;
import com.google.common.util.concurrent.AsyncCallable;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.Futures.FutureCombiner;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.microsoft.azure.documentdb.DocumentClient;
import com.microsoft.azure.documentdb.DocumentClientException;
import com.microsoft.azure.documentdb.DocumentCollection;
import com.microsoft.azure.documentdb.FeedOptions;
import com.microsoft.azure.documentdb.PartitionKeyDefinition;
import com.microsoft.azure.documentdb.PartitionKeyRange;
import com.microsoft.azure.documentdb.RetryOptions;
import com.microsoft.azure.documentdb.internal.HttpConstants;
import com.microsoft.azure.documentdb.internal.routing.CollectionRoutingMap;
import com.microsoft.azure.documentdb.internal.routing.InMemoryCollectionRoutingMap;
import com.microsoft.azure.documentdb.internal.routing.PartitionKeyInternal;
import com.microsoft.azure.documentdb.internal.routing.Range;

public class DocumentBulkImporter implements AutoCloseable {

    public static class Builder {

        private DocumentClient client;
        private String collectionLink;
        private int maxMiniBatchSize = (int) Math.floor(MAX_BULK_IMPORT_SCRIPT_INPUT_SIZE * FRACTION_OF_MAX_BULK_IMPORT_SCRIPT_INPUT_SIZE_ALLOWED);
        private final static int DEFAULT_RETRY_ATTEMPT_ON_THROTTLING_FOR_INIT = 200;
        private final static int DEFAULT_WAIT_TIME_ON_THROTTLING_FOR_INIT_IN_SECONDS = 60;

        private PartitionKeyDefinition partitionKeyDef;
        private int offerThroughput;

        private static RetryOptions DEFAULT_INIT_RETRY_OPTIONS;

        static {
            DEFAULT_INIT_RETRY_OPTIONS = new RetryOptions();
            DEFAULT_INIT_RETRY_OPTIONS.setMaxRetryAttemptsOnThrottledRequests(DEFAULT_RETRY_ATTEMPT_ON_THROTTLING_FOR_INIT);
            DEFAULT_INIT_RETRY_OPTIONS.setMaxRetryWaitTimeInSeconds(DEFAULT_WAIT_TIME_ON_THROTTLING_FOR_INIT_IN_SECONDS);
        }

        private RetryOptions retryOptions = DEFAULT_INIT_RETRY_OPTIONS;

        /**
         * Use the instance of {@link DocumentClient} to bulk import to the given instance of {@link DocumentCollection}
         * @param client an instance of {@link DocumentClient}
         * @param partitionKeyDef specifies the {@link PartitionKeyDefinition} of the collection
         * @param databaseName name of the database
         * @param collectionName name of the collection
         * @param offerThroughput specifies the collection throughput
         * @return an instance of {@link Builder}
         */
        public Builder from(DocumentClient client,
                String databaseName,
                String collectionName,
                PartitionKeyDefinition partitionKeyDef,
                int offerThroughput) {

            // TODO: validate the retry options for the client
            this.client = client;
            this.collectionLink = String.format("/dbs/%s/colls/%s", databaseName, collectionName);
            this.partitionKeyDef = partitionKeyDef;
            this.offerThroughput = offerThroughput;
            return this;
        }

        /**
         * use the given size to configure max mini batch size.
         *
         * If not specified will use the default.
         * @param size specifies the size of mini batch.
         * @return {@link Builder}
         */
        public Builder withMaxMiniBatchSize(int size) {
            Preconditions.checkArgument(size > 0, "maxMiniBatchSize cannot be negative");
            Preconditions.checkArgument(size <= MAX_BULK_IMPORT_SCRIPT_INPUT_SIZE, "maxMiniBatchSize cannot be negative");

            this.maxMiniBatchSize = size;
            return this;
        }

        /**
         * use the given retry option for initialization
         *
         * @param options an instance of {@link RetryOptions}
         * @return {@link Builder}
         */
        public Builder withInitializationRetryOptions(RetryOptions options) {
            this.retryOptions = options;
            return this;
        }

        /**
         * Instantiates {@link DocumentBulkImporter} given the configured {@link Builder}.
         *
         * @return the new builder
         * @throws Exception if there is any failure
         */
        public DocumentBulkImporter build() throws Exception {
            DocumentBulkImporter importer = new DocumentBulkImporter(client, collectionLink, partitionKeyDef, offerThroughput);
            try {
                importer.setInitializationRetryOptions(retryOptions);
                importer.setMaxMiniBatchSize(maxMiniBatchSize);

                importer.safeInit();

            } catch (Exception e) {
                importer.close();
                throw e;
            }
            return importer;
        }

        private Builder() {}
    }

    /**
     * Creates a new {@link DocumentBulkImporter.Builder} instance
     * @return an instance of {@link DocumentBulkImporter.Builder}
     */
    public static DocumentBulkImporter.Builder builder() {
        return new DocumentBulkImporter.Builder();
    }

    /**
     * The name of the system stored procedure for bulk import.
     */
    private final static String BULK_IMPORT_STORED_PROCECURE_NAME = "__.sys.commonBulkInsert";

    /**
     * The maximal sproc payload size sent (as a fraction of 2MB).
     */
    private final static int MAX_BULK_IMPORT_SCRIPT_INPUT_SIZE = (2202010 * 5) / 10;

    /**
     * The fraction of maximum sproc payload size upto which documents allowed to be fit in a mini-batch.
     */
    private final static double FRACTION_OF_MAX_BULK_IMPORT_SCRIPT_INPUT_SIZE_ALLOWED = 0.20;

    /**
     * Initialization sleep time on
     */
    private final static int INITIALIZATION_SLEEP_TIME_ON_THROTTLING = 500;

    /**
     * Logger
     */
    private final Logger logger = LoggerFactory.getLogger(DocumentBulkImporter.class);

    /**
     * Degree of parallelism for each partition which was inferred from previous batch execution.
     */
    private final Map<String, Integer> partitionKeyRangeIdToInferredDegreeOfParallelism = new ConcurrentHashMap<>();

    /**
     * Executor Service
     */
    private final ListeningExecutorService listeningExecutorService;

    /**
     * The DocumentDB client instance.
     */
    private final DocumentClient client;

    /**
     * The document collection to which documents are to be bulk imported.
     */
    private final String collectionLink;

    /**
     * Partition Key Definition of the underlying collection.
     */
    private final PartitionKeyDefinition partitionKeyDefinition;

    /**
     * Partition Key Range Ids
     */
    private List<String> partitionKeyRangeIds;

    /**
     * Collection routing map used to retrieve partition key range Ids of a given collection
     */
    private CollectionRoutingMap collectionRoutingMap;

    /**
     * Bulk Import Stored Procedure Link relevant to the given collection
     */
    private String bulkImportStoredProcLink;

    /**
     * Collection offer throughput
     */
    private int collectionThroughput;

    /**
     * Max Mini Batch Size
     */
    private int maxMiniBatchSize;

    private RetryOptions retryOptions;

    private void setMaxMiniBatchSize(int size) {
        this.maxMiniBatchSize = size;
    }

    private void setInitializationRetryOptions(RetryOptions options) {
        this.retryOptions = options;
    }

    /**
     * Initializes a new instance of {@link DocumentBulkImporter}
     *
     * @param client {@link DocumentClient} instance to use
     * @param collectionLink inserts documents to the specified collection
     * @param maxMiniBatchSize
     * @throws DocumentClientException if any failure
     */
    private DocumentBulkImporter(DocumentClient client,
            String collectionLink,
            PartitionKeyDefinition partitionKeyDefinition,
            int collectionOfferThroughput) {
        Preconditions.checkNotNull(client, "client cannot be null");
        Preconditions.checkNotNull(partitionKeyDefinition, "partitionKeyDefinition cannot be null");
        Preconditions.checkNotNull(collectionLink, "collectionLink cannot be null");
        Preconditions.checkArgument(collectionOfferThroughput > 0, "collection throughput is less than 10,000");

        this.client = client;
        this.collectionLink = collectionLink;
        this.collectionThroughput =  collectionOfferThroughput;
        this.partitionKeyDefinition = partitionKeyDefinition;
        this.listeningExecutorService = MoreExecutors.listeningDecorator(Executors.newCachedThreadPool());
    }

    private void safeInit() throws Exception {
        int count = 0;
        long startTime = System.currentTimeMillis();
        while(true) {
            try {
                initialize();
                break;
            } catch (Exception e) {
                count++;
                DocumentClientException dce = ExceptionUtils.getThrottelingException(e);
                long now = System.currentTimeMillis();
                if (count < retryOptions.getMaxRetryAttemptsOnThrottledRequests()
                        && now - startTime < (retryOptions.getMaxRetryWaitTimeInSeconds() * 1000)
                        && dce != null
                        && dce.getStatusCode() == HttpConstants.StatusCodes.TOO_MANY_REQUESTS ) {
                    Thread.sleep(count * dce.getRetryAfterInMilliseconds() + INITIALIZATION_SLEEP_TIME_ON_THROTTLING);
                    continue;
                } else {
                    throw e;
                }
            }
        }
    }

    /**
     * Releases any internal resources.
     * It is responsibility of the caller to close {@link DocumentClient}.
     */
    @Override
    public void close() {
        // disable submission of new tasks
        listeningExecutorService.shutdown();
        try {
            // wait for existing tasks to terminate
            if (!listeningExecutorService.awaitTermination(60, TimeUnit.SECONDS)) {
                // cancel any currently running executing tasks
                listeningExecutorService.shutdownNow();
                // wait for cancelled tasks to terminate
                if (!listeningExecutorService.awaitTermination(60, TimeUnit.SECONDS)) {
                    logger.error("some tasks did not terminate");
                }
            }
        } catch (InterruptedException e) {
            listeningExecutorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    /**
     * Initializes {@link DocumentBulkImporter}. This happens only once
     * @throws DocumentClientException
     */
    private void initialize() throws DocumentClientException {
        logger.debug("Initializing ...");

        this.bulkImportStoredProcLink = String.format("%s/sprocs/%s", collectionLink, BULK_IMPORT_STORED_PROCECURE_NAME);

        logger.trace("Fetching partition map of collection");
        Range<String> fullRange = new Range<String>(
                PartitionKeyInternal.MinimumInclusiveEffectivePartitionKey,
                PartitionKeyInternal.MaximumExclusiveEffectivePartitionKey,
                true,
                false);

        this.collectionRoutingMap = getCollectionRoutingMap(client);
        Collection<PartitionKeyRange> partitionKeyRanges = this.collectionRoutingMap.getOverlappingRanges(fullRange);

        this.partitionKeyRangeIds = partitionKeyRanges.stream().map(partitionKeyRange -> partitionKeyRange.getId()).collect(Collectors.toList());

        logger.debug("Initialization completed");
    }

    /**
     * Executes a bulk import in the Azure Cosmos DB database service.
     *

     *

     * {@code
     * ConnectionPolicy connectionPolicy = new ConnectionPolicy();
     * RetryOptions retryOptions = new RetryOptions();
     * // set to 0 to let bulk importer handles throttling
     * retryOptions.setMaxRetryAttemptsOnThrottledRequests(0);
     * connectionPolicy.setRetryOptions(retryOptions);
     * connectionPolicy.setMaxPoolSize(200);
     *
     * DocumentClient client = new DocumentClient(HOST, MASTER_KEY, connectionPolicy, null);
     *
     * String collectionLink = String.format("/dbs/%s/colls/%s", "mydb", "mycol");
     * DocumentCollection collection = client.readCollection(collectionLink, null).getResource();
     *
     * DocumentBulkImporter importer = DocumentBulkImporter.builder().from(client, collection).build();
     *
     * for(int i = 0; i < 10; i++) {
     *   List<String> documents = documentSource.getMoreDocuments();
     *
     *   BulkImportResponse bulkImportResponse = importer.importAll(documents, false);
     *
     *   //validate that all documents inserted to ensure no failure.
     *   // bulkImportResponse.getNumberOfDocumentsImported() == documents.size()
     *   if (bulkImportResponse.getNumberOfDocumentsImported() < documents.size()) {
     *      for(Exception e: bulkImportResponse.getFailuresIfAny()) {
     *          // validate why there were some failures
     *          // in case if you decide to re-import these batch of documents
     *          // (as some of them may already have inserted)
     *          // you should import them with upsert enable option
     *          e.printStackTrace();
     *      }
     *      break;
     *   }
     * }
     *
     * importer.close();
     * client.close();
     * }
     *

     *

     *
     * @param documents to insert
     * @param isUpsert whether enable upsert (overwrite if it exists)
     * @return an instance of {@link BulkImportResponse}
     * @throws DocumentClientException if any failure happens
     */
    public BulkImportResponse importAll(Collection<String> documents, boolean isUpsert) throws DocumentClientException {
        return executeBulkImportInternal(documents,
                isUpsert);
    }

    private BulkImportResponse executeBulkImportInternal(Collection<String> input,
            boolean isUpsert) throws DocumentClientException {
        Preconditions.checkNotNull(input, "document collection cannot be null");
        try {
            return executeBulkImportAsyncImpl(input, isUpsert).get();

        } catch (ExecutionException e) {
            logger.debug("Failed to import documents", e);
            Throwable cause = e.getCause();
            if (cause instanceof Exception) {
                throw toDocumentClientException((Exception) cause);
            } else {
                throw toDocumentClientException(e);
            }
        } catch(Exception e) {
            logger.error("Failed to import documents", e);
            throw toDocumentClientException(e);
        }
    }

    private DocumentClientException toDocumentClientException(Exception e) {
        if (e instanceof DocumentClientException) {
            return (DocumentClientException) e;
        } else {
            return new DocumentClientException(500, e);
        }
    }

    private int getDocumentSizeOrThrow(String document) {
        int documentSize = document.getBytes(Charset.forName("UTF-8")).length;
        if (documentSize > maxMiniBatchSize) {
            logger.error("Document size {} larger than script payload limit. {}", documentSize, maxMiniBatchSize);
            throw new UnsupportedOperationException("Cannot import a document whose size is larger than script payload limit.");
        }
        return documentSize;
    }

    private ListenableFuture<BulkImportResponse> executeBulkImportAsyncImpl(Collection<String> documents,
            boolean isUpsert) throws Exception {
        Stopwatch watch = Stopwatch.createStarted();

        BulkImportStoredProcedureOptions options = new BulkImportStoredProcedureOptions(true, true, null, false, isUpsert);

        logger.debug("Bucketing documents ...");

        ConcurrentHashMap<String, Set<String>> documentsToImportByPartition = new ConcurrentHashMap<String, Set<String>>();
        ConcurrentHashMap<String, List<List<String>>> miniBatchesToImportByPartition = new ConcurrentHashMap<String, List<List<String>>>();

        for (String partitionKeyRangeId: partitionKeyRangeIds) {
            documentsToImportByPartition.put(partitionKeyRangeId,  ConcurrentHashMap.newKeySet(documents.size() / partitionKeyRangeIds.size()));
            miniBatchesToImportByPartition.put(partitionKeyRangeId, new ArrayList<List<String>>(1000));
        }

        documents.parallelStream().forEach(documentAsString -> {
            PartitionKeyInternal partitionKeyValue = DocumentAnalyzer.extractPartitionKeyValue(documentAsString, partitionKeyDefinition);
            String effectivePartitionKey = partitionKeyValue.getEffectivePartitionKeyString(partitionKeyDefinition, true);
            String partitionRangeId = collectionRoutingMap.getRangeByEffectivePartitionKey(effectivePartitionKey).getId();
            documentsToImportByPartition.get(partitionRangeId).add(documentAsString);
        });

        logger.trace("Creating mini batches within each partition bucket");

        documentsToImportByPartition.entrySet().parallelStream().forEach(entry -> {

            String partitionRangeId = entry.getKey();

            Set<String> documentsToImportInPartition =  entry.getValue();

            Iterator<String> it = documentsToImportInPartition.iterator();
            ArrayList<String> currentMiniBatch = new ArrayList<String>(500);
            int currentMiniBatchSize = 0;

            while (it.hasNext()) {
                String currentDocument = it.next();
                int currentDocumentSize = getDocumentSizeOrThrow(currentDocument);

                if ((currentMiniBatchSize + currentDocumentSize <= maxMiniBatchSize)) {
                    // add the document to current batch
                    currentMiniBatch.add(currentDocument);
                    currentMiniBatchSize += currentDocumentSize;
                } else {
                    // this batch has reached its max size
                    miniBatchesToImportByPartition.get(partitionRangeId).add(currentMiniBatch);
                    currentMiniBatch = new ArrayList<String>(500);
                    currentMiniBatch.add(currentDocument);
                    currentMiniBatchSize = currentDocumentSize;
                }
            }

            if (currentMiniBatch.size() > 0) {
                // add mini batch
                miniBatchesToImportByPartition.get(partitionRangeId).add(currentMiniBatch);
            }
        });

        logger.debug("Beginning bulk import within each partition bucket");
        Map<String, BatchInserter> batchInserters = new HashMap<String, BatchInserter>();
        Map<String, CongestionController> congestionControllers = new HashMap<String, CongestionController>();

        logger.debug("Preprocessing took: " + watch.elapsed().toMillis() + " millis");
        List<ListenableFuture<Void>> futures = new ArrayList<>();

        for (String partitionKeyRangeId: this.partitionKeyRangeIds) {

            BatchInserter batchInserter = new BatchInserter(
                    partitionKeyRangeId,
                    miniBatchesToImportByPartition.get(partitionKeyRangeId),
                    this.client,
                    bulkImportStoredProcLink,
                    options);
            batchInserters.put(partitionKeyRangeId, batchInserter);

            CongestionController cc = new CongestionController(listeningExecutorService,
                    collectionThroughput / partitionKeyRangeIds.size(),
                    partitionKeyRangeId,
                    batchInserter,
                    partitionKeyRangeIdToInferredDegreeOfParallelism.get(partitionKeyRangeId));

            congestionControllers.put(partitionKeyRangeId,cc);

            // starting
            futures.add(cc.executeAllAsync());
        }

        FutureCombiner<Void> futureContainer = Futures.whenAllComplete(futures);
        AsyncCallable<BulkImportResponse> completeAsyncCallback = new AsyncCallable<BulkImportResponse>() {

            @Override
            public ListenableFuture<BulkImportResponse> call() throws Exception {

                List<Exception> failures = new ArrayList<>();

                for(String partitionKeyRangeId: partitionKeyRangeIds) {
                    CongestionController cc = congestionControllers.get(partitionKeyRangeId);
                    failures.addAll(cc.getFailures());
                    partitionKeyRangeIdToInferredDegreeOfParallelism.put(partitionKeyRangeId, cc.getDegreeOfConcurrency());
                }

                int numberOfDocumentsImported = batchInserters.values().stream().mapToInt(b -> b.getNumberOfDocumentsImported()).sum();
                double totalRequestUnitsConsumed = batchInserters.values().stream().mapToDouble(b -> b.getTotalRequestUnitsConsumed()).sum();

                watch.stop();

                BulkImportResponse bulkImportResponse = new
                        BulkImportResponse(numberOfDocumentsImported, totalRequestUnitsConsumed, watch.elapsed(), failures);

                return Futures.immediateFuture(bulkImportResponse);
            }
        };

        return futureContainer.callAsync(completeAsyncCallback, listeningExecutorService);
    }

    private CollectionRoutingMap getCollectionRoutingMap(DocumentClient client) {
        List<ImmutablePair<PartitionKeyRange, Boolean>> ranges = new ArrayList<>();

        for (PartitionKeyRange range : client.readPartitionKeyRanges(this.collectionLink, (FeedOptions) null).getQueryIterable().toList()) {
            ranges.add(new ImmutablePair<>(range, true));
        }

        CollectionRoutingMap routingMap = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(ranges,
                StringUtils.EMPTY);

        if (routingMap == null) {
            throw new IllegalStateException("Cannot create complete routing map");
        }

        return routingMap;
    }
}


/**
 * The MIT License (MIT)
 * Copyright (c) 2017 Microsoft Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package com.microsoft.azure.documentdb.bulkimport;

import java.util.Collections;
import java.util.Iterator;

import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.microsoft.azure.documentdb.PartitionKeyDefinition;
import com.microsoft.azure.documentdb.Undefined;
import com.microsoft.azure.documentdb.internal.routing.PartitionKeyInternal;

class DocumentAnalyzer {
    private final static ObjectMapper objectMapper = new ObjectMapper();
    private final static Logger LOGGER = LoggerFactory.getLogger(DocumentAnalyzer.class);
    /**
     * Extracts effective {@link PartitionKeyInternal} from serialized document.
     * @param documentAsString Serialized document to extract partition key value from.
     * @param partitionKeyDefinition Information about partition key.
     * @return PartitionKeyInternal
     */
    public static PartitionKeyInternal extractPartitionKeyValue(String documentAsString,
            PartitionKeyDefinition partitionKeyDefinition)  {

        if (partitionKeyDefinition == null || partitionKeyDefinition.getPaths().size() == 0) {
            return PartitionKeyInternal.getEmpty();
        }

        return DocumentAnalyzer.extractPartitionKeyValueInternal(documentAsString, partitionKeyDefinition);
    }

    private static PartitionKeyInternal extractPartitionKeyValueInternal(String documentAsString, PartitionKeyDefinition partitionKeyDefinition) {
        JsonNode root;
        try {
            root = objectMapper.readTree(documentAsString);

            Iterator<String> path = partitionKeyDefinition.getPaths().iterator();
            JsonNode node =  root.path(path.next().substring(1));

            while(path.hasNext() && node != null) {
                node = node.path(path.next());
            }

            Object partitionKeyValue = null;

            if (node != null) {

                switch (node.getNodeType()) {
                case BOOLEAN:
                    partitionKeyValue = node.booleanValue();
                    break;
                case MISSING:
                    partitionKeyValue = Undefined.Value();
                    break;
                case NULL:
                    partitionKeyValue = JSONObject.NULL;
                    break;
                case NUMBER:
                    partitionKeyValue = node.numberValue();
                    break;
                case STRING:
                    partitionKeyValue = node.textValue();
                    break;
                default:
                    throw new RuntimeException(String.format("undefined json type %s", node.getNodeType()));
                }
            } else {
                partitionKeyValue = Undefined.Value();
            }

            return fromPartitionKeyvalue(partitionKeyValue);

        } catch (Exception e) {
            LOGGER.error("Failed to extract partition key value from document {}", documentAsString, e);
            throw ExceptionUtils.toRuntimeException(e);
        }
    }

    public static PartitionKeyInternal fromPartitionKeyvalue(Object partitionKeyValue) {
        try {
            return PartitionKeyInternal.fromObjectArray(Collections.singletonList(partitionKeyValue), true);
        } catch (Exception e) {
            LOGGER.error("Failed to instantiate ParitionKeyInternal from {}", partitionKeyValue, e);
            throw ExceptionUtils.toRuntimeException(e);
        }
    }
}


/**
 * Copyright (c) 2011, Yubico AB.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following
 *   disclaimer in the documentation and/or other materials provided
 *   with the distribution.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
 *  CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 *  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
 *  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 *  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 *  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 *  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
 *  TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
 *  THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 *  SUCH DAMAGE.
 */
package com.yubico.jaas;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import javax.security.auth.Subject;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;
import javax.security.jacc.PolicyContext;
import javax.security.jacc.PolicyContextException;
import javax.servlet.http.HttpServletRequest;

import com.yubico.client.v2.ResponseStatus;
import com.yubico.client.v2.VerificationResponse;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.yubico.client.v2.YubicoClient;
import com.yubico.client.v2.exceptions.YubicoVerificationException;
import com.yubico.client.v2.exceptions.YubicoValidationFailure;

/**
 * A JAAS module for verifying OTPs (One Time Passwords) against a
 * Yubikey Validation Service.
 *
 * @author Fredrik Thulin (fredrik@yubico.com)
 *
 */
public class YubikeyLoginModule implements LoginModule {
	/* Options for this class.
	 * Note that the options map is shared with other classes, like YubikeyToUserMap.
	 */
	public static final String OPTION_YUBICO_CLIENT_ID		= "clientId";
	public static final String OPTION_YUBICO_CLIENT_KEY         	= "clientKey";
	public static final String OPTION_YUBICO_ID_REALM		= "id_realm";
	public static final String OPTION_YUBICO_SOFT_FAIL_NO_OTPS	= "soft_fail_on_no_otps";
	public static final String OPTION_YUBICO_WSAPI_URLS		= "wsapi_urls";
	public static final String OPTION_YUBICO_USERMAP_CLASS  	= "usermap_class";
	public static final String OPTION_YUBICO_SYNC_POLICY        	= "sync_policy";
	public static final String OPTION_YUBICO_JACC		 	= "jacc";
	public static final String JACC_ATTR_WEB_REQUEST_KEY 		= "javax.servlet.http.HttpServletRequest";
	public static final String HTTP_REQUEST_ATTR_TOTP 		= "j_otp";

	/* JAAS stuff */
	private Subject subject;
	private CallbackHandler callbackHandler;

	/* YubicoClient settings */
	private YubicoClient yc;

	private boolean soft_fail_on_no_otps;
	private boolean jacc;

	private YubikeyToUserMap ykmap;
	private String idRealm;

	private final Logger log = LoggerFactory.getLogger(YubikeyLoginModule.class);

	private ArrayList<YubikeyPrincipal> principals = new ArrayList<YubikeyPrincipal>();

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#abort()
	 */
	public boolean abort() throws LoginException {
		log.trace("In abort()");
		for (YubikeyPrincipal p : this.principals) {
			this.subject.getPrincipals().remove(p);
		}
		return true;
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#commit()
	 */
	public boolean commit() throws LoginException {
		log.trace("In commit()");
		for (YubikeyPrincipal p : this.principals) {
			log.debug("Committing principal {}", p);
			this.subject.getPrincipals().add(p);
		}
		return true;
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#logout()
	 */
	public boolean logout() throws LoginException {
		log.trace("In logout()");
		for (YubikeyPrincipal p : this.principals) {
			this.subject.getPrincipals().remove(p);
		}
		return false;
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#initialize(javax.security.auth.Subject, javax.security.auth.callback.CallbackHandler, java.util.Map, java.util.Map)
	 */
	public void initialize(Subject newSubject, CallbackHandler newCallbackHandler,
			Map<String, ?> sharedState, Map<String, ?> options) {

		log.debug("Initializing YubikeyLoginModule");
		this.subject = newSubject;
		this.callbackHandler = newCallbackHandler;

		/* Yubico verification client */
		Integer clientId = Integer.parseInt(options.get(OPTION_YUBICO_CLIENT_ID).toString());
        String clientKey = options.get(OPTION_YUBICO_CLIENT_KEY).toString();
		this.yc = YubicoClient.getClient(clientId, clientKey);

		/* Realm of principals added after authentication */
		if (options.containsKey(OPTION_YUBICO_ID_REALM)) {
			this.idRealm = options.get(OPTION_YUBICO_ID_REALM).toString();
		}

		/* Should this JAAS module be ignored when no OTPs are supplied? */
		if (options.containsKey(OPTION_YUBICO_SOFT_FAIL_NO_OTPS)) {
			if ("true".equals(options.get(OPTION_YUBICO_SOFT_FAIL_NO_OTPS).toString())) {
				this.soft_fail_on_no_otps = true;
			}
		}

		/* Should this JAAS module uses j_otp form to pick up otp */
		if (options.containsKey(OPTION_YUBICO_JACC)) {
			if ("true".equals(options.get(OPTION_YUBICO_JACC).toString())) {
				this.jacc = true;
			}
		}

		/* User-provided URLs to the Yubico validation service, separated by "|". */
		if (options.containsKey(OPTION_YUBICO_WSAPI_URLS)) {
			String in = options.get(OPTION_YUBICO_WSAPI_URLS).toString();
			String l[] = in.split("\\|");
			this.yc.setWsapiUrls(l);
		}

		if (options.containsKey(OPTION_YUBICO_SYNC_POLICY)) {
			this.yc.setSync(Integer.parseInt(options.get(OPTION_YUBICO_SYNC_POLICY).toString()));
		}

		/* Instantiate the specified usermap implementation. */
                String usermap_class_name = null;
                if (options.containsKey(OPTION_YUBICO_USERMAP_CLASS)) {
                    usermap_class_name = options.get(OPTION_YUBICO_USERMAP_CLASS).toString();
                } else {
                    usermap_class_name = "com.yubico.jaas.impl.YubikeyToUserMapImpl"; // Default implementation
                }
                try {
                    log.debug("Trying to instantiate {}",usermap_class_name);
                    this.ykmap = (YubikeyToUserMap) Class.forName(usermap_class_name).newInstance();
                    this.ykmap.setOptions(options);
                } catch (ClassNotFoundException ex) {
                    log.error("Could not create usermap from class " + usermap_class_name, ex);
                } catch (InstantiationException ex) {
                    log.error("Could not create usermap from class " + usermap_class_name, ex);
                } catch (IllegalAccessException ex) {
                    log.error("Could not create usermap from class " + usermap_class_name, ex);
                }
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#login()
	 */
	public boolean login() throws LoginException {
		NameCallback nameCb = new NameCallback("Enter username: ");

		log.debug("Begin OTP login");

		if (callbackHandler == null) {
			throw new LoginException("No callback handler available in login()");
		}

		List<String> otps = get_tokens(nameCb);
		if (otps.size() == 0) {
			if (this.soft_fail_on_no_otps) {

				log.debug("No OTPs found, and soft-fail is on. Making JAAS ignore this module.");
				return false;
			}
			throw new LoginException("YubiKey OTP authentication failed - no OTPs supplied");
		}

		if (validate_otps(otps, nameCb)) {
			return true;
		}

		log.info("None out of {} possible YubiKey OTPs for user {} validated successful",
					otps.size(), nameCb.getName());

		throw new LoginException("YubiKey OTP authentication failed");
	}

	/**
	 * Try to validate all the OTPs provided.
	 * @param otps  Possible YubiKey OTPs
	 * @param nameCb  JAAS callback to get authenticating username
	 * @return  true if one or more of the OTPs validated OK, otherwise false
	 * @throws LoginException for exceptions during validation, logging real exception at warn level
	 */
	private boolean validate_otps(List<String> otps, NameCallback nameCb) throws LoginException {
		boolean validated = false;

		for (String otp : otps) {
			log.trace("Checking OTP {}", otp);

			VerificationResponse ykr;
			try {
				ykr = this.yc.verify(otp);
			} catch (YubicoVerificationException e) {
				log.warn("Errors during validation: ", e);
				throw new LoginException("Errors during validation: " + e.getMessage());
			} catch (YubicoValidationFailure e) {
				log.warn("Something went very wrong during authentication: ", e);
				throw new LoginException("Something went very wrong during authentication: " + e.getMessage());
			}
			if (ykr != null) {
				log.trace("OTP {} verify result : {}", otp, ykr.getStatus().toString());
				if (ykr.getStatus() == ResponseStatus.OK) {
					String publicId = YubicoClient.getPublicId(otp);
					log.info("OTP verified successfully (YubiKey id {})", publicId);
					if (is_right_user(nameCb.getName(), publicId)) {
						this.principals.add(new YubikeyPrincipal(publicId, this.idRealm));
						/* Don't just return here, we want to "consume" all OTPs if
						 * more than one is provided.
						 */
						validated = true;
					}
				} else {
					log.debug("OTP validation returned {}", ykr.getStatus().toString());
				}
			}
		}

		return validated;
	}

	/**
	 * After validation of an OTP, check that it came from a YubiKey that actually
	 * belongs to the user trying to authenticate.
	 *
	 * @param username  Username to match against YubiKey publicId.
	 * @param publicId  The public ID of the authenticated YubiKey.
	 * @return true if the username matched the YubiKey, false otherwise
	 */
	private boolean is_right_user(String username, String publicId) {
		log.debug("Check if YubiKey {} belongs to user {}", publicId, username);
		return this.ykmap.is_right_user(username, publicId);
	}

	/**
	 * Get username and token(s) from the application, using the
	 * javax.security.auth.callback.CallbackHandler passed to our initialize()
	 * function.
	 *
	 * The tokens returned have been identified as plausible YubiKey OTPs.
	 *
	 * @param nameCb
	 * @return list of possible YubiKey OTPs
	 */
	private List<String> get_tokens(NameCallback nameCb) {
		MultiValuePasswordCallback mv_passCb = new MultiValuePasswordCallback("Enter authentication tokens: ", false);
		List<String> result = new ArrayList<String>();

		try {
			/* Fetch a password using the callbackHandler */
			callbackHandler.handle(new Callback[] { nameCb, mv_passCb });

			for (char[] c : mv_passCb.getSecrets()) {
				String s = new String(c);
				/* Check that OTP is at least 32 chars before we verify it. User might have entered
				 * some other password instead of an OTP, and we don't want to send that, possibly
				 * in clear text, over the network.
				 */
				if (s.length() < 32) {
					log.debug("Skipping token, not a valid YubiKey OTP (too short, {} < 32)", s.length());
				} else {
					result.add(s);
				}
			}
		} catch (UnsupportedCallbackException ex) {
			log.error("Callback type not supported", ex);
		} catch (IOException ex) {
			log.error("CallbackHandler failed", ex);
		}

		if (jacc) {
			// This is JACC specific mechanism
			try {
				HttpServletRequest request = (HttpServletRequest) PolicyContext
						.getContext(JACC_ATTR_WEB_REQUEST_KEY);
				String j_otp = request.getParameter(HTTP_REQUEST_ATTR_TOTP);

				if (j_otp == null || j_otp.length() < 32) {
					log.debug("Skipping token from j_otp, not a valid YubiKey OTP (too short, {} < 32)", j_otp == null ? 0 : j_otp.length());
				} else {
					result.add(j_otp);
				}
				log.debug("OTP from j_otp token : {}", j_otp);
			} catch (PolicyContextException e) {
				log.debug("No OTP from j_otp token)");
			}
		}

		return result;
	}
}


/**
 * Copyright (c) 2011, Yubico AB.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following
 *   disclaimer in the documentation and/or other materials provided
 *   with the distribution.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
 *  CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 *  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
 *  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 *  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 *  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 *  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
 *  TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
 *  THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 *  SUCH DAMAGE.
 */
package com.yubico.jaas.impl;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Map;
import java.util.Scanner;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.yubico.jaas.YubikeyToUserMap;

/**
 * Class to verify that a user is the rightful user of a YubiKey.
 *
 * The current implementation is very rudimentary and uses a plain text file
 * as a backend, expecting lines like "yk.vvcccfhcb.user = alice". If "bob"
 * tries to authenticate with YubiKey "vvcccfhcb" he should be denied.
 *
 * There is also support for auto provisioning (default: false). If the option
 * auto_provision_owner is set to "true", a user logging in without a prior
 * YubiKey association will have the YubiKey they are using associated with them.
 * This should probably only be used during an initial time window when YubiKey
 * tokens are distributed to users, after which the auto provisioning should be
 * turned off again.
 *
 * @author Fredrik Thulin (fredrik@yubico.com)
 *
 */
public class YubikeyToUserMapImpl implements YubikeyToUserMap {
	/* Options for this class */
	public static final String OPTION_YUBICO_AUTO_PROVISION		= "auto_provision_owner";
	public static final String OPTION_YUBICO_ID2NAME_TEXTFILE	= "id2name_textfile";
	public static final String OPTION_YUBICO_VERIFY_YK_OWNER	= "verify_yubikey_owner";
	private String id2name_textfile;
	private boolean auto_provision_owners = false;
	private boolean verify_yubikey_owner = true;
    private final Logger log = LoggerFactory.getLogger(YubikeyToUserMapImpl.class);

    /** {@inheritDoc} */
    public final void setOptions(Map<String, ?> options) {
		/* Is verification of YubiKey owners enabled? */
		this.verify_yubikey_owner = true;
		if (options.get(OPTION_YUBICO_VERIFY_YK_OWNER) != null) {
			if ("false".equals(options.get(OPTION_YUBICO_VERIFY_YK_OWNER).toString())) {
				this.verify_yubikey_owner = false;
			}
		}

		/* id2name text file */
		if (options.get(OPTION_YUBICO_ID2NAME_TEXTFILE) != null) {
			this.id2name_textfile = options.get(OPTION_YUBICO_ID2NAME_TEXTFILE).toString();
		}

		/* should we automatically assign new yubikeys to users? */
		if (options.get(OPTION_YUBICO_AUTO_PROVISION) != null) {
			if ("true".equals(options.get(OPTION_YUBICO_AUTO_PROVISION).toString())) {
				this.auto_provision_owners = true;
			}
		}
	}

	/** {@inheritDoc} */
	public boolean is_right_user(String username, String publicId) {
		if (! this.verify_yubikey_owner) {
			log.debug("YubiKey owner verification disabled, returning 'true'");
			return true;
		}
		if (this.id2name_textfile == null) {
			log.debug("No id2name configuration. Defaulting to {}.", this.verify_yubikey_owner);
			return this.verify_yubikey_owner;
		}

		String ykuser;
		try {
			ykuser = get_username_for_id(publicId, this.id2name_textfile);
		} catch (FileNotFoundException ex) {
			log.error("Yubikey to username textfile {} not found", this.id2name_textfile);
			return false;
		}

		if (ykuser != null) {
			if (! ykuser.equals(username)) {
				log.info("YubiKey " + publicId + " registered to user {}, NOT {}", ykuser, username);
				return false;
			}
			return true;
		} else {
			if (this.auto_provision_owners) {
				log.info("Registering new YubiKey " + publicId + " as belonging to {}", username);

				add_yubikey_to_user(publicId, username, this.id2name_textfile);
				return true;
			}
			log.debug("No record of YubiKey {} found. Returning 'false'.", publicId);
			return false;
		}
	}

	/**
	 * Given publicId "vvcccccfhc", scans filename for a line like "yk.vvcccccfhc.user = alice"
	 * and returns "alice" if found. Null is returned in case there is no matching line in file.
	 *
	 * @param publicId YubiKey public ID to scan for
	 * @param filename name of the file to scan
	 * @return String username
	 * @throws FileNotFoundException if the filename does not match an existing file
	 */
	private String get_username_for_id(String publicId, String filename) throws FileNotFoundException {
		Scanner sc = null;
		File file = new File(filename);
		try {
			sc = new Scanner(file);
			while (sc.hasNextLine()) {
				String line = sc.nextLine();
				if (line.startsWith("yk." + publicId + ".user")) {
					String ykuser = line.split("=")[1].trim();

					return ykuser;
				}
			}
		} finally {
			if (sc != null) {
				sc.close();
			}
		}
		return null;
	}

	/**
	 * Stores an association between username and YubiKey publicId in filename.
	 *
     * @param publicId YubiKey public ID to associate with user
	 * @param username username to associate with a YubiKey public ID
     * @param filename name of the file to store to
	 */
	private void add_yubikey_to_user(String publicId, String username, String filename) {
		try {
			File file = new File(filename);
			FileWriter writer = new FileWriter(file, true);
			writer.write("yk." + publicId + ".user = " + username
					+ System.getProperty("line.separator"));
			writer.close();
		} catch (IOException ex) {
			log.error("Failed appending entry to file {}", filename, ex);
		}
	}
}


/* Copyright (c) 2011, Linus Widströmer.  All rights reserved.

   Copyright (c) 2011, Simon Buckle.  All rights reserved.

   Copyright (c) 2014, Yubico AB.  All rights reseved.

   Redistribution and use in source and binary forms, with or without
   modification, are permitted provided that the following conditions
   are met:

   * Redistributions of source code must retain the above copyright
     notice, this list of conditions and the following disclaimer.

   * Redistributions in binary form must reproduce the above copyright
     notice, this list of conditions and the following
     disclaimer in the documentation and/or other materials provided
     with the distribution.

   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
   CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
   INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
   MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
   DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
   BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
   TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
   DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
   ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
   TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
   THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
   SUCH DAMAGE.

   Written by Linus Widströmer <linus.widstromer@it.su.se>, January 2011.

   Modified by Simon Buckle <simon@webteq.eu>, September 2011.
    - Added support for generating and validating signatures
*/

package com.yubico.client.v2.impl;

import com.yubico.client.v2.Signature;
import com.yubico.client.v2.YubicoClient;
import com.yubico.client.v2.VerificationResponse;
import com.yubico.client.v2.VerificationRequester;
import com.yubico.client.v2.exceptions.YubicoSignatureException;
import com.yubico.client.v2.exceptions.YubicoVerificationException;
import com.yubico.client.v2.exceptions.YubicoValidationFailure;

import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.*;
import java.util.Map.Entry;

import static com.yubico.client.v2.HttpUtils.toQueryString;
import static com.yubico.client.v2.ResponseStatus.BAD_SIGNATURE;

public class YubicoClientImpl extends YubicoClient {
    private final VerificationRequester validationService;

    YubicoClientImpl(VerificationRequester validationService) {
        this.validationService = validationService;
    }

    /**
     * Creates a YubicoClient that will be using the given Client ID.
     *
     * @param clientId Retrieved from https://upgrade.yubico.com/getapikey
     */
    public YubicoClientImpl(Integer clientId) {
        this(new VerificationRequester());
        this.clientId = clientId;
    }

    /**
     * Creates a YubicoClient that will be using the given Client ID and API key.
     *
     * @param clientId Retrieved from https://upgrade.yubico.com/getapikey
     * @param apiKey Retrieved from https://upgrade.yubico.com/getapikey
     */
    public YubicoClientImpl(Integer clientId, String apiKey) {
        this(clientId);
        setKey(apiKey);
    }

    /**
     * Creates a YubicoClient that will be using the given Client ID and API key.
     *
     * @param clientId Retrieved from https://upgrade.yubico.com/getapikey
     * @param apiKey Retrieved from https://upgrade.yubico.com/getapikey
     * @param sync A value 0 to 100 indicating percentage of syncing required by client, or strings "fast" or "secure"
     *             to use server-configured values; if absent, let the server decide
     */
    public YubicoClientImpl(Integer clientId, String apiKey, Integer sync) {
        this(clientId, apiKey);
        setSync(sync);
    }

    /**
     * {@inheritDoc}
     */
    public VerificationResponse verify(String otp) throws YubicoVerificationException, YubicoValidationFailure {
        if (!isValidOTPFormat(otp)) {
            throw new IllegalArgumentException("The OTP is not a valid format");
        }
        Map<String, String> requestMap = new TreeMap<String, String>();
        String nonce = UUID.randomUUID().toString().replaceAll("-", "");
        requestMap.put("nonce", nonce);
        requestMap.put("id", clientId.toString());
        requestMap.put("otp", otp);
        requestMap.put("timestamp", "1");
        if (sync != null) {
            requestMap.put("sl", sync.toString());
        }

        String queryString;
        try {
            queryString = toQueryString(requestMap);
        } catch (UnsupportedEncodingException e) {
            throw new YubicoVerificationException("Failed to encode parameter.", e);
        }

        if (key != null) {
            queryString = sign(queryString);
        }

        String[] wsapiUrls = this.getWsapiUrls();
        List<String> validationUrls = new ArrayList<String>();
        for (String wsapiUrl : wsapiUrls) {
            warnIfDeprecatedUrl(wsapiUrl);
            validationUrls.add(wsapiUrl + "?" + queryString);
        }

        VerificationResponse response = validationService.fetch(validationUrls, userAgent);

        if (key != null) {
            verifySignature(response);
        }

        // NONCE/OTP fields are not returned to the client when sending error codes.
        // If there is an error response, don't need to check them.
        if (!response.getStatus().isError()) {
            if (response.getOtp() == null || !otp.equals(response.getOtp())) {
                throw new YubicoValidationFailure("OTP mismatch in response, is there a man-in-the-middle?");
            }
            if (response.getNonce() == null || !nonce.equals(response.getNonce())) {
                throw new YubicoValidationFailure("Nonce mismatch in response, is there a man-in-the-middle?");
            }
        }
        return response;
    }

    private void verifySignature(VerificationResponse response) throws YubicoValidationFailure, YubicoVerificationException {
        StringBuilder keyValueStr = new StringBuilder();
        for (Entry<String, String> entry : response.getKeyValueMap().entrySet()) {
            if ("h".equals(entry.getKey())) {
                continue;
            }
            if (keyValueStr.length() > 0) {
                keyValueStr.append("&");
            }
            keyValueStr
                    .append(entry.getKey())
                    .append("=")
                    .append(entry.getValue());
        }
        try {
            String signature = Signature.calculate(keyValueStr.toString(), key).trim();
            if (!response.getH().equals(signature) &&
                    !response.getStatus().equals(BAD_SIGNATURE)) {
                // don't throw a ValidationFailure if the server said bad signature, in that
                // case we probably have the wrong key/id and want to check it.
                throw new YubicoValidationFailure("Signatures do not match");
            }
        } catch (YubicoSignatureException e) {
            throw new YubicoVerificationException("Failed to calculate the response signature.", e);
        }
    }

    private String sign(String queryString) throws YubicoVerificationException {
        try {
            queryString += "&h=" + URLEncoder.encode(Signature.calculate(queryString, key), "UTF-8");
        } catch (YubicoSignatureException e) {
            throw new YubicoVerificationException("Failed signing of request", e);
        } catch (UnsupportedEncodingException e) {
            throw new YubicoVerificationException("Failed to encode signature", e);
        }
        return queryString;
    }
}


/*
 * Copyright 2011, Yubico AB. All rights reserved.
 * This file is derivative work from YubikeyToUserMapImpl.java in the
 * Yubico Java client. The following copyright applies to the
 * derivative work:
 * Copyright 2011, Université de Lausanne. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in
 *   the documentation and/or other materials provided with the
 *   distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package com.yubico.jaas.impl;

import com.yubico.jaas.YubikeyToUserMap;
import edu.vt.middleware.ldap.Ldap;
import edu.vt.middleware.ldap.LdapConfig;
import edu.vt.middleware.ldap.SearchFilter;
import java.util.Iterator;
import java.util.Map;
import javax.naming.NamingException;
import javax.naming.directory.Attributes;
import javax.naming.directory.SearchResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Class to verify that a user is the rightful owner of a YubiKey.
 *
 * This implementation uses a LDAP directory to look up the Yubikey's
 * publicId and fetch the associated username.
 *
 * @author Etienne Dysli (etienne.dysli@unil.ch)
 */
public class YubikeyToUserLDAPMap implements YubikeyToUserMap {
    /* Supported JAAS configuration options */
    /** Enable verification of Yubikey owner (default: true) */
    public static final String OPTION_YUBICO_VERIFY_YK_OWNER = "verify_yubikey_owner";
    /** Name of the LDAP attribute containing the Yubikey publicId (default: empty) */
    public static final String OPTION_YUBICO_LDAP_PUBLICID_ATTRIBUTE = "ldap_publicid_attribute";
    /** Name of the LDAP attribute containing the username (default: "uid") */
    public static final String OPTION_YUBICO_LDAP_USERNAME_ATTRIBUTE = "ldap_username_attribute";
    /** URL of the LDAP directory (default: empty) */
    public static final String OPTION_LDAP_URL = "ldap_url";
    /** Base DN for LDAP searches (default: empty) */
    public static final String OPTION_LDAP_BASE_DN = "ldap_base_dn";
    /** DN used to log into the LDAP directory (default: empty) */
    public static final String OPTION_LDAP_BIND_DN = "ldap_bind_dn";
    /** Password for the bind DN (default: empty) */
    public static final String OPTION_LDAP_BIND_CREDENTIAL = "ldap_bind_credential";
    private boolean verify_yubikey_owner = true;
    private String publicid_attribute = "";
    private String username_attribute = "uid";
    private String ldap_url = "";
    private String ldap_base_dn = "";
    private String ldap_bind_dn = "";
    private String ldap_bind_credential = "";
    private Ldap ldap;
    private final Logger log = LoggerFactory.getLogger(YubikeyToUserLDAPMap.class);

    /** {@inheritDoc} */
    public boolean is_right_user(String username, String publicId) {
        log.trace("In is_right_user()");
        if (!this.verify_yubikey_owner) {
            log.debug("YubiKey owner verification disabled, returning 'true'");
            return true;
        }
        String ykuser = null;
        try {
            SearchFilter filter = new SearchFilter("({0}={1})", new String[]{this.publicid_attribute, publicId});
            log.debug("Searching for YubiKey publicId with filter: {}", filter.toString());
            Iterator<SearchResult> results = ldap.search(filter, new String[]{this.username_attribute});
            if (results.hasNext()) {
                Attributes results_attributes = results.next().getAttributes();
                log.debug("Found attributes: {}", results_attributes.toString());
                ykuser = results_attributes.get(this.username_attribute).get().toString();
            } else {
                log.debug("No search results");
            }
        } catch (NamingException ex) {
            log.error(ex.getMessage(), ex);
            return false;
        }
        if (ykuser != null) {
            if (!ykuser.equals(username)) {
                log.info("YubiKey " + publicId + " registered to user {}, NOT {}", ykuser, username);
                return false;
            } else {
                log.info("YubiKey " + publicId + " registered to user {}", ykuser);
                return true;
            }
        } else {
            log.info("No record of YubiKey {} found. Returning 'false'.", publicId);
            return false;
        }
    }

    /** {@inheritDoc} */
    public final void setOptions(Map<String, ?> options) {
        /* Is verification of YubiKey owners enabled? */
        this.verify_yubikey_owner = true;
        if (options.get(OPTION_YUBICO_VERIFY_YK_OWNER) != null) {
            if ("false".equals(options.get(OPTION_YUBICO_VERIFY_YK_OWNER).toString())) {
                this.verify_yubikey_owner = false;
            }
        }
        if (options.get(OPTION_YUBICO_LDAP_PUBLICID_ATTRIBUTE) != null) {
            this.publicid_attribute = options.get(OPTION_YUBICO_LDAP_PUBLICID_ATTRIBUTE).toString();
        }
        if (options.get(OPTION_YUBICO_LDAP_USERNAME_ATTRIBUTE) != null) {
            this.username_attribute = options.get(OPTION_YUBICO_LDAP_USERNAME_ATTRIBUTE).toString();
        }
        if (options.get(OPTION_LDAP_URL) != null) {
            this.ldap_url = options.get(OPTION_LDAP_URL).toString();
        }
        if (options.get(OPTION_LDAP_BASE_DN) != null) {
            this.ldap_base_dn = options.get(OPTION_LDAP_BASE_DN).toString();
        }
        if (options.get(OPTION_LDAP_BIND_DN) != null) {
            this.ldap_bind_dn = options.get(OPTION_LDAP_BIND_DN).toString();
        }
        if (options.get(OPTION_LDAP_BIND_CREDENTIAL) != null) {
            this.ldap_bind_credential = options.get(OPTION_LDAP_BIND_CREDENTIAL).toString();
        }
        LdapConfig config = new LdapConfig(this.ldap_url, this.ldap_base_dn);
        config.setBindDn(this.ldap_bind_dn);
        config.setBindCredential(this.ldap_bind_credential);
        this.ldap = new Ldap(config);
    }
}


/**
 * Copyright (c) 2011, Yubico AB.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * * Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following
 *   disclaimer in the documentation and/or other materials provided
 *   with the distribution.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
 *  CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 *  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 *  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
 *  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 *  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 *  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 *  ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
 *  TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
 *  THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 *  SUCH DAMAGE.
 */
package com.yubico.jaas;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.URL;
import java.net.URLConnection;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import javax.security.auth.Subject;
import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.auth.login.LoginException;
import javax.security.auth.spi.LoginModule;

import org.apache.commons.codec.binary.Base64;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A JAAS module for verifying OATH OTPs (One Time Passwords) using
 * HTTP Basic Auth to a backend server doing the real authentication.
 *
 * @author Fredrik Thulin (fredrik@yubico.com)
 *
 */
public class HttpOathOtpLoginModule implements LoginModule {

	/* Options */
	public static final String OPTION_HTTPOATHOTP_PROTECTED_URL			= "protectedUrl";
	public static final String OPTION_HTTPOATHOTP_EXPECTED_OUTPUT		= "expectedOutput";
	public static final String OPTION_HTTPOATHOTP_MIN_LENGTH			= "minLength";
	public static final String OPTION_HTTPOATHOTP_MAX_LENGTH			= "maxLength";
	public static final String OPTION_HTTPOATHOTP_REQUIRE_ALL_DIGITS	= "requireAllDigits";
	public static final String OPTION_HTTPOATHOTP_ID_REALM				= "id_realm";

	public String protectedUrl;
	public String expectedOutput = "Authenticated OK";
	public int minLength = 6;
	public int maxLength = 12;
	public boolean requireAllDigits = true;
	public String idRealm;

	/* JAAS stuff */
	private Subject subject;
	private CallbackHandler callbackHandler;

	private final Logger log = LoggerFactory.getLogger(HttpOathOtpLoginModule.class);

	private YubikeyPrincipal principal;

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#abort()
	 */
	public boolean abort() throws LoginException {
		subject.getPrincipals().remove(this.principal);
		return true;
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#commit()
	 */
	public boolean commit() throws LoginException {
		log.trace("In commit()");
		subject.getPrincipals().add(this.principal);
		return true;
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#initialize(javax.security.auth.Subject, javax.security.auth.callback.CallbackHandler, java.util.Map, java.util.Map)
	 */
	public void initialize(Subject newSubject, CallbackHandler newCallbackHandler,
			Map<String, ?> sharedState, Map<String, ?> options) {

		log.trace("Initializing HTTP OATH OTP LoginModule");

		this.subject = newSubject;
		this.callbackHandler = newCallbackHandler;

		this.protectedUrl = options.get(OPTION_HTTPOATHOTP_PROTECTED_URL).toString();

		if (options.get(OPTION_HTTPOATHOTP_EXPECTED_OUTPUT) != null) {
			this.expectedOutput = options.get(OPTION_HTTPOATHOTP_EXPECTED_OUTPUT).toString();
		}

		if (options.get(OPTION_HTTPOATHOTP_MIN_LENGTH) != null) {
			this.minLength = Integer.parseInt(options.get(OPTION_HTTPOATHOTP_MIN_LENGTH).toString());
		}
		if (options.get(OPTION_HTTPOATHOTP_MAX_LENGTH) != null) {
			this.maxLength = Integer.parseInt(options.get(OPTION_HTTPOATHOTP_MAX_LENGTH).toString());
		}
		if (options.get(OPTION_HTTPOATHOTP_REQUIRE_ALL_DIGITS) != null) {
			String s = options.get(OPTION_HTTPOATHOTP_REQUIRE_ALL_DIGITS).toString();
			if (s.equals("true")) {
				this.requireAllDigits = true;
			} else if (s.equals("false")) {
				this.requireAllDigits = false;
			} else {
				log.error("Bad value for option {}", OPTION_HTTPOATHOTP_REQUIRE_ALL_DIGITS);
			}
		}
		/* Realm of principals added after authentication */
		if (options.get(OPTION_HTTPOATHOTP_ID_REALM) != null) {
			this.idRealm = options.get(OPTION_HTTPOATHOTP_ID_REALM).toString();
		}
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#login()
	 */
	public boolean login() throws LoginException {
		log.trace("Begin OATH OTP login");

		if (callbackHandler == null) {
			throw new LoginException("No callback handler available in login()");
		}

		NameCallback nameCb = new NameCallback("Enter username: ");

		List<String> otps = get_tokens(nameCb);

		for (String otp : otps) {
			String userName = nameCb.getName();

			log.trace("Checking OATH OTP for user {}", userName);

			if (verify_otp(userName, otp)) {
				log.info("OATH OTP verified successfully");
				principal = new YubikeyPrincipal(userName, this.idRealm);
				return true;
			}
			log.info("OATH OTP did NOT verify");
		}
		throw new LoginException("OATH OTP verification failed");
	}

	/**
	 * Access protectedUrl using userName and otp for basic auth.
	 * Check if what we get back contains expectedOutput.
	 * @param userName
	 * @param otp
	 * @return boolean
	 */
	boolean verify_otp(String userName, String otp) {
		try {
			String authString = userName + ":" + otp;
			String authStringEnc = Base64.encodeBase64URLSafeString(authString.getBytes());

			BufferedReader in = attemptAuthentication(authStringEnc);

			String inputLine;
			while ((inputLine = in.readLine()) != null) {
				if (inputLine.contains(expectedOutput)) {
					return true;
				}
			}
		} catch (Exception ex) {
			log.error("Failed verifying OATH OTP :", ex);
		}
		return false;
	}

	BufferedReader attemptAuthentication(String authStringEnc) throws IOException {
		URL url = new URL(this.protectedUrl);
		URLConnection conn = url.openConnection();
		conn.setRequestProperty("Authorization", "Basic " + authStringEnc);
		conn.connect();

		return new BufferedReader(new InputStreamReader(
                conn.getInputStream()
		));
	}

	private List<String> get_tokens(NameCallback nameCb) throws LoginException {
		MultiValuePasswordCallback mv_passCb = new MultiValuePasswordCallback("Enter authentication tokens: ", false);
		List<String> result = new ArrayList<String>();

		try {
			/* Fetch a password using the callbackHandler */
			callbackHandler.handle(new Callback[] { nameCb, mv_passCb });

			for (char[] c : mv_passCb.getSecrets()) {
				String s = new String(c);
				/* Check that OTP looks like OATH before we verify it. User might have entered
				 * some other password instead of an OTP, and we don't want to send that, possibly
				 * in clear text, over the network.
				 */
				if (s.length() < this.minLength) {
					log.info("Skipping token, not a valid OATH OTP token (too short, {} < {})", s.length(), this.minLength);
				} else if (s.length() > this.maxLength) {
					log.info("Skipping token, not a valid OATH OTP token (too long, {} > {})", s.length(), this.maxLength);
				} else {
					if (this.requireAllDigits) {
						if (s.matches("^[0-9]+$")) {
							result.add(s);
						} else {
							log.info("Skipping token, not a valid OATH OTP token (non-digits not allowed)");
						}
					} else {
						result.add(s);
					}
				}
			}
		} catch (UnsupportedCallbackException ex) {
			log.error("Callback type not supported", ex);
		} catch (IOException ex) {
			log.error("CallbackHandler failed", ex);
		}

		return result;
	}

	/* (non-Javadoc)
	 * @see javax.security.auth.spi.LoginModule#logout()
	 */
	public boolean logout() throws LoginException {
		log.trace("In logout()");
		subject.getPrincipals().remove(this.principal);
		return false;
	}
}