Enhance OTP script with QR code generation and improve secret key handling

.gitignore

@@ -1,2 +1,3 @@
 .venv/
 ft_otp.key
+key2.hex

ft_otp.py

@@ -1,66 +1,102 @@
 import argparse
 from cryptography.fernet import Fernet
 from hashlib import sha256
-from base64 import urlsafe_b64encode
+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="hexfile", type=str, help="Generate a new OTP key and save it to the specified file.")
+	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_otp_key(args.g)
+		generate_secret_key(args.g)
 	elif args.k:
 		# Call the function to generate a OTP from the specified key file
-		generate_otp_from_key(args.k)
+		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()
+		print("QR code generated successfully.")
 	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 is_hex(s):
+    try:
+        int(s, 16)
+        return True
+    except ValueError:
+        return False
 
-def generate_otp_from_key(keyfile):
-	import time
+def generate_secret_key(secret_file):
+	print(f"Generating OTP key from {secret_file}...")
+	machine_id = get_machine_hash()
+	with open(secret_file, 'rb') as f:
+		secret = f.read().strip()
+	if not is_hex(secret.decode()):
+		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)
+	with open('./ft_otp.key', 'wb') as f:
+		f.write(secret_enc)
+
+def read_secret_key(keyfile):
+	"""Read the secret key from the specified file."""
 	machine_id = get_machine_hash()
 	with open(keyfile, 'rb') as f:
-		otp_key = f.read().strip()
+		secret_enc = f.read().strip()
 	fernet = Fernet(machine_id)
-	otp = fernet.decrypt(otp_key)
+	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)
+	print(f"Secret key read from {keyfile}: {secret.decode()}")
+	base32 = base64.b32encode(base64.b16decode(secret.upper())).decode().strip('=')
+	print(str(base32))
+	otp_uri = f"otpauth://totp/FTOTP?secret={base32}&issuer=FTOTP&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(otp, counter)
+	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 urlsafe_b64encode(sha256_hash)
+		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"""
-	import hmac
-	import struct
-	import hashlib
-
 	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

key.hex

@@ -1 +1 @@
-d779574a5815a54089b260694be7b29838e28a72e943bdd41e47656885118c7a3f29f6611ca1383f7e67e9868a80b8fd68c71be79b418e8984bdaef4991932e4
+d779574a5815a54089b260694be7b29838e28a72e943bdd41e47656885118c7a

requirements.txt

@@ -0,0 +1,3 @@
+cffi==1.17.1
+cryptography==45.0.4
+pycparser==2.22