# -*- coding: utf-8 -*-

import hashlib,hmac
import struct
import time
import base64

class HOTP:

	def __init__(self, secret = None, counter = 0):
		self.secret = secret.encode() if secret != None else None
		self.counter = counter
	
	def setSecret(self, secret, format = None):
		self.secret = self._convertSecret(secret, format)
	
	def gen(self, counter = None):
		if counter == None:
			if self.counter == None:
				self.counter = 0
			counter = self.counter
		return self._gen(self.secret, counter)
	
	def verify(self, token, counter = None):
		if counter == None:
			if self.counter == None:
				self.counter = 0
			counter = self.counter
			if (token == self.gen(counter)):
				self.counter += 1
				return true
			return false
		return token == self.gen(counter)
	
	@staticmethod
	def _gen(secret, counter):
		data = counter.to_bytes(8, byteorder='big', signed=False)
		hash = hmac.new(secret, data, hashlib.sha1).digest()
		offset = hash[19] & 0xf
		num = (hash[offset] & 0x7F) << 24 | \
			(hash[offset + 1] & 0xFF) << 16 | \
			(hash[offset + 2] & 0xFF) << 8  | \
			(hash[offset + 3] & 0xFF)
		num = num % 1000000
		return "%06d"%(num)
	
	@staticmethod
	def _convertSecret(secret, format):
		if type(secret) == bytes:
			return secret
		if type(secret) == str:
			if 'hex' == format:
				return bytes.fromhex(secret)
			elif 'base64' == format:
				return base64.b64decode(secret.encode())
			elif 'string' == format:
				return secret.encode()
			else:
				raise Exception('参数错误')
		elif type(secret) == bytes:
			return secret
		else:
			raise Exception('参数错误')

class TOTP:

	def __init__(self, secret = None, window = 3, interval = 30):
		self.hotp = HOTP(secret, 0)
		self.window = window
		self.interval = interval
	
	def setSecret(self, secret, format = None):
		self.hotp.setSecret(secret, format)
	
	def gen(self, timestamp = None):
		if timestamp == None:
			timestamp = round(time.time() * 1000)
		counter = int(timestamp / 1000 / self.interval)
		return self.hotp.gen(counter)
	
	def verify(self, token, timestamp = None):
		if timestamp == None:
			timestamp = round(time.time() * 1000)
		counter = int(timestamp / 1000 / self.interval)
		
		i = 0
		w = 0
		while(True):
			if token == self.hotp.gen(counter - i):
				return True
			w += 1
			if w >= self.window:
				break
			if i > 0:
				if token == self.hotp.gen(counter + i):
					return True
				w += 1
				if w >= self.window:
					break
			i += 1
		
		return False

# -------------------------------
"""
if __name__ == "__main__":
	p1 = HOTP("123456", 0)
	print(p1.gen()) # 186818
	print(p1.gen(1234)) # 263197
	print(p1.verify('263197', 1234))
	
	p2 = TOTP("123456", 3, 30)
	print(p2.gen(1614054214456)) # 330925
	print(p2.verify('330925', 1614054214456))

	p3 = TOTP('helloworld', 3, 30)
	p3.setSecret("helloworld", "string")
	p3.setSecret("68656c6c6f776f726c64", "hex")
	p3.setSecret("aGVsbG93b3JsZA==", "base64")
	print(p3.gen(1614054214456)) # 759440
	print(p3.verify('759440', 1614054214456))
"""