Generates 2FA Codes from TOTP URLs

totp.py

import base64
import hmac
import struct
import sys
import time
from urllib.parse import unquote

def generate_hotp_code(secret, counter, hash_algo='sha1', size=6):
    """
    Generates a HOTP code.
    
    secret: The secret string, encoded as base32. Padding is optional.
    
    counter: The integer value of the counter.
    
    hash_algo: The name of the hashing algorithm to use when computing the
               HMAC. SHA1 is used by default.
               
    size: The number of digits to output, 6 by default. If the output would
          be less than this number, than it is zero-padded on the right.
    """
    # Some HOTP secrets can be lower-case, but the standard base32 alphabet is
    # upper case
    secret = secret.upper()

    # HOTP secrets should not include padding by default
    padding = 8 - (len(secret) % 8)
    if padding == 8:
        padding = 0
        
    secret = secret + ('=' * padding)
    secret_bytes = base64.b32decode(secret)

    mac = hmac.new(secret_bytes, digestmod=hash_algo)

    counter_bytes = struct.pack('>q', counter)

    mac.update(counter_bytes)
    digest = mac.digest()

    offset = digest[-1] & 0xf

    trailing_int = (((digest[offset] & 0x7f) << 24) |
                    ((digest[offset + 1] & 0xff) << 16) |
                    ((digest[offset + 2] & 0xff) << 8) |
                    (digest[offset + 3] & 0xff))

    return str(trailing_int % (10**digits)).rjust(6, '0')

def generate_totp_code(secret, hash_algo='sha1', size=6, period=30):
    """
    Generates a TOTP code based upon the current UNIX timestamp.
    
    secret: The secret string, encoded as base32. Padding is optional.        
    
    hash_algo: The name of the hashing algorithm to use when computing the
               HMAC. SHA1 is used by default.
               
    size: The number of digits to output, 6 by default. If the output would
          be less than this number, than it is zero-padded on the right.
          
    period: How long a code is to remain valid in seconds, 30 by default.
    """
    counter = int(time.time() // period)
    return generate_hotp_code(secret, counter, hash_algo, size)

if __name__ == '__main__':
    try:
        url = sys.argv[1]
    except IndexError:
        print(sys.argv[0], 'URL', file=sys.stderr)
        sys.exit(1)

    (otpauth, rest) = url.split('://', 1)
    if otpauth != 'otpauth':
        print('Not OTPAuth URL', file=sys.stderr)
        sys.exit(1)

    (method, rest) = rest.split('/', 1)
    if method != 'totp':
        print('Only TOTP is supported currently', file=sys.stderr)
        sys.exit(1)

    (label, rest) = rest.split('?', 1)
    print('Label:', label)

    issuer = None
    secret = None
    algorithm = 'sha1'
    digits = 6
    period = 30

    params = rest.split('&')
    for param in params:
        (key, value) = param.split('=')
        if key == 'secret':
            secret = value
        elif key == 'issuer':
            issuer = unquote(value)
        elif key == 'algorithm':
            algorithm = unquote(algorithm)
        elif key == 'digits':
            digits = int(digits)
        elif key == 'period':
            period = int(period)

    if issuer is not None:
        print('Issuer:', issuer)

    print(generate_totp_code(secret, algorithm, digits, period))