/*
 * Copyright (c) 2018-2019 yingtingxu(徐应庭). All rights reserved.
 */

package com.arch.totp;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;

/**
 * The implementation of TOTP: Time-Based One-Time Password Algorithm
 * <p>
 * see:
 * TOTP: https://tools.ietf.org/html/rfc6238
 */
public class Totp {
    public enum HashAlgorithm {
        HmacSHA1("HmacSHA1"), HmacSHA256("HmacSHA256"), HmacSHA512("HmacSHA512");

        private String name;

        HashAlgorithm(String name) {
            this.name = name;
        }

        public String getName() {
            return name;
        }

        @Override
        public String toString() {
            return getName();
        }
    }

    // default hash algorithm
    public static final HashAlgorithm DEFAULT_HASH_ALGORITHM = HashAlgorithm.HmacSHA1;

    // default time step in seconds
    public static final int DEFAULT_TIME_STEP = 30;

    // default number of digits
    public static final int DEFAULT_DIGITS = 8;

    // T0 is the Unix time to start counting time steps
    // (default value is 0, i.e., the Unix epoch)
    public static final int DEFAULT_T0 = 0;

    private static final int[] DIGITS_POWER
            // 0  1   2    3     4      5       6        7         8
            = {1, 10, 100, 1000, 10000, 100000, 1000000, 10000000, 100000000};

    private int timeStep;
    private int t0;
    private int digits;
    private HashAlgorithm hashAlgorithm;

    private Totp(int timeStep, int t0, int digits, HashAlgorithm hashAlgorithm) {
        this.timeStep = timeStep;
        this.t0 = t0;
        this.digits = digits;
        this.hashAlgorithm = hashAlgorithm;
    }

    public static Totp newInstance() {
        return new Totp(
                DEFAULT_TIME_STEP,
                DEFAULT_T0,
                DEFAULT_DIGITS,
                DEFAULT_HASH_ALGORITHM);
    }

    public static Totp newInstance(
            int timeStep,
            int t0,
            int digits,
            HashAlgorithm hashAlgorithm) {
        return new Totp(timeStep, t0, digits, hashAlgorithm);
    }

    /**
     * This method generates a TOTP value for the given secret.
     *
     * @param secret: the shared secret, HEX encoded
     * @return: a numeric String in base 10 that includes truncated digits
     */
    public String generateTotp(String secret) {
        return generateTotp(secret, getCurrentTimeStepNumber());
    }

    /**
     * This method generates a TOTP value for the given secret.
     *
     * @param secret:         the shared secret, HEX encoded
     * @param timeStepNumber: the number of time step between the initial time T0 and the current unix time.
     * @return: a numeric String in base 10 that includes truncated digits
     */
    private String generateTotp(String secret, long timeStepNumber) {
        Assert.hasText(secret, "secret cannot be null or empty");

        // Using the counter
        // First 8 bytes are for the movingFactor
        // Compliant with base RFC 4226 (HOTP)
        byte[] timeStepBytes = ByteBuffer.allocate(8)
                .putLong(timeStepNumber)
                .array();
                
        // TODO: if the secret is HEX, not using UTF-8 encoding.
        byte[] secretBytes = secret.getBytes(StandardCharsets.UTF_8);

        // the output would be a 20 byte long in RFC4226, see: https://tools.ietf.org/html/rfc4226
        byte[] hmacHash = getHmacSHA(secretBytes, timeStepBytes);

        // put selected bytes into result int
        int offset = hmacHash[hmacHash.length - 1] & 0xf;

        int truncatedHash = ((hmacHash[offset] & 0x7f) << 24) |
                ((hmacHash[offset + 1] & 0xff) << 16) |
                ((hmacHash[offset + 2] & 0xff) << 8) |
                (hmacHash[offset + 3] & 0xff);

        int otp = truncatedHash % DIGITS_POWER[digits];

        String result = Integer.toString(otp);
        while (result.length() < digits) {
            result = "0" + result;
        }
        return result;
    }

    /**
     * Validates the TOTP for the specified secret
     *
     * @param secret: the shared secret, HEX encoded
     * @param totp:   the TOTP generated by the secret
     * @return {@code true} if the TOTP is valid
     */
    public boolean validateTotp(String secret, String totp) {
        if (!StringUtils.hasText(secret) || !StringUtils.hasText(totp)) {
            return false;
        }

        // Allow a variance of no greater than 90 seconds in either direction
        long timeStepNumber = getCurrentTimeStepNumber();
        for (int i = -2; i <= 2; i++) {
            String totp2 = generateTotp(secret, timeStepNumber + i);
            if (totp.equals(totp2)) {
                return true;
            }
        }

        // No match
        return false;
    }

    private long getCurrentTimeStepNumber() {
        return (System.currentTimeMillis() / 1000L - t0) / timeStep;
    }

    /**
     * This method uses the JCE to provide the crypto algorithm.
     * HMAC computes a Hashed Message Authentication Code with the
     * crypto hash algorithm as a parameter.
     *
     * @param secretBytes: the bytes to use for the HMAC key
     * @param textBytes:   the message or text to be authenticated
     */
    private byte[] getHmacSHA(byte[] secretBytes, byte[] textBytes) {
        try {
            Mac hmac = Mac.getInstance(hashAlgorithm.getName());
            SecretKeySpec spec = new SecretKeySpec(secretBytes, "RAW");
            hmac.init(spec);
            return hmac.doFinal(textBytes);
        } catch (GeneralSecurityException gse) {
            throw new IllegalStateException(gse);
        }
    }
}