import argparse
from cryptography.fernet import Fernet
from hashlib import sha256
import base64
import hmac
import struct
import hashlib
import qrcode_terminal
import time



def main():
	parser = argparse.ArgumentParser(description="Run the FT OTP script.")
	parser.add_argument("-g", metavar="secret_file", 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.")
	parser.add_argument("-q", action="store_true", help="Generate a QR code for the OTP key.")
	args = parser.parse_args()

	if args.g:
		# Call the function to generate a new OTP key
		generate_secret_key(args.g)
	elif args.k:
		# Call the function to generate a OTP from the specified key file
		generate_otp_from_secret_key(args.k)
	elif args.q:
		# Call the function to generate a QR code for the OTP key
		generate_qr_code()
	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 is_hex(s):
    try:
        int(s, 16)
        return True
    except ValueError:
        return False

def generate_secret_key(secret_file):
	print(f"Generating OTP key from {secret_file}...")
	machine_id = get_machine_hash()
	try:
		with open(secret_file, 'r') as f:
			secret = f.read().strip()
	except FileNotFoundError:
		print(f"Secret file {secret_file} not found. Please provide a valid file.")
		return
	if not secret:
		print("The provided secret file is empty. Please provide a valid secret.")
		return
	if not is_hex(secret):
		print("The provided secret is not a valid hexadecimal string.")
		return
	if len(secret) < 64:
		print("The provided secret is too short. It should be at least 64 characters long.")
		return
	fernet = Fernet(machine_id)
	secret_enc = fernet.encrypt(secret.encode())
	try:
		with open('./ft_otp.key', 'wb') as f:
			f.write(secret_enc)
		print("OTP key generated and saved to ft_otp.key")
	except IOError as e:
		print(f"Error writing to file: {e}")
		return

def read_secret_key(keyfile):
	"""Read the secret key from the specified file."""
	machine_id = get_machine_hash()
	if machine_id is None:
		print("Machine ID could not be retrieved. Exiting.")
		return None
	try:
		with open(keyfile, 'rb') as f:
			secret_enc = f.read().strip()
	except FileNotFoundError:
		print(f"Key file {keyfile} not found. Please provide a valid file.")
		return None
	if not secret_enc:
		print("The provided key file is empty. Please provide a valid key file.")
		return None
	fernet = Fernet(machine_id)
	secret = fernet.decrypt(secret_enc)
	return secret

def generate_qr_code():
	"""Generate a QR code for the OTP key."""
	keyfile = './ft_otp.key'
	secret = read_secret_key(keyfile)
	base32 = base64.b32encode(base64.b16decode(secret.upper())).decode().strip('=')
	otp_uri = f"otpauth://totp/FTOTP?secret={base32}&issuer=whaffman&algo=SHA1&digits=6&period=30"
	qrcode_terminal.draw(otp_uri)

def generate_otp_from_secret_key(keyfile):
	secret = read_secret_key(keyfile)
	counter = int(time.time() // 30)  # Using time-based counter for HOTP
	otp = HOTP(secret, 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 base64.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"""
	counter = struct.pack('>Q', counter)
	key = base64.b16decode(key.upper())
	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()