import base64, time, struct, hmac, hashlib
import os


def get_google_code_by_date(secret_key):
    """
    基于时间的算法
    :param secret_key:
    :return:
    """
    decode_secret = base64.b32decode(secret_key, True)
    # 解码 Base32 编码过的 bytes-like object 或 ASCII 字符串 s 并返回解码过的 bytes。

    interval_number = int(time.time() // 30)

    message = struct.pack(">Q", interval_number)
    digest = hmac.new(decode_secret, message, hashlib.sha1).digest()
    index = ord(chr(digest[19])) % 16
    google_code = (struct.unpack(">I", digest[index:index + 4])[0] & 0x7fffffff) % 1000000

    return "%06d" % google_code


def get_google_code_by_time(secret_key, num):
    """
    基于计数器的算法
    :param secret_key:
    :return:
    """
    decode_secret = base64.b32decode(secret_key, True)
    # 解码 Base32 编码过的 bytes-like object 或 ASCII 字符串 s 并返回解码过的 bytes。

    result = {}

    for interval_number in range(num, num + 5):
        message = struct.pack(">Q", interval_number)
        digest = hmac.new(decode_secret, message, hashlib.sha1).digest()
        index = ord(chr(digest[19])) % 16
        google_code = (struct.unpack(">I", digest[index:index + 4])[0] & 0x7fffffff) % 1000000

        result[interval_number] = "%06d" % google_code

    return result


def get_secret_key():
    """
    随机生成secret_key
    :return:
    """
    return base64.b32encode(os.urandom(10)).decode('utf-8')