2-factor authentication terminal app in Python

2fa

#!/usr/bin/env python
import os, os.path, stat, sys, base64

# TOTP lib inlined
import time, hmac, base64, hashlib, struct


def pack_counter(t):
    return struct.pack(">Q", t)


def dynamic_truncate(raw_bytes, length):
    offset = raw_bytes[19] & 0x0f
    decimal_value = ((raw_bytes[offset] & 0x7f) << 24) | (
        (raw_bytes[offset + 1] & 0xff) << 16
    ) | ((raw_bytes[offset + 2] & 0xFF) << 8) | (raw_bytes[offset + 3] & 0xFF)
    return str(decimal_value)[-length:]


def hotp(secret, counter, length=6):
    if type(counter) != bytes: counter = pack_counter(int(counter))
    if type(secret) != bytes: secret = base64.b32decode(secret)
    digest = hmac.new(secret, counter, hashlib.sha1).digest()
    return dynamic_truncate(digest, length)


def totp(secret, length=6):
    """TOTP is implemented as HOTP, but with the counter being the floor of
           the division of the Unix timestamp by 30."""
    counter = pack_counter(round(time.time() // 30))
    return hotp(secret, counter, length)


# end TOTP lib inline

KEYS_PATH = os.path.expanduser("~/.2fa")
if os.path.exists(KEYS_PATH) and stat.S_IMODE(
        os.stat(KEYS_PATH).st_mode) != 0o600:
    print("#########################")
    print("# Insecure keys file!   #")
    print("# Set your .2fa file    #")
    print("# to be read-write by   #")
    print("# you only!             #")
    print("# Command:              #")
    print("# chmod go=,u=rw ~/.2fa #")
    print("#########################")
    print()
    print("Refusing to use insecure key file")
    sys.exit(-1)


class DirtyDict:
    def __init__(self, d=None, **kwargs):
        self.dict = d if d is not None else kwargs
        self.dirty = False

    def __getitem__(self, k):
        return self.dict[k]

    def __setitem__(self, k, v):
        self.dict[k] = v
        self.dirty = True

    def __contains__(self, i):
        return i in self.dict


KEYS = DirtyDict()
if os.path.exists(KEYS_PATH):
    with open(KEYS_PATH) as f:
        for n, line in enumerate(f, 1):
            line = line.strip().split("\t")
            if len(line) != 2:
                print(f"Invalid entry on line {n}: expects `name<tab>key`")
            try:
                KEYS[line[0]] = base64.b32decode(line[1].encode("ascii"))
            except:
                print(f"Invalid entry on line {n}: invalid key")
    KEYS.dirty = False  # reset dirty flag

args = list(sys.argv[1:])


def new_key_wizard():
    print("Hello, friend! Let's get your 2fa account set up!")
    print(
        "First, give this key a name. I would suggest using the name of the site, or maybe"
    )
    print("the account name if you have more than one account on a site.")
    name = input("Key name: ").strip()
    print()
    print("Alright, now that we have a name, let's get that key.")
    print(
        "On the site you're trying to add, you should see a QR code. I can't see that, so"
    )
    print("you should click on the option to manually enter the key.")
    print()
    print(
        "It should give you a bunch of letters and numbers, somewhere in the ballpark of"
    )
    print("16-32 characters is common. Type that in here.")
    entry = True
    while entry:
        key = input("Enter key: ").strip()
        try:
            key = base64.b32decode(key)
            KEYS[name] = key
            entry = False
        except KeyboardInterrupt:
            print()
            return
        except:
            print("That didn't work. Try typing it again?")
    print()
    print("Great! That works!")
    print("Your verification code is", totp(key))
    return


def delete_key_wizard():
    for i, key in enumerate(KEYS.dict, 1):
        print(f"{i}.) {key}")
    entry = True
    while entry:
        try:
            index = int(
                input("Which key do you want to delete?: ").strip()) - 1
            name = list(KEYS.dict.keys())[index]
            entry = False
        except KeyboardInterrupt:
            print()
            return
        except Exception as e:
            print("Invalid index!")
    print(f"Are you sure you want to delete the key for `{name}`?")
    print(f"Once it's gone, there's no going back!")
    c = input(f"Delete key `{name}`?(y/N): ")
    if not c: c = "n"
    if not c.lower()[0] == "y": return
    del KEYS.dict[name]
    print(f"Deleted key `{name}`.")


if len(args) == 1:
    if args[0] in KEYS:  # `2fa reddit`
        print("Your code: " + totp(KEYS[args[0]]))
    elif args[0] in ("wizard",
                     "new"):  # `2fa wizard` or `2fa new` without args
        new_key_wizard()
    elif args[0] in ("del", "delete", "remove", "rm"):
        delete_key_wizard()
    elif args[0] == "help":
        print("2fa - 2 factor authentication app")
        print("Usage:")
        print("2fa <key> - generate code for key `key`")
        print("2fa new - new key wizard")
        print("2fa wizard - direct link to the above")
        print("2fa <del/delete/rm/remove> - key removal wizard")
        print(
            "2fa new <name> [key] - direct key addition. If key is not provided, it will be asked for."
        )
elif len(args) == 2 and args[0] == "new":  # `2fa new reddit`
    entry = True
    while entry:
        key = input("Enter Base32 key: ").strip()
        try:
            key = base64.b32decode(key.encode("ascii"))
            entry = False
        except:
            print("Invalid key!")
    KEYS[args[1]] = key
elif len(args) == 3 and args[0] == "new":  # `2fa new reddit ABCDEFGHIJKLMNOP`
    try:
        key = base64.b32decode(args[2].encode("ascii"))
    except:
        print("Invalid key!")
        sys.exit(1)
    KEYS[args[1]] = key

if KEYS.dirty:
    r = os.path.exists(KEYS_PATH)
    with open(KEYS_PATH, "w") as f:
        for key in KEYS.dict:
            f.write("\t".join(
                [key, base64.b32encode(KEYS[key]).decode("ascii")]))
            f.write("\n")
    if not r:
        os.chmod(KEYS_PATH, 0o600)
    print("Saved!")

README.md

2fa

A 2-factor auth app in the terminal. Coded in Python.

Usage

2fa - 2 factor authentication app
Usage:
2fa <key> - generate code for key `key`
2fa new - new key wizard
2fa wizard - direct link to the above
2fa <del/delete/rm/remove> - key removal wizard
2fa new <name> [key] - direct key addition. If key is not provided, it will be asked for.

Bit of a nitpick but you can simplify DirtyDict and inherit dict features with:

class DirtyDict(dict):
    def __init__(self, d=None, **kwargs):
        self.dict.__init__(self, d, **kwargs)
        self.dirty = False

    def __setitem__(self, k, v):
        dict.__setitem__(self, k, v)
        self.dirty = True

Meh, it's not that important to me. It works, and that's the important part. That would probably be a more Pythonic way to do it (although the real Pythonic way would most likely be subclassing collections.UserDict), but it works and that's the important thing.