authenticator.py

# Authors: 
#   Piers Lauder        (getpass()      Original)
#   Guido van Rossum    (getpass()      Windows support and cleanup)
#   Gregory P. Smith    (getpass()      tty support & GetPassWarning)
#   geeksforgeeks.org   (print colors python terminal)
#   codeigo.com         (aes-128 encryption and decryption)
#   Susam Pal           (totp())
#   dev01 Resnova       (Rest all)


import os
import sys


def prRed(skk): print("\033[91m {}\033[00m" .format(skk))
def prCyan(skk): print("\033[96m {}\033[00m" .format(skk))
def prGreen(skk): print("\033[92m {}\033[00m" .format(skk))

try:
    from Crypto.Cipher import AES
except ModuleNotFoundError as me:
    prRed("To resolve, type: \n pip3 install pycryptodome")
    raise
    exit()

from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
from Crypto.Hash import SHA1

import base64
import hmac
import struct
import time




############################################
###  AES-128 Encrpytion and Decryption   ###
############################################
def encrypt(secret_raw_bytes, key):
    # Create an AES cipher object with the key and AES.MODE_ECB mode
    cipher = AES.new(key, AES.MODE_ECB)
    # Pad the secret_raw_bytes and encrypt it
    ciphertext = cipher.encrypt(pad(secret_raw_bytes, AES.block_size))
    return ciphertext
 
def decrypt(ciphertext, key):
    # Create an AES cipher object with the key and AES.MODE_ECB mode
    cipher = AES.new(key, AES.MODE_ECB)
    # Decrypt the ciphertext and remove the padding
    try:
        decrypted_data = unpad(cipher.decrypt(ciphertext), AES.block_size)
    except ValueError as ve:
        if hasattr(ve, 'message'):
            raise ValueError(e.message)
        else:
            raise ValueError(ve)
    
    return decrypted_data


############################################
###            TOTP                      ###
############################################
def hotp(key, counter, digits=6, digest='sha1'):
    key = base64.b32decode(key.upper() + '=' * ((8 - len(key)) % 8))
    counter = struct.pack('>Q', counter)
    mac = hmac.new(key, counter, digest).digest()
    offset = mac[-1] & 0x0f
    binary = struct.unpack('>L', mac[offset:offset+4])[0] & 0x7fffffff
    return str(binary)[-digits:].zfill(digits)

def totp(key, time_step=30, digits=6, digest='sha1'):
    return hotp(key, int(time.time() / time_step), digits, digest)



############################################
###             getpass                  ###
############################################
def fallback_getpass(prompt='Password: ', stream=None):
    if not stream:
        stream = sys.stderr
    return _raw_input(prompt, stream)


def _raw_input(prompt="", stream=None, input=None):
    # This doesn't save the string in the GNU readline history.
    if not stream:
        stream = sys.stderr
    if not input:
        input = sys.stdin
    prompt = str(prompt)
    if prompt:
        try:
            stream.write(prompt)
        except UnicodeEncodeError:
            # Use replace error handler to get as much as possible printed.
            prompt = prompt.encode(stream.encoding, 'replace')
            prompt = prompt.decode(stream.encoding)
            stream.write(prompt)
        stream.flush()
    # NOTE: The Python C API calls flockfile() (and unlock) during readline.
    line = input.readline()
    if not line:
        raise EOFError
    if line[-1] == '\n':
        line = line[:-1]
    return line
    
def win_getpass(prompt='Password: ', stream=None):
    """Prompt for password with echo off, using Windows getwch()."""
    if sys.stdin is not sys.__stdin__:
        return fallback_getpass(prompt, stream)

    prompt = f'{prompt} (Typing will be Hidden) :'
    for c in prompt:
        msvcrt.putwch(c)
    pw = ""
    while 1:
        c = msvcrt.getwch()
        if c == '\r' or c == '\n':
            break
        if c == '\003':
            raise KeyboardInterrupt
        if c == '\b':
            pw = pw[:-1]
        else:
            pw = pw + c
    msvcrt.putwch('\r')
    msvcrt.putwch('\n')
    return pw

def unix_getpass(prompt='Password: ', stream=None):
    """Prompt for a password, with echo turned off.

    Args:
      prompt: Written on stream to ask for the input.  Default: 'Password: '
      stream: A writable file object to display the prompt.  Defaults to
              the tty.  If no tty is available defaults to sys.stderr.
    Returns:
      The seKr3t input.
    Raises:
      EOFError: If our input tty or stdin was closed.
      GetPassWarning: When we were unable to turn echo off on the input.

    Always restores terminal settings before returning.
    """
    passwd = None
    with contextlib.ExitStack() as stack:
        try:
            # Always try reading and writing directly on the tty first.
            fd = os.open('/dev/tty', os.O_RDWR|os.O_NOCTTY)
            tty = io.FileIO(fd, 'w+')
            stack.enter_context(tty)
            input = io.TextIOWrapper(tty)
            stack.enter_context(input)
            if not stream:
                stream = input
        except OSError:
            # If that fails, see if stdin can be controlled.
            stack.close()
            try:
                fd = sys.stdin.fileno()
            except (AttributeError, ValueError):
                fd = None
                passwd = fallback_getpass(prompt, stream)
            input = sys.stdin
            if not stream:
                stream = sys.stderr

        if fd is not None:
            try:
                old = termios.tcgetattr(fd)     # a copy to save
                new = old[:]
                new[3] &= ~termios.ECHO  # 3 == 'lflags'
                tcsetattr_flags = termios.TCSAFLUSH
                if hasattr(termios, 'TCSASOFT'):
                    tcsetattr_flags |= termios.TCSASOFT
                try:
                    prompt = f'{prompt} (Typing will be Hidden) :'
                    termios.tcsetattr(fd, tcsetattr_flags, new)
                    passwd = _raw_input(prompt, stream, input=input)
                finally:
                    termios.tcsetattr(fd, tcsetattr_flags, old)
                    stream.flush()  # issue7208
            except termios.error:
                if passwd is not None:
                    # _raw_input succeeded.  The final tcsetattr failed.  Reraise
                    # instead of leaving the terminal in an unknown state.
                    raise
                # We can't control the tty or stdin.  Give up and use normal IO.
                # fallback_getpass() raises an appropriate warning.
                if stream is not input:
                    # clean up unused file objects before blocking
                    stack.close()
                passwd = fallback_getpass(prompt, stream)

        stream.write('\n')
        return passwd

import contextlib
import io
try:
    import msvcrt
except ImportError:
    try:
        import termios
        # it's possible there is an incompatible termios from the
        # McMillan Installer, make sure we have a UNIX-compatible termios
        termios.tcgetattr, termios.tcsetattr
    except (ImportError, AttributeError):    
        getpass = fallback_getpass
    else:
        getpass = unix_getpass          
else:
    getpass = win_getpass


############################################
###               Main                   ###
############################################
def cal_tabs_len(strlen, max_tabs=4, tab_chars = 8):
    max_tabs = max_tabs
    tablen = 1 + max_tabs - int((strlen+2)/tab_chars)
    tablen = 0 if tablen < 0 else tablen
    return tablen

def valid_secret_key(key):
    # A  remainder of 1 or 3 or 6 characters is not a legal base32 encoding
    extra_after_mod_8 =  len(key) % 8
    if extra_after_mod_8 in [1,3,6]:
        return False    # not valid remainders
    else:
        return True

def create_key(txt):
    # Find the hash of the tezt
    hash = SHA1.new(txt.strip().encode('ASCII'))

    # Key-length accepted: 16, 24, and 32 bytes.
    keybytes = hash.digest()[:16]  
    return keybytes

def store(site, user, secret, filename='secretStore.txt', prefix='<TOTP_REC_START>',suffix='<TOTP_REC_END>'):
    rec = f'\r\n{prefix} {site} {user} {secret} {suffix}'
    # get the current working directory
    cwd = os.getcwd()

    file_path = os.path.join(cwd,filename)

    # Open the file in write mode
    with open(file_path, 'a') as file:
        # Write content to the file
        file.write(rec)
        prGreen("Saved the secret. Remember the paraphrase used for encryption")
 
def read_from_file(filename='secretStore.txt'):
    # get the current working directory
    cwd = os.getcwd()
    file_path = os.path.join(cwd,filename)
    contents = ""
    try:
        with open(file_path, 'r') as fp:
            contents = fp.readlines()
    except FileNotFoundError as fe:
        prRed(f'No record')
        exit()
        
    return contents
        
def get_secret_key_from_store(site, user, key, filename='secretStore.txt', prefix='<TOTP_REC_START>',suffix='<TOTP_REC_END>'):
    # initialize
    decrypted_Secret = ""
    found_secrtky = False
    status_msg = ""
    

    records = list_rec()
    for tokens in records:
        if(tokens[1]==site and tokens[2]==user):
            found_secrtky = True
            try:                        
                encr = bytes.fromhex(tokens[3]) # Convert hex string to bytes object
                # Decryption
                decrypted_Secret = decrypt(encr, key)
            except ValueError as ve:
                # print(ve)
                if hasattr(ve, 'message'):
                    status_msg = ve.message
                else:
                    status_msg = ve
                continue
                
            # if here,  no errors in decryption
            if isinstance(decrypted_Secret, bytes):
                 decrypted_Secret = decrypted_Secret.decode("ASCII")
            return decrypted_Secret, True, ""
                    
    return decrypted_Secret, found_secrtky, status_msg # Could not get the secret

def list_rec(filename='secretStore.txt', prefix='<TOTP_REC_START>',suffix='<TOTP_REC_END>',printout=False):
    contents = read_from_file(filename)
    count = 0
    records = []

    if(printout):
        print("\nNo.", "\t", "Website", "\t\t\t\t", "User")

    for line in contents:
        if len(line) > (len(prefix) + len(suffix)):
            tokens = line.split()
            if(tokens[0]==prefix and tokens[4]==suffix):
                records.append(tokens)
                if printout:
                    tablen = cal_tabs_len(len(tokens[1]))
                    print(count, "\t", tokens[1], "\t" * tablen, tokens[2])
                    count +=1
    return records


def get_token(site, user):
    # paraphrase = input('Enter the paraphrase to retrieve the token:')  
    paraphrase = getpass('Enter the paraphrase to retrieve the token')
    key = create_key(paraphrase)

    secret_key, found_rec, status_msg = get_secret_key_from_store(site, user, key)
    if(found_rec and len(secret_key) and len(status_msg) == 0):
        token = totp(secret_key)
        prGreen (f'token = {token}')
        print("Press any key to show token")
        print("Press 'q' to exit")
        for inp in sys.stdin:
            inp = inp.strip()
            if inp=='Q' or inp=='q' :
                break # quit
            else:
                token = totp(secret_key)
                prGreen (f'token = {token}')
                print("Press any key to show token")
                print("Press 'q' to exit")
            
    elif not found_rec:
        prCyan("Could not find the record")
    else:
        # print("Error",status_msg)
        prRed(f'Error {status_msg}')
 
def select_rec(records):
    print("Press 'q' for Exit \n")
    try:
        rec_selectd = int(input('Select a No. to show token:'))
    except Exception as e:
        prRed(f'Error Invalid No.')
        exit()

    if(rec_selectd>=0 and rec_selectd<len(records)):
        get_token(records[rec_selectd][1],records[rec_selectd][2])
    else:
        prRed(f'Error Invalid No.')
 
def store_rec():
    # paraphrase = input('Enter the paraphrase to encrypt the secret:')
    paraphrase = getpass('Enter the paraphrase to encrypt the secret')
    
    website = input('Enter the Website:').strip()
    user = input('Enter the User:').strip()
    # secret_raw_txt = input('Enter the Secret Key:')    
    secret_raw_txt = getpass('Enter the Secret Key')
    if not valid_secret_key(secret_raw_txt):
        prRed("Not a valid Secret Key, Not allowed remainder length after mod 8 are [1,3,6]")
        exit()
    
    
    key = create_key(paraphrase)   

    secret_raw_bytes = secret_raw_txt.strip().encode('ASCII')

    # Encryption
    encrypted_secret_key = encrypt(secret_raw_bytes, key)

    store(website, user, encrypted_secret_key.hex()) 

def get_rec():
    website = input('Enter the Website:').strip()
    user = input('Enter the User:').strip()
    get_token(website,user)

def select_menu():
    print("1.", "\t", "List")
    print("2.", "\t", "Store")
    print("3.", "\t", "Get")
    print("4.", "\t", "Exit")
    try:
        menu_selectd = int(input('Select an option:'))
    except Exception as e:
        prRed(f'Error Invalid No.')
        exit()
   
    if(menu_selectd>=1 and menu_selectd<=4):
        if(menu_selectd==1):
            select_rec(list_rec(printout=True))
        elif(menu_selectd==2):
            store_rec()        
        elif(menu_selectd==3): 
            get_rec()
        elif(menu_selectd==4): 
            exit()
         
    else:
        prRed(f'Error Invalid No.')
            

## ENCRYPT and STORE the SECRET of TOTP for user from website
## GET the TOKEN CODE FROM File after decrypting SECRET of User from website 

##   Authenticator App
## ######################
## This app is to show the token code mostly used in multi factor authentication. 
## The token code shown is based in totp and it changes every 30s.
## With this app, user can store the secret key from which the token code is generated.
## The secret key is encrypted and then stored.
## This app can manage token code for several website and users.
## A paraphrase is used to encrypt the secret key, 
## so user should remeber the paraphrase for every token to be generated.
 

select_menu()