Building a Secure 2FA Authenticator with Python

class Authenticator:
    def __init__(self):
        self.data_file = "auth_data.enc"
        self.salt_file = "salt.bin"
        self.key = None
        self.fernet = None
        self.websites = {}
        self.initialize_encryption()
        self.load_data()

This method initializes file names for storing encrypted data and salt, sets up variables for encryption keys and the Fernet encryption object, and creates an empty dictionary to store website information. It then calls two important methods: initialize_encryption() and load_data().

Encryption Initialization

This method sets up the encryption using a user-provided master password and a salt. If a salt does not exist, it generates a new one.

    def initialize_encryption(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        if not os.path.exists(self.salt_file):
            salt = os.urandom(16)
            with open(self.salt_file, "wb") as f:
                f.write(salt)
        else:
            with open(self.salt_file, "rb") as f:
                salt = f.read()

        password = getpass("Enter Master Password: ")
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        self.key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        self.fernet = Fernet(self.key)

This method performs several important security operations:

1. It checks for an existing salt file. If not found, it generates a new 16-byte random salt and saves it.
2. It prompts the user for a master password using getpass for secure input.
3. It uses PBKDF2HMAC (Password-Based Key Derivation Function 2 with HMAC) to derive an encryption key from the master password and salt.
4. It creates a Fernet encryption object using the derived key.

The use of a salt and PBKDF2 adds significant security to the encryption process, making it resistant to rainbow table attacks and slowing down brute-force attempts.

Loading and Saving Data

The script includes methods to load and save encrypted authentication data. These methods ensure that sensitive information is securely stored.

    def load_data(self):
        if os.path.exists(self.data_file):
            try:
                with open(self.data_file, "rb") as f:
                    encrypted_data = f.read()
                    decrypted_data = self.fernet.decrypt(encrypted_data)
                    self.websites = json.loads(decrypted_data.decode())
            except Exception as e:
                print(f"Error loading data: {e}")
                self.websites = {}

    def save_data(self):
        try:
            encrypted_data = self.fernet.encrypt(json.dumps(self.websites).encode())
            with open(self.data_file, "wb") as f:
                f.write(encrypted_data)
        except Exception as e:
            print(f"Error saving data: {e}")

The script manages data through two primary methods: load_data() and save_data(). The load_data() method attempts to load existing encrypted data. This method checks for an existing data file, reads the encrypted data, decrypts it using the Fernet object, and loads the decrypted JSON data into the websites dictionary.

The save_data() method handles saving the current state of the websites dictionary. This method serializes the websites dictionary to JSON, encrypts it using the Fernet object, and writes the encrypted data to the file.

    def clear_screen(self):
        os.system('cls' if os.name == 'nt' else 'clear')

User Interface

The script provides a command-line interface that allows users to interact with the script. The main menu offers options to generate OTPs, add new websites, and delete websites. The display_menu() method presents the main menu options:

    def display_menu(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        print("1. Generate OTP")
        print("2. Add New Website")
        print("3. Delete Website")
        print("4. Exit\n")
        return input("Choose an Option: ")

Generating OTP

This method generates OTPs for the selected website. It uses the pyotp library to generate the OTP based on the stored secret key.

    def generate_otp(self):
        while True:
            self.clear_screen()
            print("2FA Authenticator\n")
            if not self.websites:
                print("No websites saved. Please add a website first.")
                input("\nPress Enter to return to Main Menu...")
                return

            print("Choose a website to generate an OTP code for:\n")
            for idx, website in enumerate(self.websites.keys(), 1):
                print(f"{idx}. {website}")
            print(f"\n{len(self.websites) + 1}. Return to Main Menu")

            try:
                choice = int(input("\nChoose a website (number): "))
                if choice == len(self.websites) + 1:
                    return
                if choice < 1 or choice > len(self.websites):
                    raise ValueError

                website = list(self.websites.keys())[choice - 1]
                secret = self.websites[website]

                totp = pyotp.TOTP(secret)
                otp = totp.now()

                self.clear_screen()
                print("2FA Authenticator\n")
                print(f"Website: {website}")
                print(f"OTP: {otp}")
                print(f"Valid for: {30 - int(time.time()) % 30} seconds")
                input("\nPress Enter to generate another OTP or return to Main Menu...")

            except ValueError:
                print("Invalid input. Please enter a valid number.")
                input("\nPress Enter to try again...")
            except Exception as e:
                print(f"An error occurred: {e}")
                input("\nPress Enter to try again...")

This method:

1. Displays a list of saved websites
2. Prompts the user to choose a website
3. Retrieves the secret key for the chosen website
4. Uses the pyotp library to generate a Time-based One-Time Password (TOTP)
5. Displays the OTP along with its remaining validity time

The use of TOTP ensures that the generated passwords are time-sensitive and change every 30 seconds, providing an additional layer of security.

Adding a Website

This method allows users to add a new website and its corresponding secret key. The data is then saved securely.

    def add_website(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        website = input("Enter Website Name: ")
        secret = input("Enter Private key: ")
        self.websites[website] = secret
        self.save_data()
        print("Website added successfully!\n")
        input("Press Enter to continue...")

Deleting a Website

Users can delete a website and its corresponding secret key using this method. The updated data is then saved.

    def delete_website(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        if not self.websites:
            print("No websites saved. Nothing to delete.")
            input("\nPress Enter to return to Main Menu...")
            return

        print("Choose a website to delete:\n")
        for idx, website in enumerate(self.websites.keys(), 1):
            print(f"{idx}. {website}")
        print(f"\n{len(self.websites) + 1}. Return to Main Menu")

        try:
            choice = int(input("\nChoose a website to delete (number): "))
            if choice == len(self.websites) + 1:
                return
            if choice < 1 or choice > len(self.websites):
                raise ValueError

            website = list(self.websites.keys())[choice - 1]
            confirm = input(f"Are you sure you want to delete {website}? (y/n): ").lower()
            if confirm == 'y':
                del self.websites[website]
                self.save_data()
                print(f"\n{website} has been deleted.")
            else:
                print("\nDeletion cancelled.")

        except ValueError:
            print("Invalid input. Please enter a valid number.")
        except Exception as e:
            print(f"An error occurred: {e}")

        input("\nPress Enter to continue...")

This method displays a list of saved websites, prompts the user to choose one for deletion, asks for confirmation, and then removes the selected website from the websites dictionary if confirmed.

Main Loop

The script’s main loop continuously displays the menu and processes user input until the user decides to exit.

    def run(self):
        while True:
            choice = self.display_menu()
            if choice == "1":
                self.generate_otp()
            elif choice == "2":
                self.add_website()
            elif choice == "3":
                self.delete_website()
            elif choice == "4":
                break
            else:
                print("Invalid choice. Please try again.\n")
                input("Press Enter to continue...")

Script Execution

The script execution starts by initializing the Authenticator class and running its run method.

if __name__ == "__main__":
    auth = Authenticator()
    auth.run()

Full Script

Here is the complete script for your reference:

#!/usr/bin/env python3
import base64
import json
import os
import pyotp
import time
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from getpass import getpass

class Authenticator:
    def __init__(self):
        # Initializes variables and encryption setup.
        self.data_file = "auth_data.enc"  # File to store encrypted authentication data.
        self.salt_file = "salt.bin"       # File to store the salt for key derivation.
        self.key = None                   # Variable to store the encryption key.
        self.fernet = None                # Variable to store the Fernet encryption object.
        self.websites = {}                # Dictionary to store website and their secrets.
        self.initialize_encryption()      # Initialize encryption setup.
        self.load_data()                  # Load existing data if available.

    def initialize_encryption(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        if not os.path.exists(self.salt_file):
            # Generate and save a new salt if it doesn't exist.
            salt = os.urandom(16)
            with open(self.salt_file, "wb") as f:
                f.write(salt)
        else:
            # Load the existing salt.
            with open(self.salt_file, "rb") as f:
                salt = f.read()

        password = getpass("Enter Master Password: ")  # Prompt user for master password.
        # Derive encryption key from password using PBKDF2HMAC.
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        self.key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        self.fernet = Fernet(self.key)  # Initialize Fernet with the derived key.

    def load_data(self):
        if os.path.exists(self.data_file):
            try:
                # Load and decrypt existing authentication data.
                with open(self.data_file, "rb") as f:
                    encrypted_data = f.read()
                    decrypted_data = self.fernet.decrypt(encrypted_data)
                    self.websites = json.loads(decrypted_data.decode())
            except Exception as e:
                print(f"Error loading data: {e}")
                self.websites = {}

    def save_data(self):
        try:
            # Encrypt and save authentication data.
            encrypted_data = self.fernet.encrypt(json.dumps(self.websites).encode())
            with open(self.data_file, "wb") as f:
                f.write(encrypted_data)
        except Exception as e:
            print(f"Error saving data: {e}")

    def clear_screen(self):
        # Clear the console screen.
        os.system('cls' if os.name == 'nt' else 'clear')

    def display_menu(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        # Display the main menu options.
        print("1. Generate OTP")
        print("2. Add New Website")
        print("3. Delete Website")
        print("4. Exit\n")
        return input("Choose an Option: ")

    def generate_otp(self):
        while True:
            self.clear_screen()
            print("2FA Authenticator\n")
            if not self.websites:
                # Inform user if no websites are saved.
                print("No websites saved. Please add a website first.")
                input("\nPress Enter to return to Main Menu...")
                return

            # Display list of saved websites.
            print("Choose a website to generate an OTP code for:\n")
            for idx, website in enumerate(self.websites.keys(), 1):
                print(f"{idx}. {website}")
            print(f"\n{len(self.websites) + 1}. Return to Main Menu")

            try:
                choice = int(input("\nChoose a website (number): "))
                if choice == len(self.websites) + 1:
                    return
                if choice < 1 or choice > len(self.websites):
                    raise ValueError

                website = list(self.websites.keys())[choice - 1]
                secret = self.websites[website]

                # Generate and display OTP for the selected website.
                totp = pyotp.TOTP(secret)
                otp = totp.now()

                self.clear_screen()
                print("2FA Authenticator\n")
                print(f"Website: {website}")
                print(f"OTP: {otp}")
                print(f"Valid for: {30 - int(time.time()) % 30} seconds")
                input("\nPress Enter to generate another OTP or return to Main Menu...")

            except ValueError:
                # Handle invalid input.
                print("Invalid input. Please enter a valid number.")
                input("\nPress Enter to try again...")
            except Exception as e:
                # Handle other errors.
                print(f"An error occurred: {e}")
                input("\nPress Enter to try again...")

    def add_website(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        # Prompt user to enter website name and secret key.
        website = input("Enter Website Name: ")
        secret = input("Enter Private key: ")
        self.websites[website] = secret
        self.save_data()  # Save the updated data.
        print("Website added successfully!\n")
        input("Press Enter to continue...")

    def delete_website(self):
        self.clear_screen()
        print("2FA Authenticator\n")
        if not self.websites:
            # Inform user if no websites are saved.
            print("No websites saved. Nothing to delete.")
            input("\nPress Enter to return to Main Menu...")
            return

        # Display list of saved websites for deletion.
        print("Choose a website to delete:\n")
        for idx, website in enumerate(self.websites.keys(), 1):
            print(f"{idx}. {website}")
        print(f"\n{len(self.websites) + 1}. Return to Main Menu")

        try:
            choice = int(input("\nChoose a website to delete (number): "))
            if choice == len(self.websites) + 1:
                return
            if choice < 1 or choice > len(self.websites):
                raise ValueError

            website = list(self.websites.keys())[choice - 1]
            confirm = input(f"Are you sure you want to delete {website}? (y/n): ").lower()
            if confirm == 'y':
                # Delete the selected website and save data.
                del self.websites[website]
                self.save_data()
                print(f"\n{website} has been deleted.")
            else:
                print("\nDeletion cancelled.")

        except ValueError:
            # Handle invalid input.
            print("Invalid input. Please enter a valid number.")
        except Exception as e:
            # Handle other errors.
            print(f"An error occurred: {e}")

        input("\nPress Enter to continue...")

    def run(self):
        # Main loop to run the authenticator.
        while True:
            choice = self.display_menu()
            if choice == "1":
                self.generate_otp()
            elif choice == "2":
                self.add_website()
            elif choice == "3":
                self.delete_website()
            elif choice == "4":
                break
            else:
                print("Invalid choice. Please try again.\n")
                input("Press Enter to continue...")

if __name__ == "__main__":
    # Initialize and run the Authenticator.
    auth = Authenticator()
    auth.run()

Security Considerations

This script implements several security measures to protect sensitive 2FA data:

1. Encryption: All stored data is encrypted using the Fernet symmetric encryption scheme, based on AES in CBC mode with a 128-bit key for encryption and HMAC using SHA256 for authentication.
2. Key Derivation: The encryption key is derived from a user-provided master password using PBKDF2HMAC, which helps protect against brute-force and rainbow table attacks.
3. Salt: A unique salt is generated and stored, adding an extra layer of security to the key derivation process.
4. Secure Password Input: The getpass module securely inputs the master password without displaying it on the screen.
5. Time-based OTP: The use of TOTP ensures that generated passwords are valid only for a short period, typically 30 seconds.

Best Practices for Using This Authenticator

To make the most of this 2FA Authenticator script, users should follow these best practices:

1. Use a strong, unique master password.
2. Keep the master password and the salt file secure and separate from the encrypted data file.
3. Regularly backup the encrypted data file and salt file.
4. Update the script regularly to ensure you have the latest security improvements.
5. Be cautious when entering 2FA secrets, ensuring you’re in a secure environment.
6. Consider using this in conjunction with other security measures, not as a standalone solution.

Conclusion

Creating your own 2FA authenticator with Python provides a practical and powerful way to enhance your understanding of both security measures and Python programming. This project takes you through the essentials of generating one-time passwords (OTPs), securely storing authentication data, and managing 2FA tokens for various websites. By using the pyotp library for OTP generation and the cryptography library for data encryption, you ensure that your authentication information remains secure.

As cyber threats continue to evolve, the importance of robust security practices cannot be overstated. This 2FA authenticator not only serves as a valuable tool for protecting your online accounts but also offers a solid foundation for learning about encryption and secure coding practices. Whether you’re a seasoned developer or just starting out, this project provides a comprehensive look at creating a customizable security solution tailored to your needs.

Moreover, building your own security tools allows you to gain deeper insights into how these systems work, empowering you to implement and advocate for stronger security measures in your personal and professional life. You now have the knowledge and the code to expand and customize this authenticator further. You could integrate it with more sophisticated user interfaces, add additional security features, or even turn it into a web-based application.

In summary, the journey of building this 2FA authenticator highlights the significance of multi-layered security in safeguarding digital assets. Embrace the challenge, enhance your skills, and contribute to a more secure digital world. Now get out there and start coding!