OTP Python scripting

.bahrc

# .bashrc or .zshrc
# ...

# Set environment variables
export OTP_VPN="SOMESECRET"
export OTP_EC2="ANOTHERSECRET"

# Set aliases for executing python script from any path
alias otp="~/Documents/sandbox/otp/otp.sh"

# ...

otp.py

import pyotp
import pyperclip
import sys
import os
import time

TIMEOUT = 10

YELLOW = '\033[0;33m'
CLEAR = '\033[0m

# Secret keys dictionary
tokens = {
    "vpn": os.environ.get('OTP_VPN'),
    "ec2": os.environ.get('OTP_EC2')
}

# Check if the required environment variables are set
missing_envs = [key for key, value in tokens.items() if value is None]
if missing_envs:
    misenv = f"OTP_{', '.join(missing_envs).upper()}"
    print(f"Missing or empty environment variable: {misenv} ")
    print(f"Check that SYSTEM ENV is set by typing 'echo ${misenv}")
    print(f"Check that 'export {misenv}='token' ' argument is set in your .bashrc, .zshrc files") 
    sys.exit(1)

# Get the name of the service for the requested OTP
requestedService = sys.argv[1] if len(sys.argv) > 1 else "vpn"

# Validate the requested service
if requestedService not in tokens:
    print("ERROR: Wrong argument was given.")
    print(f"Try the following values: {list(tokens.keys())}")
    sys.exit(1)

# Get the requested token
requestedToken = tokens[requestedService]

# Create a PyOTP TOTP object with the secret key
# Generate the current OTP code
totp = pyotp.TOTP(requestedToken)
code = totp.now()

# Copy the current code to the clipboard using shell commands
pyperclip.copy(code)

# Print the current code and the countdown timer
try:
    for countdown in range(TIMEOUT, -1, -1):
        message = f'Current code for {requested_service}: {code} will be cleared in {YELLOW}{countdown}{CLEAR} seconds'
        print(f'\r{message}', end=' ', flush=True)
        time.sleep(1)
except KeyboardInterrupt:
    # Exit by hitting Ctrl+C
    sys.exit(0)
finally:
    # Clear the terminal output
    print('\r', end='', flush=True)

    # Clear the clipboard
    pyperclip.copy("")

otp.sh

#!/bin/bash

# Change dir to the python code dir
cd ~/Documents/sandbox/otp

# Activate the virtual environment
source otpenv/bin/activate

# Execute the python code
python3 otp.py $1

# Deactivate the virtual environment
deactivate