#!/usr/bin/env python3
"""Import TOTP credentials from otp-keys-export.xlsx into a VivoKey Apex OTP
(ykneo-oath) applet over PC/SC.

Protocol verified against the apex-totp (com.vivokey.otp.YkneoOath) and
javacard-memory applet sources. See
docs/superpowers/specs/2026-07-08-apex-otp-import-design.md.
"""

import argparse
import base64
import csv
import datetime
import getpass
import hashlib
import hmac
import os
import sys
from dataclasses import dataclass, field

# ykneo-oath enforces a 14-byte minimum HMAC key; the Yubico host right-zero-pads
# short secrets to this length. Padding is HMAC-identical (HMAC zero-pads the key
# to the block size regardless), so TOTP output is unchanged.
HMAC_MIN_KEY = 14
# The applet stores the HMAC key in a 64-byte ipad/opad buffer and rejects longer.
MAX_KEY = 64
MAX_NAME = 64
DEFAULT_PERIOD = 30

# ykneo-oath tags / type|algorithm nibbles (verified against YkneoOath/OathObj).
TAG_NAME = 0x71
TAG_NAME_LIST = 0x72
TAG_KEY = 0x73
TAG_CHALLENGE = 0x74
TAG_RESPONSE = 0x75
TAG_VERSION = 0x79
TAG_ALGORITHM = 0x7B

HOTP_TYPE = 0x10
TOTP_TYPE = 0x20
HMAC_MASK = 0x0F   # low nibble of a type|algo byte selects the HMAC algorithm
HMAC_SHA1 = 0x01
HMAC_SHA256 = 0x02
HMAC_SHA512 = 0x03  # rejected by the applet; used only to flag bad input


# ---------------------------------------------------------------------------
# Credential field helpers
# ---------------------------------------------------------------------------

def decode_secret(secret: str) -> bytes:
    """Base32-decode an OATH secret, tolerating lowercase, spaces and missing
    padding."""
    s = secret.strip().replace(" ", "").upper()
    s += "=" * ((8 - len(s) % 8) % 8)
    return base64.b32decode(s)


def pad_key(raw: bytes) -> bytes:
    """Right-zero-pad a raw HMAC key to the applet's 14-byte minimum."""
    if len(raw) >= HMAC_MIN_KEY:
        return raw
    return raw + b"\x00" * (HMAC_MIN_KEY - len(raw))


def format_cred_id(issuer: str, account: str, period: int = DEFAULT_PERIOD) -> str:
    """Build the on-applet credential name: '[period/]issuer:account'.
    The period prefix is only added when it differs from the 30s default."""
    label = f"{issuer}:{account}" if issuer else account
    if period and period != DEFAULT_PERIOD:
        return f"{period}/{label}"
    return label


@dataclass
class Credential:
    issuer: str
    account: str
    secret: str                 # Base32 text
    oath_type: int = TOTP_TYPE
    algorithm: int = HMAC_SHA1
    digits: int = 6
    period: int = DEFAULT_PERIOD
    source: str = ""

    @property
    def cred_id(self) -> str:
        return format_cred_id(self.issuer, self.account, self.period)

    @property
    def type_algo(self) -> int:
        return self.oath_type | self.algorithm

    def raw_key(self) -> bytes:
        """Base32-decode the secret (raises on invalid Base32)."""
        return decode_secret(self.secret)

    def put_data(self) -> bytes:
        return build_put_data(self.cred_id, self.type_algo, self.digits,
                              pad_key(self.raw_key()))


_TYPE_MAP = {"TOTP": TOTP_TYPE, "HOTP": HOTP_TYPE}
_ALGO_MAP = {"SHA1": HMAC_SHA1, "SHA256": HMAC_SHA256, "SHA512": HMAC_SHA512}


def _match_columns(header):
    """Map the OTP Keys header row to column indices by keyword (order-agnostic)."""
    col = {}
    for i, h in enumerate(header):
        if h is None:
            continue
        k = str(h).strip().lower()
        if "issuer" in k:            col.setdefault("issuer", i)
        elif "account" in k:         col.setdefault("account", i)
        elif "secret" in k:          col.setdefault("secret", i)
        elif k.startswith("type"):   col.setdefault("type", i)
        elif "algorithm" in k:       col.setdefault("algorithm", i)
        elif "digit" in k:           col.setdefault("digits", i)
        elif "period" in k:          col.setdefault("period", i)
        elif "source" in k:          col.setdefault("source", i)
    return col


def load_credentials(xlsx_path, sheet="OTP Keys"):
    """Load credentials from the curated columns of the export spreadsheet."""
    from openpyxl import load_workbook
    wb = load_workbook(xlsx_path, read_only=True, data_only=True)
    ws = wb[sheet]
    rows = ws.iter_rows(values_only=True)
    header = next(rows)
    col = _match_columns(header)

    def cell(row, key, default=None):
        i = col.get(key)
        return row[i] if i is not None and i < len(row) else default

    def as_int(v, default):
        try:
            return int(v)
        except (TypeError, ValueError):
            return default

    creds = []
    for row in rows:
        if row is None:
            continue
        secret = cell(row, "secret")
        issuer = cell(row, "issuer") or ""
        account = cell(row, "account") or ""
        if not secret and not issuer and not account:
            continue  # skip blank rows
        creds.append(Credential(
            issuer=str(issuer).strip(),
            account=str(account).strip(),
            secret=str(secret or "").strip(),
            oath_type=_TYPE_MAP.get(str(cell(row, "type", "TOTP")).strip().upper(), TOTP_TYPE),
            algorithm=_ALGO_MAP.get(str(cell(row, "algorithm", "SHA1")).strip().upper(), HMAC_SHA1),
            digits=as_int(cell(row, "digits"), 6),
            period=as_int(cell(row, "period"), DEFAULT_PERIOD),
            source=str(cell(row, "source", "") or "").strip(),
        ))
    wb.close()
    return creds


def validate_input(creds):
    """Return a list of human-readable problems that must be fixed in the source
    data before importing. An empty list means the input is clean."""
    problems = []
    seen = {}
    for c in creds:
        cid = c.cred_id
        # duplicate credential ids (OATH keys by name, so a dup silently clobbers)
        if cid in seen:
            problems.append(
                f"Duplicate credential name '{cid}': {seen[cid]} and {c.source} "
                f"map to the same applet entry (different secrets would collide)."
            )
        else:
            seen[cid] = c.source

        # invalid characters in issuer/account. The applet stores raw bytes, but
        # control characters and the OATH name separators (':' / 'period/')
        # corrupt display and round-trip in authenticator apps.
        for fname, fval in (("issuer", c.issuer), ("account", c.account)):
            ctrl = next((ch for ch in fval if ord(ch) < 0x20 or ord(ch) == 0x7F), None)
            if ctrl is not None:
                problems.append(
                    f"Control character (0x{ord(ctrl):02X}) in {fname} for "
                    f"'{cid}' ({c.source}).")
            if ":" in fval:
                problems.append(
                    f"Character ':' in {fname} for '{cid}' ({c.source}); it "
                    f"collides with the issuer:account separator.")
            if "/" in fval:
                problems.append(
                    f"Character '/' in {fname} for '{cid}' ({c.source}); it "
                    f"collides with the 'period/' name prefix.")

        # composed name length (the applet's hard 64-byte limit)
        name_bytes = len(cid.encode("utf-8"))
        if name_bytes > MAX_NAME:
            problems.append(
                f"Composed name too long ({name_bytes}>{MAX_NAME} bytes): "
                f"'{cid}' ({c.source}); shorten the issuer or account.")

        # digits / period sanity
        if c.digits not in (6, 7, 8):
            problems.append(
                f"Unusual digit count {c.digits} for '{cid}' ({c.source}); "
                f"OATH codes are 6-8 digits.")
        if not isinstance(c.period, int) or c.period <= 0:
            problems.append(f"Invalid period {c.period!r} for '{cid}' ({c.source}).")

        # algorithm supported by the applet
        if c.algorithm not in (HMAC_SHA1, HMAC_SHA256):
            problems.append(
                f"Unsupported algorithm 0x{c.algorithm:02X} (SHA512?) for '{cid}' "
                f"({c.source}); the applet accepts only SHA1/SHA256."
            )

        # secret decodes and fits
        try:
            raw = c.raw_key()
        except Exception as e:
            problems.append(f"Invalid Base32 secret for '{cid}' ({c.source}): {e}.")
            continue
        if len(raw) == 0:
            problems.append(f"Empty secret for '{cid}' ({c.source}).")
        elif len(raw) > MAX_KEY:
            problems.append(
                f"Key too long ({len(raw)}>{MAX_KEY} bytes) for '{cid}' ({c.source})."
            )
    return problems


# ---------------------------------------------------------------------------
# BER-TLV helpers
# ---------------------------------------------------------------------------

def tlv(tag: int, value: bytes) -> bytes:
    """Encode a single-byte-tag BER-TLV with 1/2/3-byte length forms."""
    n = len(value)
    if n < 0x80:
        length = bytes([n])
    elif n <= 0xFF:
        length = bytes([0x81, n])
    else:
        length = bytes([0x82, (n >> 8) & 0xFF, n & 0xFF])
    return bytes([tag]) + length + value


_HMAC_HASH = {HMAC_SHA1: "sha1", HMAC_SHA256: "sha256"}


def derive_key(password: str, salt: bytes) -> bytes:
    """Derive the 16-byte OATH access key from a password, per the ykneo-oath
    convention: PBKDF2-HMAC-SHA1(password, salt = device id from SELECT, 1000
    rounds). PBKDF2 always uses SHA1 here, regardless of the challenge algo."""
    return hashlib.pbkdf2_hmac("sha1", password.encode("utf-8"), salt, 1000, 16)


def compute_response(key: bytes, challenge: bytes, algo: int = HMAC_SHA1) -> bytes:
    """HMAC(key, challenge) using the algorithm named by the SELECT ALGORITHM
    tag (SHA1 in practice)."""
    return hmac.new(key, challenge, _HMAC_HASH[algo]).digest()


def build_put_data(name: str, type_algo: int, digits: int, key: bytes) -> bytes:
    """PUT payload: NAME(0x71) + KEY(0x73, [type|algo, digits, key...]).
    `key` must already be padded (see pad_key). PROPERTY/IMF are omitted
    (TOTP, no touch)."""
    return (
        tlv(TAG_NAME, name.encode("utf-8"))
        + tlv(TAG_KEY, bytes([type_algo, digits]) + key)
    )


def build_validate_data(response: bytes, client_challenge: bytes) -> bytes:
    """VALIDATE payload: RESPONSE(0x75, our HMAC of the device challenge) +
    CHALLENGE(0x74, our fresh challenge for mutual auth)."""
    return tlv(TAG_RESPONSE, response) + tlv(TAG_CHALLENGE, client_challenge)


def build_set_code_data(key: bytes, challenge: bytes, algo: int = HMAC_SHA1) -> bytes:
    """SET CODE payload: KEY(0x73, [type|algo, 16-byte key]) + CHALLENGE(0x74) +
    RESPONSE(0x75, HMAC(key, challenge)) so the applet can confirm the host
    derives the same key."""
    return (
        tlv(TAG_KEY, bytes([TOTP_TYPE | algo]) + key)
        + tlv(TAG_CHALLENGE, challenge)
        + tlv(TAG_RESPONSE, compute_response(key, challenge, algo))
    )


INS_SEND_REMAINING = 0xA5


class CardChannel:
    """Thin wrapper over a pyscard connection that transceives APDUs and
    transparently follows 61xx status words with SEND REMAINING."""

    def __init__(self, connection):
        self.conn = connection

    def send(self, cla, ins, p1, p2, data: bytes = b"", le=None):
        """Transmit one APDU; returns (response_bytes, sw) with sw = (sw1<<8)|sw2.
        Reassembles chained 61xx responses."""
        apdu = [cla, ins, p1, p2]
        if data:
            apdu += [len(data)] + list(data)
        if le is not None:
            apdu += [le]
        resp, sw1, sw2 = self.conn.transmit(apdu)
        out = bytes(resp)
        while sw1 == 0x61:
            resp, sw1, sw2 = self.conn.transmit(
                [0x00, INS_SEND_REMAINING, 0x00, 0x00, sw2])
            out += bytes(resp)
        return out, (sw1 << 8) | sw2


def parse_tlvs(data: bytes):
    """Parse a concatenation of single-byte-tag BER-TLVs into an ordered list
    of (tag, value) tuples. Preserves repeated tags (e.g. LIST 0x72 entries)."""
    out = []
    i = 0
    n = len(data)
    while i < n:
        tag = data[i]
        i += 1
        first = data[i]
        if first < 0x80:
            length = first
            i += 1
        elif first == 0x81:
            length = data[i + 1]
            i += 2
        elif first == 0x82:
            length = (data[i + 1] << 8) | data[i + 2]
            i += 3
        else:
            raise ValueError(f"unsupported TLV length byte 0x{first:02X}")
        out.append((tag, data[i:i + length]))
        i += length
    return out


# ---------------------------------------------------------------------------
# OATH session (ykneo-oath) over a CardChannel
# ---------------------------------------------------------------------------

# Applet-instance AID first, then the shorter Yubico prefixes (the applet
# prefix-selects, so any of these reaches it).
OATH_AIDS = [
    bytes.fromhex("A0000005272101014150455801"),
    bytes.fromhex("A000000527210101"),
    bytes.fromhex("A0000005272101"),
]

INS_SELECT = 0xA4
INS_PUT = 0x01
INS_SET_CODE = 0x03
INS_RESET = 0x04
INS_LIST = 0xA1
INS_VALIDATE = 0xA3

_SW_MESSAGES = {
    0x9000: "OK",
    0x6A84: "applet full (no space)",
    0x6982: "authentication required",
    0x6A80: "wrong data / incorrect password",
    0x6984: "no such object",
    0x6D00: "instruction not supported (wrong applet?)",
    0x6E00: "class not supported",
}


def sw_message(sw: int) -> str:
    return _SW_MESSAGES.get(sw, f"error SW={sw:04X}")


class OathError(Exception):
    def __init__(self, sw: int, message: str = ""):
        self.sw = sw
        super().__init__(message or sw_message(sw))


class OathSession:
    """ykneo-oath session: SELECT (with AID fallback), password VALIDATE, LIST,
    PUT. All wire encoding verified against the applet source."""

    def __init__(self, channel: "CardChannel", rng=os.urandom, aids=OATH_AIDS):
        self.ch = channel
        self.rng = rng
        self.aids = aids
        self.version = None
        self.device_id = None
        self.challenge = None
        self.algorithm = HMAC_SHA1
        self._key = None

    def select(self):
        last_sw = None
        for aid in self.aids:
            data, sw = self.ch.send(0x00, INS_SELECT, 0x04, 0x00, data=aid, le=0x00)
            if sw == 0x9000 and data and data[0] == TAG_VERSION:
                self._parse_select(data)
                return self
            last_sw = sw
        raise OathError(last_sw or 0, "OATH applet not found (SELECT failed)")

    def _parse_select(self, data):
        self.challenge = None
        for tag, val in parse_tlvs(data):
            if tag == TAG_VERSION:
                self.version = tuple(val)
            elif tag == TAG_NAME:
                self.device_id = val
            elif tag == TAG_CHALLENGE:
                self.challenge = val
            elif tag == TAG_ALGORITHM:
                # applet reports the full type|algo byte; keep only the algo nibble
                self.algorithm = (val[0] & HMAC_MASK) if val else HMAC_SHA1

    @property
    def locked(self) -> bool:
        return self.challenge is not None

    def validate(self, password: str):
        """Mutually authenticate against a password-protected applet."""
        if not self.locked:
            return
        key = derive_key(password, self.device_id)
        our_response = compute_response(key, self.challenge, self.algorithm)
        client_challenge = self.rng(8)
        data, sw = self.ch.send(
            0x00, INS_VALIDATE, 0x00, 0x00,
            data=build_validate_data(our_response, client_challenge), le=0x00)
        if sw != 0x9000:
            raise OathError(sw, "authentication failed (incorrect password?)")
        device_response = dict(parse_tlvs(data)).get(TAG_RESPONSE)
        if device_response != compute_response(key, client_challenge, self.algorithm):
            raise OathError(sw, "mutual authentication check failed")
        self._key = key
        return

    def list_names(self):
        data, sw = self.ch.send(0x00, INS_LIST, 0x00, 0x00, le=0x00)
        if sw != 0x9000:
            raise OathError(sw, "LIST failed")
        names = []
        for tag, val in parse_tlvs(data):
            if tag == TAG_NAME_LIST and len(val) >= 1:
                names.append(val[1:].decode("utf-8", "replace"))  # drop type byte
        return names

    def put(self, cred: "Credential"):
        data, sw = self.ch.send(0x00, INS_PUT, 0x00, 0x00, data=cred.put_data())
        if sw != 0x9000:
            raise OathError(sw)
        return

    def set_password(self, password: str):
        """SET CODE: install an access password (SHA1) on the applet. Requires
        the applet be unlocked, or already authenticated this session."""
        key = derive_key(password, self.device_id)
        challenge = self.rng(8)
        data, sw = self.ch.send(
            0x00, INS_SET_CODE, 0x00, 0x00,
            data=build_set_code_data(key, challenge, HMAC_SHA1))
        if sw != 0x9000:
            raise OathError(sw, "SET CODE failed")
        return

    def reset(self):
        """RESET: wipe all credentials and any password (P1P2 = 0xDEAD).
        Unlike Yubico's OATH, this applet gates RESET behind authentication when
        a password is set, so a locked applet must be VALIDATEd first."""
        data, sw = self.ch.send(0x00, INS_RESET, 0xDE, 0xAD)
        if sw != 0x9000:
            raise OathError(sw, f"RESET failed ({sw_message(sw)})")
        return


# ---------------------------------------------------------------------------
# Optional javacard-memory applet (for capacity checks / profiling)
# ---------------------------------------------------------------------------

MEMORY_AID = bytes.fromhex("A0000008466D656D6F727901")
INS_GET_BATCH = 0x01


def parse_memory(data: bytes) -> dict:
    """Parse the 12-byte SELECT response of the javacard-memory applet."""
    return {
        "available_persistent": int.from_bytes(data[0:4], "big"),
        "total_persistent": int.from_bytes(data[4:8], "big"),
        "transient_reset": int.from_bytes(data[8:10], "big"),
        "transient_deselect": int.from_bytes(data[10:12], "big"),
    }


class MemorySession:
    """Optional javacard-memory applet. select() returns None if absent."""

    def __init__(self, channel: "CardChannel"):
        self.ch = channel

    def select(self):
        data, sw = self.ch.send(0x00, INS_SELECT, 0x04, 0x00, data=MEMORY_AID, le=0x00)
        if sw != 0x9000 or len(data) < 12:
            return None
        return parse_memory(data)

    def available(self):
        """SELECT the applet and return available persistent bytes (or None)."""
        m = self.select()
        return m["available_persistent"] if m else None

    def batch(self):
        data, sw = self.ch.send(0x00, INS_GET_BATCH, 0x00, 0x00, le=0x00)
        return data.hex() if sw == 0x9000 else ""


# ---------------------------------------------------------------------------
# CLI / orchestration (hardware I/O)
# ---------------------------------------------------------------------------

DEFAULT_INPUT = "otp-keys-export.xlsx"
DEFAULT_MEMLOG = "apex-otp-memlog.csv"
MEM_PER_CRED_ESTIMATE = 250  # rough EEPROM per credential (3x64B arrays + overhead)


def get_uid(connection) -> str:
    """Contactless/contact UID via the PC/SC GET DATA pseudo-APDU."""
    try:
        resp, sw1, sw2 = connection.transmit([0xFF, 0xCA, 0x00, 0x00, 0x00])
        if sw1 == 0x90:
            return bytes(resp).hex().upper()
    except Exception:
        pass
    return ""


def _readers_with_cards():
    from smartcard.System import readers
    out = []
    for r in readers():
        try:
            conn = r.createConnection()
            conn.connect()
            out.append((str(r), conn))
        except Exception:
            pass
    return out


def pick_reader(reader_filter=None):
    """Return (reader_name, connection) for the single Apex holding an OATH
    applet, prompting if several are present. (None, None) if none found."""
    candidates = _readers_with_cards()
    if reader_filter:
        candidates = [(n, c) for (n, c) in candidates
                      if reader_filter.lower() in n.lower()]
    oath = []
    for name, conn in candidates:
        try:
            OathSession(CardChannel(conn)).select()
            oath.append((name, conn))
        except Exception:
            try:
                conn.disconnect()
            except Exception:
                pass
    if not oath:
        print("No Apex OTP applet found on any reader with a card.")
        print("Place your Apex on a reader (or insert the test card) and re-run.")
        return None, None
    if len(oath) == 1:
        return oath[0]
    print("Multiple readers present an OATH applet:")
    for i, (name, _) in enumerate(oath):
        print(f"  [{i}] {name}")
    idx = int(input("Choose target reader index: ").strip())
    chosen = oath[idx]
    for j, (_, conn) in enumerate(oath):
        if j != idx:
            try:
                conn.disconnect()
            except Exception:
                pass
    return chosen


class ConflictPolicy:
    """Decide overwrite/skip for a credential already on the applet."""

    def __init__(self, args):
        self.always = ("overwrite" if args.overwrite_all
                       else "skip" if args.skip_existing else None)

    def decide(self, name):
        if self.always:
            return self.always
        if not sys.stdin.isatty():
            raise SystemExit(
                f"'{name}' already exists on the applet and no "
                f"--overwrite-all/--skip-existing was given (non-interactive).")
        while True:
            ans = input(f"  '{name}' exists. [o]verwrite / [s]kip / "
                        f"[a]ll overwrite / [x] skip all: ").strip().lower()
            if ans == "o":
                return "overwrite"
            if ans == "s":
                return "skip"
            if ans == "a":
                self.always = "overwrite"
                return "overwrite"
            if ans == "x":
                self.always = "skip"
                return "skip"


def _put_apdu_hex(cred):
    data = cred.put_data()
    return (bytes([0x00, INS_PUT, 0x00, 0x00, len(data)]) + data).hex().upper()


def parse_args(argv=None):
    p = argparse.ArgumentParser(
        description="Import TOTP credentials from the export spreadsheet into a "
                    "VivoKey Apex OTP (ykneo-oath) applet over PC/SC.")
    p.add_argument("--input", default=DEFAULT_INPUT,
                   help=f"export .xlsx (default: {DEFAULT_INPUT})")
    p.add_argument("--reader", help="substring to select a specific PC/SC reader")
    p.add_argument("--dry-run", action="store_true",
                   help="validate and print the APDUs; no card I/O")
    p.add_argument("--list", action="store_true",
                   help="read-only: list credentials already on the applet")
    g = p.add_mutually_exclusive_group()
    g.add_argument("--overwrite-all", action="store_true",
                   help="overwrite any credential that already exists")
    g.add_argument("--skip-existing", action="store_true",
                   help="skip any credential that already exists")
    p.add_argument("--profile-memory", action="store_true",
                   help="record per-credential EEPROM consumption to the mem-log")
    p.add_argument("--memlog", default=DEFAULT_MEMLOG,
                   help=f"profiling CSV path (default: {DEFAULT_MEMLOG})")
    p.add_argument("--password", help="applet password (else prompted if locked)")
    p.add_argument("--reset", action="store_true",
                   help="wipe all credentials and any password, then exit")
    p.add_argument("--set-password", metavar="PASS",
                   help="after importing, set this applet password "
                        "(only if the applet has none)")
    return p.parse_args(argv)


def _reauth(oath, pw):
    oath.select()
    if oath.locked:
        oath.validate(pw)


def main(argv=None):
    args = parse_args(argv)

    # --reset is a standalone admin action; it needs no input spreadsheet.
    if args.reset:
        reader, connection = pick_reader(args.reader)
        if connection is None:
            return 1
        oath = OathSession(CardChannel(connection))
        oath.select()
        if oath.locked:
            # This applet requires authentication before RESET (unlike Yubico).
            pw = args.password or getpass.getpass("Applet password (required to reset): ")
            oath.validate(pw)
        oath.reset()
        print(f"Applet reset on '{reader}' (all credentials and any password wiped).")
        return 0

    # 1. Load + validate BEFORE touching any reader.
    creds = load_credentials(args.input)
    problems = validate_input(creds)
    if problems:
        print(f"Input validation FAILED ({len(problems)} problem(s)) - "
              f"fix the spreadsheet and re-run:")
        for p in problems:
            print("  -", p)
        return 2
    print(f"Loaded and validated {len(creds)} credentials from {args.input}.")

    # 2. Dry-run: no card.
    if args.dry_run:
        print("\n--dry-run: PUT APDUs that WOULD be sent (no card I/O):")
        for c in creds:
            print(f"  {c.cred_id}")
            print(f"    {_put_apdu_hex(c)}")
        return 0

    # 3. Reader/card detection.
    reader, connection = pick_reader(args.reader)
    if connection is None:
        return 1
    channel = CardChannel(connection)

    # 4. Capacity pre-check (selects the memory applet, before the OATH session).
    if not args.list:
        avail = MemorySession(channel).available()
        if avail is not None:
            est = MEM_PER_CRED_ESTIMATE * len(creds)
            msg = (f"Memory applet: {avail} bytes free "
                   f"(rough estimate ~{est} bytes for {len(creds)} credentials).")
            print("WARNING: " + msg + " They may not all fit." if avail < est else msg)

    # 5. OATH session.
    oath = OathSession(channel)
    oath.select()
    applet_was_locked = oath.locked
    print(f"OATH applet on '{reader}': v{'.'.join(map(str, oath.version))}, "
          f"{'LOCKED' if oath.locked else 'unlocked'}.")
    pw = None
    if oath.locked:
        pw = args.password or getpass.getpass("Applet password: ")
        oath.validate(pw)
        print("Authenticated OK.")
    existing = set(oath.list_names())

    if args.list:
        print(f"\n{len(existing)} credential(s) currently on the applet:")
        for n in sorted(existing):
            print("  ", n)
        return 0

    # 6. Import.
    policy = ConflictPolicy(args)
    writer = uid = batch = ms = None
    memlog_fh = None
    prev_avail = None
    if args.profile_memory:
        uid = get_uid(connection)
        ms = MemorySession(channel)
        m = ms.select()
        prev_avail = m["available_persistent"] if m else None
        batch = ms.batch()
        _reauth(oath, pw)
        new_file = not os.path.exists(args.memlog)
        memlog_fh = open(args.memlog, "a", newline="", encoding="utf-8")
        writer = csv.writer(memlog_fh)
        if new_file:
            writer.writerow(["timestamp_iso", "card_uid", "mem_batch", "key_name",
                             "avail_before", "avail_after", "consumed_bytes"])

    ok = skipped = failed = 0
    try:
        for c in creds:
            if c.cred_id in existing and policy.decide(c.cred_id) == "skip":
                skipped += 1
                print(f"  skip   {c.cred_id} (already present)")
                continue
            try:
                if args.profile_memory:
                    oath.put(c)
                    cur = ms.available()          # SELECT memory (GC) -> deselects OATH
                    _reauth(oath, pw)             # re-select OATH for the next PUT
                    if prev_avail is not None and cur is not None:
                        before, after, consumed = prev_avail, cur, prev_avail - cur
                        prev_avail = cur
                    else:
                        before = after = consumed = ""
                    writer.writerow([datetime.datetime.now().isoformat(timespec="seconds"),
                                     uid, batch, c.cred_id, before, after, consumed])
                    print(f"  ok     {c.cred_id}  (consumed ~{consumed} B)")
                else:
                    oath.put(c)
                    print(f"  ok     {c.cred_id}")
                existing.add(c.cred_id)
                ok += 1
            except OathError as e:
                failed += 1
                print(f"  FAIL   {c.cred_id}: {e}")
                if e.sw == 0x6A84:
                    print("  Applet reports FULL (0x6A84) - stopping.")
                    break
    finally:
        if memlog_fh:
            memlog_fh.close()

    # Post-import: optionally set an access password on a not-yet-protected applet.
    if args.set_password and failed == 0:
        if applet_was_locked:
            print("Applet already password-protected; --set-password skipped.")
        else:
            oath.select()  # ensure OATH is the active applet (profiling deselects it)
            oath.set_password(args.set_password)
            print("Applet password set.")

    print(f"\nDone: {ok} written, {skipped} skipped, {failed} failed.")
    if args.profile_memory:
        print(f"Memory profile appended to {args.memlog}.")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except OathError as e:
        print(f"OATH error: {e}")
        sys.exit(1)
    except KeyboardInterrupt:
        sys.exit(130)