ft_otp/ft_otp.py

126 lines
3.9 KiB
Python

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()