import argparse
from cryptography.fernet import Fernet
from hashlib import sha256
from base64 import urlsafe_b64encode

def main():
	parser = argparse.ArgumentParser(description="Run the FT OTP script.")
	parser.add_argument("-g", metavar="hexfile", type=str, help="Generate a new OTP key and save it to the specified file.")
	parser.add_argument("-k", metavar="keyfile", type=str, help="Generate a OTP from the specified key file.")
	args = parser.parse_args()

	if args.g:
		# Call the function to generate a new OTP key
		generate_otp_key(args.g)
	elif args.k:
		# Call the function to generate a OTP from the specified key file
		generate_otp_from_key(args.k)
	else:
		print("No action specified. Use -g to generate a new OTP key or -k to generate a OTP from a key file.")
		parser.print_help()

def generate_otp_key(hexfile):
	print(f"Generating OTP key from {hexfile}...")
	machine_id = get_machine_hash()
	with open(hexfile, 'rb') as f:
		hex_key = f.read().strip()
	fernet = Fernet(machine_id)
	otp_key = fernet.encrypt(hex_key)
	with open('./ft_otp.key', 'wb') as f:
		f.write(otp_key)

def generate_otp_from_key(keyfile):
	import time
	machine_id = get_machine_hash()
	with open(keyfile, 'rb') as f:
		otp_key = f.read().strip()
	fernet = Fernet(machine_id)
	otp = fernet.decrypt(otp_key)
	counter = int(time.time() // 30)  # Using time-based counter for HOTP
	otp = HOTP(otp, counter)
	# Print the generated OTP
	print(f"Generated OTP: {otp}")

def get_machine_hash() -> bytes:
	"""Retrieve the machine ID from /etc/machine-id."""
	try:
		with open('/etc/machine-id', 'r') as f:
			machine_id =  f.read().strip()
		sha256_hash = sha256(machine_id.encode()).digest()
		return urlsafe_b64encode(sha256_hash)
	except FileNotFoundError:
		print("Machine ID file not found. Please ensure the system is properly configured.")
		return None



def HOTP(key, counter):
	"""Generate a one-time password using the HOTP algorithm. using RFC4226"""
	import hmac
	import struct
	import hashlib

	counter = struct.pack('>Q', counter)
	hmac_hash = hmac.new(key, counter, hashlib.sha1).digest()
	offset = hmac_hash[-1] & 0x0F
	otp = (struct.unpack('>I', hmac_hash[offset:offset + 4])[0] & 0x7FFFFFFF) % 1000000
	return str(otp).zfill(6)


if __name__ == "__main__":
	main()