Virtual MFA device-in-a-script for AWS IAM, RFC6238 TOTP

aws-2fa-session.py

#!/usr/bin/env python

"""
Use TOTP secret with RFC6238 algorithm for AWS MFA login.

$AWS_PROFILE is expected to be set, naming a section in ~/.aws/credentials

That same section must also contain 2 extra values for this script:

[my-awesome-shiny-aws-account]
aws_access_key_id = AKIA....
aws_secret_access_key = ............
# custom values for aws-2fa-session.py:
aws_2fa_device = arn:aws:iam::012345678901:mfa/MyVirtual2FADevice
aws_2fa_secret = RRVTIXW..................................................77XFUCY

The device ARN ("serial number") can be found in IAM console, or with:
    aws iam list-virtual-mfa-devices

The TOTP secret can be scanned off the enrollment QR-code, or extracted from
(a backup of) Google Authenticator app storage, or whatever analogous app you use.
It consists of 64 base-32 characters, encoding 160 bits of entropy.
"""

import argparse
import hashlib, hmac
import json
import logging
import os, sys
from base64 import b32decode
from binascii import hexlify
from configparser import ConfigParser
from time import time

__author__ = "Max Ulidtko"
__email__ = "ulidtko@gmail.com"
__copyright__ = "Copyright © 2024, All Right Reserved"
__license__ = "BSD-3-Clause-Attribution"

try:
    import boto3
    from botocore.exceptions import ClientError

    def get_session_token(mfa_arn, mfa_code):
        try:
            response = boto3.client('sts').get_session_token(
                DurationSeconds=3600 * opts.hours,
                SerialNumber=mfa_arn,
                TokenCode=fmt_code(mfa_code)
            )
            return response['Credentials']
        except ClientError as e:
            # likely a Permission Denied; strip traceback
            print(e, file=sys.stderr)
            sys.exit(1)

except ImportError:
    import subprocess

    def get_session_token(mfa_arn, mfa_code):
        try:
            output = subprocess.check_output(
                args=[
                    'aws', 'sts', 'get-session-token',
                    '--serial-number', str(mfa_arn),
                    '--token-code', fmt_code(mfa_code),
                    '--duration-seconds', str(3600 * opts.hours),
                    '--output', 'json'
                ], shell=False
            )
            return json.loads(output)['Credentials']
        except subprocess.CalledProcessError:
            # awscli prints the error; strip traceback
            sys.exit(1)

argP = argparse.ArgumentParser(
    usage='%(prog)s [..] | eval    -- in bash\n       %(prog)s [..] | source  -- in fish',
    epilog="For those places where DevOps enforce 2FA by a cargo-cult instead of reason.",
    description=__doc__.strip(),
    formatter_class=argparse.RawTextHelpFormatter
)
argP.add_argument('--debug', default=False, action='store_true',
    help="write debug log to stderr")
argP.add_argument('--hours', type=int, default=18,
    help="session duration to request [default 18]")
argP.add_argument('--fish', default=False, action='store_true',
    help="output fish syntax [default bash]")
argP.add_argument('--code-only', default=False, action='store_true',
    help="just output the TOTP code -- skip calling aws sts get-session-token.\n" +
         "this avoids consuming it, so it's still good for e.g. Console login.")

EVAL_TEMPLATE = lambda isfish: """
set -x AWS_ACCESS_KEY_ID     {AccessKeyId}
set -x AWS_SECRET_ACCESS_KEY {SecretAccessKey}
set -x AWS_SESSION_TOKEN     {SessionToken}
""" if isfish else """
export AWS_ACCESS_KEY_ID={AccessKeyId}
export AWS_SECRET_ACCESS_KEY={SecretAccessKey}
export AWS_SESSION_TOKEN={SessionToken}
"""

def main():
    awscreds = read_aws_creds(os.getenv('AWS_PROFILE'))
    totp_device, totp_secret = awscreds['aws_2fa_device'], awscreds['aws_2fa_secret']

    key = b32decode(totp_secret.encode('ascii'))
    now = int(time())
    totp_code = compute_totp(key, now)
    print("TOTP code %s" % fmt_code(totp_code), file=sys.stderr)

    #-- NOTE: AWS STS will invalidate ("use up") the 2FA code on use. In other words,
    #-- after get_session_token() the code will become "consumed" and no more good.
    if opts.code_only:
        return

    session_creds = get_session_token(totp_device, totp_code)
    output = EVAL_TEMPLATE(opts.fish) .format(**session_creds)
    print(output)

def configure_logging(debug: bool):
    if debug:
        logging.basicConfig(level=logging.DEBUG, stream=sys.stderr)

def read_aws_creds(profile):
    ini = ConfigParser()
    ini.read(os.getenv('HOME') + '/.aws/credentials')
    return ini if profile is None else ini[profile]

def fmt_code(code: int) -> str:
    return "%06d" % code

#-------------------------------------------------------#
# The TOTP 2FA protocol is simple, its client side is   #
# implementable in just 12 lines of code (read below).  #
# Here's a high-level diagram how it works in general.  #
#-------------------------------------------------------#
#  ┌────────────┐                        ┌───────────┐  #
#  │ 2FA server │                        │  2FA app  │  #
#  └────────────┘                        └───────────┘  #
#              Enrollment (QR-code usually)             #
# ───────────────────────────────────────────────────── #
#  selects parameters:                                  #
#  - interval [30 secs]                                 #
#  - t0 [unix epoch]                                    #
#  - hmac hash alg [sha1]                               #
#  - num of code digits [6]                             #
#                                                       #
#  generates a secret                                   #
#  (160-bit for sha1)                                   #
#                                                       #
#  sends the secret &                                   #
#  params to client app ──────────► scans the secret &  #
#                                   params off the QR,  #
#                                   saves securely      #
#                                                       #
#  verifies code;  ◄─────────────── generates a code    #
#  enrollment done                  to confirm          #
#                                                       #
#                RFC6238 2FA code auth                  #
# ───────────────────────────────────────────────────── #
#                   ┌───────────┐                       #
#      now()   ◄────┤ UTC clock ├──► (now() - t0)       #
#                   └───────────┘     / interval → ctr  #
#                                                       #
#  receive digits  ◄───────────────  HOTP(secret, ctr)  #
#                                     ▲                 #
#  verify match, ± clock skew,        └─── RFC4226      #
#  against stored secret & now()                        #
#                                                       #
#-------------------------------------------------------#

def compute_totp(hotp_secret: bytes, unix_time: int, sysparam_t0=0, sysparam_ts=30):
    """ https://www.rfc-editor.org/rfc/rfc6238#section-4 """
    t = (unix_time - sysparam_t0) // sysparam_ts #-- // floors to int
    return compute_hotp(hotp_secret, t)

def compute_hotp(hotp_secret: bytes, counter: int, sysparam_digits=6):
    """ https://www.rfc-editor.org/rfc/rfc4226#section-5.3 """
    hmac_sha1 = lambda k, x: hmac.new(k, x, 'sha1').digest()
    ctr = int(counter).to_bytes(8, 'big')
    hash = hmac_sha1(hotp_secret, ctr)
    dtix = 15 & hash[-1]
    dt = 0x7fffffff & int.from_bytes(hash[dtix:dtix+4], 'big')
    return dt % (10 ** sysparam_digits)
# test vectors, https://www.rfc-editor.org/rfc/rfc4226#page-32
assert compute_hotp(b'12345678901234567890', 0) == 755224
assert compute_hotp(b'12345678901234567890', 1) == 287082
assert compute_hotp(b'12345678901234567890', 2) == 359152
assert compute_hotp(b'12345678901234567890', 3) == 969429
assert compute_hotp(b'12345678901234567890', 4) == 338314
assert compute_hotp(b'12345678901234567890', 5) == 254676
assert compute_hotp(b'12345678901234567890', 6) == 287922
assert compute_hotp(b'12345678901234567890', 7) == 162583
assert compute_hotp(b'12345678901234567890', 8) == 399871
assert compute_hotp(b'12345678901234567890', 9) == 520489

if __name__ == "__main__":
    opts = argP.parse_args()
    configure_logging(opts.debug)
    log = logging.getLogger('script')
    sys.exit(main())