#!/usr/bin/env python3
"""RFC 6238 TOTP generator. Reads base32 secret from the 0600 creds env file
(key AGENT_MS_TOTP_SECRET) and prints the current 6-digit code. The secret is
never printed. Usage: python3 totp_gen.py [step_offset]
  step_offset: integer number of 30s steps to add (default 0). Use -1/+1 to probe skew.
"""
import sys, os, time, hmac, hashlib, base64, struct

CREDS = "/home/claude/.config/amd-agent/credentials.env"


def load_secret():
    with open(CREDS) as f:
        for line in f:
            line = line.strip()
            if line.startswith("AGENT_MS_TOTP_SECRET="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    raise RuntimeError("AGENT_MS_TOTP_SECRET not found")


def totp(secret, when=None, step=30, digits=6, offset=0):
    if when is None:
        when = time.time()
    counter = int(when // step) + offset
    # base32 decode (pad to multiple of 8)
    s = secret.upper().replace(" ", "")
    s += "=" * (-len(s) % 8)
    key = base64.b32decode(s)
    msg = struct.pack(">Q", counter)
    h = hmac.new(key, msg, hashlib.sha1).digest()
    o = h[-1] & 0x0F
    code = (struct.unpack(">I", h[o:o + 4])[0] & 0x7FFFFFFF) % (10 ** digits)
    return str(code).zfill(digits)


if __name__ == "__main__":
    off = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    sec = load_secret()
    print(totp(sec, offset=off))