python totp代码

import time
import datetime
import math
import hmac
import base64
from hashlib import sha1

class GoogleAuthenticator(object):

    def __init__(self,secretKey=None,digits=6,interval=30):
        self.secretKey = secretKey
        self.digits = digits
        self.interval = interval

    def __str_extend(self,old_str,length,extend_str):
        new_strs = None
        if len(old_str) < length:
            clen = length - len(old_str)
            s0 = ''.join([extend_str for i in range(clen)])
            new_strs = s0 + old_str
        elif len(old_str) == length:
            new_strs =  old_str
        return new_strs

    def __str_split(self,old_str,split_len,prefix='',suffix=''):
        array = []
        str_len = len(old_str)
        for i in range(0,str_len,4):
            if i + 4 < str_len:
                array.append(prefix+old_str[i:i+4]+suffix)
            else:
                array.append(prefix+old_str[i:str_len]+suffix)
        return array

    def __base32tohex(self,base32str):
        base32chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
        bits = ""
        for i in range(len(base32str)):
            char = str(base32str[i]).upper()
            idx = int(base32chars.index(char))
            bit = self.__str_extend(format(idx,'b'),5,'0')
            if bit is None:
                raise Exception('bit={},len={}'.format(bit,len(bit)))
            bits += bit
        bitArray = self.__str_split(bits,4,'0b')
        hexstr = ''.join([format(int(i,2),'x') for i in bitArray])
        return hexstr

    def __get_HexSecret(self):
        hexstr = self.__base32tohex(self.secretKey)
        return hexstr

    def __byte_secret(self):
        missing_padding = len(self.secretKey) % 8
        if missing_padding != 0:
            self.secret += '=' * (8 - missing_padding)
        return base64.b32decode(self.secretKey, casefold=True)

    def __int_to_bytestring(self,i, padding=8):
        result = bytearray()
        while i != 0:
            result.append(i & 0xFF)
            i >>= 8
        return bytes(bytearray(reversed(result)).rjust(padding, b'\0'))

    def __timecode(self, for_time):
        i = time.mktime(for_time.timetuple())
        return int(i / self.interval)

    def get_QR_url(self):
    	QR_url = "https://chart.googleapis.com/chart?chs=200x200&cht=qr&chl=200x200&chld=M|0&cht=qr&chl=otpauth://totp/[email protected]\%3Fsecret%3D" + self.secretKey
    	return QR_url

    def getTotp(self):
        key = self.secretKey
        timestamp = self.__timecode(datetime.datetime.now())
        hasher = hmac.new(self.__byte_secret(), self.__int_to_bytestring(timestamp), sha1)
        # print(hasher.hexdigest())
        hmac_hash = bytearray(hasher.digest())
        offset = hmac_hash[-1] & 0xf
        code = ((hmac_hash[offset] & 0x7f) << 24 |
                (hmac_hash[offset + 1] & 0xff) << 16 |
                (hmac_hash[offset + 2] & 0xff) << 8 |
                (hmac_hash[offset + 3] & 0xff))
        str_code = str(code % 10 ** self.digits)
        while len(str_code) < self.digits:
            str_code = '0' + str_code
        return str_code

    def verifyTotp(self,str_code):
    	gTotp = self.getTotp()
    	if gTotp == str_code:
    		return True
    	else:
    		return False