Build Your Own 2FA for the Windows Logon

#include "totp.h"
#include <bcrypt.h>

#pragma comment(lib, "bcrypt.lib")

namespace tac
{
    static const char kAlphabet[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";

    static int Base32Value(char c)
    {
        if (c >= 'A' && c <= 'Z') return c - 'A';
        if (c >= 'a' && c <= 'z') return c - 'a';
        if (c >= '2' && c <= '7') return c - '2' + 26;
        return -1;
    }

    bool Base32Decode(const std::string& in, std::vector<BYTE>& out)
    {
        uint32_t buffer = 0;
        int bits = 0;
        out.clear();

        for (char c : in)
        {
            if (c == '=' || c == ' ')
                continue;
            int v = Base32Value(c);
            if (v < 0)
                return false;

            buffer = (buffer << 5) | static_cast<uint32_t>(v);
            bits += 5;
            if (bits >= 8)
            {
                bits -= 8;
                out.push_back(static_cast<BYTE>(buffer >> bits));
                buffer &= (1u << bits) - 1;   // keep only the bits not consumed yet
            }
        }
        return true;
    }

    std::string Base32Encode(const BYTE* data, size_t len)
    {
        std::string out;
        uint32_t buffer = 0;
        int bits = 0;

        for (size_t i = 0; i < len; ++i)
        {
            buffer = (buffer << 8) | data[i];
            bits += 8;
            while (bits >= 5)
            {
                bits -= 5;
                out += kAlphabet[(buffer >> bits) & 0x1F];
            }
            buffer &= (1u << bits) - 1;
        }
        if (bits > 0)
            out += kAlphabet[(buffer << (5 - bits)) & 0x1F];
        return out;
    }

    static bool HmacSha1(const std::vector<BYTE>& key, const BYTE* msg, ULONG cbMsg, BYTE mac[20])
    {
        BCRYPT_ALG_HANDLE  hAlg  = nullptr;
        BCRYPT_HASH_HANDLE hHash = nullptr;
        bool ok = false;

        if (BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_SHA1_ALGORITHM, nullptr,
                                        BCRYPT_ALG_HANDLE_HMAC_FLAG) != 0)
            return false;

        if (BCryptCreateHash(hAlg, &hHash, nullptr, 0, const_cast<PUCHAR>(key.data()),
                             static_cast<ULONG>(key.size()), 0) == 0)
        {
            ok = BCryptHashData(hHash, const_cast<PUCHAR>(msg), cbMsg, 0) == 0 &&
                 BCryptFinishHash(hHash, mac, 20, 0) == 0;
            BCryptDestroyHash(hHash);
        }
        BCryptCloseAlgorithmProvider(hAlg, 0);
        return ok;
    }

    bool TotpAt(const std::vector<BYTE>& key, uint64_t unixTime,
                uint32_t step, int digits, uint32_t& code)
    {
        // The counter is the number of time steps since 1970, as 8 bytes big-endian.
        uint64_t counter = unixTime / step;
        BYTE msg[8];
        for (int i = 7; i >= 0; --i)
        {
            msg[i] = static_cast<BYTE>(counter & 0xFF);
            counter >>= 8;
        }

        BYTE mac[20];
        if (!HmacSha1(key, msg, sizeof(msg), mac))
            return false;

        // Dynamic truncation (RFC 4226, 5.3): the low nibble of the last byte
        // selects 4 bytes; the top bit is dropped to avoid signed values.
        int offset = mac[19] & 0x0F;
        uint32_t bin = ((mac[offset]     & 0x7Fu) << 24) |
                       ((mac[offset + 1] & 0xFFu) << 16) |
                       ((mac[offset + 2] & 0xFFu) <<  8) |
                        (mac[offset + 3] & 0xFFu);

        uint32_t mod = 1;
        for (int i = 0; i < digits; ++i)
            mod *= 10;
        code = bin % mod;
        return true;
    }

    bool ValidateTotp(const std::vector<BYTE>& key, const std::wstring& input,
                      uint64_t now, uint64_t lastStep, uint64_t& matchedStep,
                      uint32_t step, int digits, int window)
    {
        matchedStep = 0;

        // Exactly `digits` ASCII digits. wcstoul would also accept "12abc" or "0".
        if (input.size() != static_cast<size_t>(digits))
            return false;
        uint32_t entered = 0;
        for (wchar_t c : input)
        {
            if (c < L'0' || c > L'9')
                return false;
            entered = entered * 10 + static_cast<uint32_t>(c - L'0');
        }

        const uint64_t current = now / step;
        for (int w = -window; w <= window; ++w)
        {
            if (w < 0 && current < static_cast<uint64_t>(-w))
                continue;
            const uint64_t counter = current + static_cast<int64_t>(w);

            // "<=" and not "==": a step at or before the last used one is never
            // accepted again. This also stops an old code after the clock was
            // turned back.
            if (counter <= lastStep)
                continue;

            uint32_t expected = 0;
            if (TotpAt(key, counter * step, step, digits, expected) && expected == entered)
            {
                matchedStep = counter;
                return true;
            }
        }
        return false;
    }
}

import hmac, hashlib, struct, sys

BASE32="ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"

# --- faithful port of my C++ Base32Encode ---
def b32enc(data):
    out=""; buf=0; bits=0
    for b in data:
        buf=(buf<<8)|b; bits+=8
        while bits>=5:
            bits-=5; out+=BASE32[(buf>>bits)&0x1F]
        buf &= (1<<bits)-1
    if bits>0: out+=BASE32[(buf<<(5-bits))&0x1F]
    return out

# --- faithful port of my C++ Base32Decode ---
def b32dec(s):
    def val(c):
        if 'A'<=c<='Z': return ord(c)-65
        if 'a'<=c<='z': return ord(c)-97
        if '2'<=c<='7': return ord(c)-ord('2')+26
        return -1
    buf=0; bits=0; out=bytearray()
    for c in s:
        if c in '= ': continue
        v=val(c)
        if v<0: return None
        buf=(buf<<5)|v; bits+=5
        if bits>=8:
            bits-=8; out.append((buf>>bits)&0xFF); buf&=(1<<bits)-1
    return bytes(out)

# --- faithful port of my C++ TotpAt ---
def totp_at(key, unixtime, step=30, digits=6):
    counter = unixtime//step
    msg = struct.pack(">Q", counter)          # big-endian 8 bytes == my loop
    h = hmac.new(key, msg, hashlib.sha1).digest()
    off = h[19] & 0x0F
    binv = ((h[off]&0x7F)<<24)|((h[off+1]&0xFF)<<16)|((h[off+2]&0xFF)<<8)|(h[off+3]&0xFF)
    mod = 10**digits
    return binv % mod

seed = b"12345678901234567890"  # RFC 6238 appendix B (SHA1)
vectors = {59:94287082, 1111111109:7081804, 1111111111:14050471,
           1234567890:89005924, 2000000000:69279037, 20000000000:65353130}

print("== RFC 6238 (8-digit) ==")
ok_all=True
for t,exp in vectors.items():
    got8 = totp_at(seed, t, digits=8)
    got6 = totp_at(seed, t, digits=6)
    ok = (got8==exp)
    ok_all &= ok
    print(f"T={t:<12} got8={got8:08d} exp={exp:08d} {'OK' if ok else 'FAIL'}   (6-digit={got6:06d})")

print("\n== Base32 round-trip ==")
enc=b32enc(seed)
dec=b32dec(enc)
print("encode(seed)=",enc)
print("decode==seed:", dec==seed)

# Known: base32 of ASCII "12345678901234567890" is GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ
print("matches known GEZD... :", enc=="GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ")
ok_all &= dec==seed and enc=="GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"

# --- faithful port of my C++ ValidateTotp (with replay protection) ---
def validate(key, code, now, last_step, step=30, digits=6, window=1):
    if len(code)!=digits or not all('0'<=c<='9' for c in code): return None
    current = now//step
    for w in range(-window, window+1):
        if w<0 and current < -w: continue
        counter = current + w
        if counter <= last_step: continue          # "<=", not "=="
        if totp_at(key, counter*step, step, digits) == int(code): return counter
    return None

print("\n== Replay protection ==")
T=1800000000
c=f"{totp_at(seed,T):06d}"
s1=validate(seed,c,T,0)
checks = {
    "fresh code accepted":             s1==T//30,
    "same code rejected":              validate(seed,c,T+5,s1) is None,
    "previous step rejected":          validate(seed,f"{totp_at(seed,T-30):06d}",T,s1) is None,
    "next step accepted":              validate(seed,f"{totp_at(seed,T+30):06d}",T,s1)==T//30+1,
    "old code after clock turned back":validate(seed,f"{totp_at(seed,T-3600):06d}",T-3600,s1) is None,
    "'12abc' rejected":                validate(seed,"12abc",T,0) is None,
}
for k,v in checks.items(): print(f"{k:<34} {'OK' if v else 'FAIL'}")
ok_all &= all(checks.values())

print("\nALL CHECKS PASS:", ok_all)

# CI runs this script on its own; without a failing exit code it could never go red.
sys.exit(0 if ok_all else 1)

#pragma once
#include <windows.h>
#include <cstdint>
#include <string>
#include <vector>

namespace tac
{
    // "PC\Alice", ".\alice" and "alice" all become "alice". Case folding is
    // locale-invariant, so non-ASCII names fold the same way everywhere.
    std::wstring NormalizeUser(const std::wstring& user);

    // The same, for code inside lsass: `user` is counted (cch characters, it
    // may also end early at a NUL), the result goes into a caller buffer, and
    // nothing allocates. Returns false for an empty name or one that does not
    // fit. Lives in store_index.cpp.
    bool NormalizeUserInto(const WCHAR* user, SIZE_T cch, WCHAR* out, SIZE_T outCap);

    // Looks up COMPUTERNAME\name and returns the SID string ("S-1-5-21-...")
    // if it is a local user account. Everything else returns false.
    bool ResolveLocalUserSid(const std::wstring& user, std::wstring& sid);

    // The last sub-authority of a SID string, the RID ("...-1001" -> 1001).
    // Returns false for anything that is not a well-formed SID string.
    bool RidFromSidString(const std::wstring& sid, DWORD& rid);

    // Is this a syntactically valid SID string ("S-1-...")? Lets enroll accept
    // a SID where a name no longer resolves.
    bool IsSidString(const std::wstring& sid);

    // The secret is stored per SID, not per name. A renamed or re-created
    // account is a different account. StoreSecret also mirrors two enrollment
    // indexes for code that runs inside lsass: one keyed by RID (see
    // EnrollmentByRid) and one keyed by name (see EnrollmentByName).
    bool StoreSecret(const std::wstring& sid, const std::string& base32Secret);
    bool LoadSecretKey(const std::wstring& sid, std::vector<BYTE>& key);

    // Hot-path enrollment checks for code that runs INSIDE lsass. Each is a
    // single registry read: no DPAPI decrypt and no account-name or SID lookup,
    // so nothing calls back into LSA and the logon path cannot deadlock.
    // Neither allocates: a caller in lsass must not risk std::bad_alloc.
    //
    // Three outcomes, because a caller that VETOES must not read a failed read
    // as "not enrolled" - that is the one wrong answer that lets an enrolled
    // account through. A missing index key is No (nothing indexed on this
    // machine at all); any other read error is Unknown.
    enum class Enrollment { No, Yes, Unknown };

    // EnrollmentByRid is the one to use whenever the caller has the account's
    // RID from the SAM (the MSV1_0 sub-authentication filter gets it in
    // USER_ALL_INFORMATION). A RID survives a rename and the SAM never hands
    // out a RID twice, so this index can neither miss a renamed account nor
    // match a different, re-created one.
    Enrollment EnrollmentByRid(DWORD rid);

    // IsEnrolledByRid is EnrollmentByRid(rid) == Yes, so it reads every failure
    // as "not enrolled". Only for callers that do not veto.
    bool IsEnrolledByRid(DWORD rid);

    // EnrollmentByName is only for callers that have nothing but the name a
    // client typed (the authentication package). `name` is counted, not
    // null-terminated, and is normalized into a fixed buffer here.
    //
    // It is weaker than the RID check and fails OPEN in one case: after an
    // enrolled account is renamed, the index still holds the old name, so the
    // new name reads as "not enrolled" until `enroll /reindex` runs. A
    // re-created account with an old name reads as enrolled (more
    // restrictive). Prefer EnrollmentByRid wherever the RID is available.
    Enrollment EnrollmentByName(const WCHAR* name, SIZE_T cch);

    // Removes an enrollment completely: the DPAPI secret, the state and the RID
    // index entry under the SID, and the name-keyed index entry. Both halves are
    // removed independently on purpose. The name entry has to go even when the
    // Windows account was already deleted and the SID can no longer be resolved -
    // otherwise that name stays "enrolled" forever. (A RID entry of a deleted
    // account is harmless: the SAM never reuses the RID.)
    // An empty user or sid simply skips that half. A value that is not there
    // counts as removed.
    bool RemoveEnrollment(const std::wstring& user, const std::wstring& sid);

    // Is at least one account enrolled? Counts only stored secrets: REG_BINARY
    // values whose name is a SID string. A stray value of another type or name
    // does not count, because the credential provider's filter hides every
    // other sign-in option as soon as this says yes - and it must not do that
    // when nothing can actually authenticate.
    // Returns false if the store exists but cannot be read; a store that does
    // not exist yet is "readable, nothing enrolled".
    bool QueryAnyEnrollment(bool& any);

    // One enrolled account, for `enroll /list`.
    struct EnrollmentEntry
    {
        std::wstring sid;
        std::wstring name;      // empty when the SID no longer resolves
        DWORD        rid;
        bool         ridIndexed;
        bool         nameIndexed;
    };

    // Every stored secret with its index state. Returns false if the store
    // cannot be read; an empty store is "readable, nothing enrolled".
    bool ListEnrollments(std::vector<EnrollmentEntry>& out);

    // Removes everything this project keeps in the registry: secrets, state
    // and both indexes. For a full uninstall. A tree that is not there counts
    // as removed.
    bool PurgeAllEnrollments();

    // Re-applies the "SYSTEM and Administrators only" ACL to the store keys
    // that already exist. New keys get it when they are created; this repairs
    // keys left over from an older version with a weaker ACL. Called by
    // enroll.exe, never from the logon path - it needs WRITE_DAC and writes to
    // the hive, and neither belongs in a logon.
    bool EnsureStoreAcls();

    // Rebuilds the RID and name indexes from the stored secrets. For accounts
    // that were enrolled before the RID index existed, and after a rename.
    // Stale name entries of renamed accounts are left alone (they only make the
    // name check stricter). Returns the number of enrollments indexed, -1 if
    // the enrollments cannot be read or a RID entry cannot be written.
    int ReindexEnrollments();

    // Per-account state for replay protection and rate limiting.
    struct UserState
    {
        uint64_t lastStep;      // last accepted TOTP time step, 0 = none yet
        uint32_t failures;      // wrong codes in a row
        uint32_t reserved;
        uint64_t lockedUntil;   // unix time, 0 = not locked
    };
    static_assert(sizeof(UserState) == 24, "UserState is stored as-is in the registry");

    // A missing value is a fresh state (all zero). A broken value is an error.
    // A read-modify-write of the state must hold a StateLock (statelock.h).
    bool LoadState(const std::wstring& sid, UserState& state);
    bool SaveState(const std::wstring& sid, const UserState& state);
    bool DeleteState(const std::wstring& sid);
}

#include "store.h"
#include "storekeys.h"
#include "totp.h"
#include <wincrypt.h>
#include <sddl.h>

#pragma comment(lib, "crypt32.lib")
#pragma comment(lib, "advapi32.lib")

// The index reads that run inside lsass (EnrollmentByRid, EnrollmentByName,
// NormalizeUserInto) live in store_index.cpp, so the LSA DLLs do not have to
// link DPAPI and TOTP code they never use. Everything here runs in LogonUI or
// enroll.exe only.

namespace tac
{
    using keys::kSecrets;
    using keys::kState;
    using keys::kNames;
    using keys::kRids;
    using keys::kMaxNameCch;

    // SYSTEM and Administrators only. "P" blocks the inherited ACL of
    // HKLM\SOFTWARE, which would let every user read the key.
    static const wchar_t kKeySddl[] = L"D:P(A;OICI;KA;;;SY)(A;OICI;KA;;;BA)";

    std::wstring NormalizeUser(const std::wstring& user)
    {
        WCHAR buf[kMaxNameCch];
        if (!NormalizeUserInto(user.c_str(), user.size(), buf, ARRAYSIZE(buf)))
            return std::wstring();
        return std::wstring(buf);
    }

    static bool Protect(const std::string& plain, std::vector<BYTE>& blob)
    {
        DATA_BLOB in  = { static_cast<DWORD>(plain.size()),
                          reinterpret_cast<BYTE*>(const_cast<char*>(plain.data())) };
        DATA_BLOB out = {};

        // Machine scope, so LogonUI (SYSTEM) can decrypt what an admin encrypted.
        if (!CryptProtectData(&in, L"tac-2fa", nullptr, nullptr, nullptr,
                              CRYPTPROTECT_LOCAL_MACHINE | CRYPTPROTECT_UI_FORBIDDEN, &out))
            return false;

        blob.assign(out.pbData, out.pbData + out.cbData);
        LocalFree(out.pbData);
        return true;
    }

    static bool Unprotect(const std::vector<BYTE>& blob, std::string& plain)
    {
        DATA_BLOB in  = { static_cast<DWORD>(blob.size()), const_cast<BYTE*>(blob.data()) };
        DATA_BLOB out = {};

        if (!CryptUnprotectData(&in, nullptr, nullptr, nullptr, nullptr,
                                CRYPTPROTECT_UI_FORBIDDEN, &out))
            return false;

        plain.assign(reinterpret_cast<char*>(out.pbData), out.cbData);
        SecureZeroMemory(out.pbData, out.cbData);
        LocalFree(out.pbData);
        return true;
    }

    // Resolves the name the same way LSA will: as COMPUTERNAME\name against the
    // local SAM. Only a real user account counts, not a group or an alias.
    bool ResolveLocalUserSid(const std::wstring& user, std::wstring& sid)
    {
        sid.clear();
        std::wstring name = NormalizeUser(user);
        if (name.empty())
            return false;

        WCHAR computer[MAX_COMPUTERNAME_LENGTH + 1] = {};
        DWORD cchComputer = ARRAYSIZE(computer);
        if (!GetComputerNameW(computer, &cchComputer))
            return false;

        std::wstring qualified = std::wstring(computer) + L"\\" + name;
        BYTE sidBuf[SECURITY_MAX_SID_SIZE];
        DWORD cbSid = sizeof(sidBuf);
        WCHAR domain[256] = {};
        DWORD cchDomain = ARRAYSIZE(domain);
        SID_NAME_USE use;
        if (!LookupAccountNameW(nullptr, qualified.c_str(), sidBuf, &cbSid,
                                domain, &cchDomain, &use))
            return false;

        // The account has to live in this computer's SAM, not in a domain.
        if (use != SidTypeUser || _wcsicmp(domain, computer) != 0)
            return false;

        PWSTR pszSid = nullptr;
        if (!ConvertSidToStringSidW(sidBuf, &pszSid))
            return false;
        sid = pszSid;
        LocalFree(pszSid);
        return true;
    }

    // Only SIDs are used as value names. This keeps odd input out of the registry.
    bool IsSidString(const std::wstring& sid)
    {
        return sid.size() > 4 && sid.compare(0, 4, L"S-1-") == 0;
    }

    bool RidFromSidString(const std::wstring& sid, DWORD& rid)
    {
        rid = 0;
        if (!IsSidString(sid))
            return false;
        size_t dash = sid.find_last_of(L'-');
        if (dash == std::wstring::npos || dash + 1 >= sid.size() || sid.size() - dash - 1 > 10)
            return false;

        uint64_t value = 0;
        for (size_t i = dash + 1; i < sid.size(); ++i)
        {
            if (sid[i] < L'0' || sid[i] > L'9')
                return false;
            value = value * 10 + static_cast<uint64_t>(sid[i] - L'0');
        }
        if (value == 0 || value > 0xFFFFFFFFull)
            return false;
        rid = static_cast<DWORD>(value);
        return true;
    }

    // SID string -> normalized local account name. Used only at enroll time (a
    // console app), never on the LSA hot path.
    static bool NameFromSid(const std::wstring& sidStr, std::wstring& name)
    {
        name.clear();
        PSID psid = nullptr;
        if (!ConvertStringSidToSidW(sidStr.c_str(), &psid))
            return false;

        WCHAR n[256] = {}; DWORD cn = ARRAYSIZE(n);
        WCHAR d[256] = {}; DWORD cd = ARRAYSIZE(d);
        SID_NAME_USE use;
        BOOL ok = LookupAccountSidW(nullptr, psid, n, &cn, d, &cd, &use);
        LocalFree(psid);
        if (!ok)
            return false;

        name = NormalizeUser(n);
        return !name.empty();
    }

    // Creates or opens a key below HKLM with our ACL. The security attributes
    // only apply when the key is CREATED; an existing key keeps whatever ACL it
    // has. Repairing one needs WRITE_DAC and a hive write, neither of which
    // belongs on the logon path (SaveState runs there on every code attempt),
    // so that is EnsureStoreAcls' job and enroll.exe calls it.
    static LSTATUS OpenProtectedKey(const wchar_t* path, HKEY* phKey)
    {
        PSECURITY_DESCRIPTOR psd = nullptr;
        if (!ConvertStringSecurityDescriptorToSecurityDescriptorW(kKeySddl, SDDL_REVISION_1,
                                                                  &psd, nullptr))
            return static_cast<LSTATUS>(GetLastError());

        SECURITY_ATTRIBUTES sa = { sizeof(sa), psd, FALSE };
        LSTATUS status = RegCreateKeyExW(HKEY_LOCAL_MACHINE, path, 0, nullptr, 0,
                                         KEY_SET_VALUE | KEY_QUERY_VALUE,
                                         &sa, phKey, nullptr);
        LocalFree(psd);
        return status;
    }

    bool EnsureStoreAcls()
    {
        PSECURITY_DESCRIPTOR psd = nullptr;
        if (!ConvertStringSecurityDescriptorToSecurityDescriptorW(kKeySddl, SDDL_REVISION_1,
                                                                  &psd, nullptr))
            return false;

        bool ok = true;
        const wchar_t* paths[] = { kSecrets, kState, kNames, kRids };
        for (const wchar_t* path : paths)
        {
            HKEY hKey = nullptr;
            LSTATUS status = RegOpenKeyExW(HKEY_LOCAL_MACHINE, path, 0, WRITE_DAC, &hKey);
            if (status == ERROR_FILE_NOT_FOUND)
                continue;                   // not there yet, nothing to repair
            if (status != ERROR_SUCCESS)
            {
                ok = false;
                continue;
            }
            if (RegSetKeySecurity(hKey, DACL_SECURITY_INFORMATION, psd) != ERROR_SUCCESS)
                ok = false;
            RegCloseKey(hKey);
        }

        LocalFree(psd);
        return ok;
    }

    // Deletes one value and treats "was not there" as done.
    static bool DeleteValue(const wchar_t* path, const std::wstring& value)
    {
        HKEY hKey = nullptr;
        LSTATUS status = RegOpenKeyExW(HKEY_LOCAL_MACHINE, path, 0, KEY_SET_VALUE, &hKey);
        if (status == ERROR_FILE_NOT_FOUND)
            return true;
        if (status != ERROR_SUCCESS)
            return false;

        status = RegDeleteValueW(hKey, value.c_str());
        RegCloseKey(hKey);
        return status == ERROR_SUCCESS || status == ERROR_FILE_NOT_FOUND;
    }

    // Writes a presence flag (REG_DWORD 1) under one of the index keys.
    static bool SetFlag(const wchar_t* path, const std::wstring& value)
    {
        HKEY hKey = nullptr;
        if (OpenProtectedKey(path, &hKey) != ERROR_SUCCESS)
            return false;
        DWORD one = 1;
        bool ok = RegSetValueExW(hKey, value.c_str(), 0, REG_DWORD,
                                 reinterpret_cast<const BYTE*>(&one), sizeof(one)) == ERROR_SUCCESS;
        RegCloseKey(hKey);
        return ok;
    }

    // Writes the RID index and, if the SID still resolves to a name, the name
    // index for one enrolled SID.
    static void WriteIndexes(const std::wstring& sid, bool& ridOk, bool& nameOk)
    {
        ridOk = nameOk = false;

        DWORD rid = 0;
        if (RidFromSidString(sid, rid))
            ridOk = SetFlag(kRids, std::to_wstring(rid));

        std::wstring name;
        if (NameFromSid(sid, name))
            nameOk = SetFlag(kNames, name);
    }

    // Every REG_BINARY value under a SID name in the secret store. Shared by
    // QueryAnyEnrollment, ListEnrollments and ReindexEnrollments so all three
    // agree on what counts as an enrollment.
    static bool CollectEnrolledSids(std::vector<std::wstring>& sids)
    {
        sids.clear();

        HKEY hKey = nullptr;
        LSTATUS status = RegOpenKeyExW(HKEY_LOCAL_MACHINE, kSecrets, 0, KEY_QUERY_VALUE, &hKey);
        if (status == ERROR_FILE_NOT_FOUND)
            return true;                    // nothing enrolled yet
        if (status != ERROR_SUCCESS)
            return false;

        for (DWORD i = 0;; ++i)
        {
            WCHAR name[256];
            DWORD cch = ARRAYSIZE(name);
            DWORD type = 0;
            status = RegEnumValueW(hKey, i, name, &cch, nullptr, &type, nullptr, nullptr);
            if (status == ERROR_NO_MORE_ITEMS)
                break;
            if (status == ERROR_MORE_DATA)
                continue;                   // not one of ours, SIDs are short
            if (status != ERROR_SUCCESS)
            {
                RegCloseKey(hKey);
                return false;
            }
            if (type == REG_BINARY && IsSidString(std::wstring(name, cch)))
                sids.emplace_back(name, cch);
        }

        RegCloseKey(hKey);
        return true;
    }

    bool StoreSecret(const std::wstring& sid, const std::string& base32Secret)
    {
        if (!IsSidString(sid))
            return false;

        std::vector<BYTE> blob;
        if (!Protect(base32Secret, blob))
            return false;

        // The indexes the LSA hot path reads come first. If they cannot be
        // written, nothing else is touched, so a re-enrollment that fails
        // leaves the old enrollment working. If the secret then fails, the
        // index entries stay behind: that only makes the LSA packages refuse
        // MORE (network logon of an account without a secret), never less,
        // and `enroll /remove` clears them.
        bool ridOk = false, nameOk = false;
        WriteIndexes(sid, ridOk, nameOk);
        if (!ridOk || !nameOk)
            return false;

        HKEY hKey = nullptr;
        LSTATUS status = OpenProtectedKey(kSecrets, &hKey);
        if (status == ERROR_SUCCESS)
        {
            status = RegSetValueExW(hKey, sid.c_str(), 0, REG_BINARY,
                                    blob.data(), static_cast<DWORD>(blob.size()));
            RegCloseKey(hKey);
        }
        return status == ERROR_SUCCESS;
    }

    bool LoadSecretKey(const std::wstring& sid, std::vector<BYTE>& key)
    {
        key.clear();
        if (!IsSidString(sid))
            return false;

        HKEY hKey = nullptr;
        if (RegOpenKeyExW(HKEY_LOCAL_MACHINE, kSecrets, 0, KEY_QUERY_VALUE, &hKey) != ERROR_SUCCESS)
            return false;

        DWORD type = 0;
        DWORD cb = 0;
        std::vector<BYTE> blob;
        LSTATUS status = RegQueryValueExW(hKey, sid.c_str(), nullptr, &type, nullptr, &cb);
        if (status == ERROR_SUCCESS && type == REG_BINARY && cb > 0)
        {
            blob.resize(cb);
            status = RegQueryValueExW(hKey, sid.c_str(), nullptr, &type, blob.data(), &cb);
        }
        RegCloseKey(hKey);
        if (status != ERROR_SUCCESS || blob.empty())
            return false;

        std::string base32;
        if (!Unprotect(blob, base32))
            return false;

        bool ok = Base32Decode(base32, key);
        SecureZeroMemory(&base32[0], base32.size());

        // An empty or very short key would make the code predictable. Our
        // secrets have 20 bytes; anything under 16 counts as broken.
        if (!ok || key.size() < 16)
        {
            if (!key.empty())
                SecureZeroMemory(key.data(), key.size());
            key.clear();
            return false;
        }
        return true;
    }

    bool RemoveEnrollment(const std::wstring& user, const std::wstring& sid)
    {
        bool ok = true;

        // The name-keyed index first: it is the half that can still be found
        // when the account no longer exists. The RID entry of a deleted
        // account is harmless, the SAM never hands that RID out again.
        std::wstring typed = NormalizeUser(user);
        if (!typed.empty())
            ok = DeleteValue(kNames, typed) && ok;

        if (IsSidString(sid))
        {
            // StoreSecret keys the index by the name it got from NameFromSid,
            // not by whatever was typed. Those are normally the same, but
            // LookupAccountName accepts spellings the SAM does not store, so
            // remove the canonical one as well. Deleting a value twice is free.
            std::wstring canonical;
            if (NameFromSid(sid, canonical) && !canonical.empty() && canonical != typed)
                ok = DeleteValue(kNames, canonical) && ok;

            DWORD rid = 0;
            if (RidFromSidString(sid, rid))
                ok = DeleteValue(kRids, std::to_wstring(rid)) && ok;

            ok = DeleteValue(kSecrets, sid) && ok;
            ok = DeleteValue(kState, sid) && ok;
        }
        return ok;
    }

    bool QueryAnyEnrollment(bool& any)
    {
        any = false;

        // Counts stored secrets, not values: a stray value of another type or
        // name is not an enrollment. The credential provider's filter hides
        // every other sign-in option as soon as this says yes, and it must not
        // do that while nothing on the machine can actually authenticate.
        std::vector<std::wstring> sids;
        if (!CollectEnrolledSids(sids))
            return false;

        any = !sids.empty();
        return true;
    }

    bool ListEnrollments(std::vector<EnrollmentEntry>& out)
    {
        out.clear();

        std::vector<std::wstring> sids;
        if (!CollectEnrolledSids(sids))
            return false;

        for (const std::wstring& sid : sids)
        {
            EnrollmentEntry e;
            e.sid         = sid;
            e.rid         = 0;
            e.ridIndexed  = false;
            e.nameIndexed = false;

            if (RidFromSidString(sid, e.rid))
                e.ridIndexed = EnrollmentByRid(e.rid) == Enrollment::Yes;

            std::wstring name;
            if (NameFromSid(sid, name))
            {
                e.name = name;
                e.nameIndexed = EnrollmentByName(name.c_str(), name.size()) == Enrollment::Yes;
            }

            out.push_back(e);
        }
        return true;
    }

    bool PurgeAllEnrollments()
    {
        // HKLM\SOFTWARE\TheAdminCafe\2FA and everything below it, then the
        // vendor key if nothing else lives there.
        LSTATUS status = RegDeleteTreeW(HKEY_LOCAL_MACHINE, kSecrets);
        if (status != ERROR_SUCCESS && status != ERROR_FILE_NOT_FOUND)
            return false;
        status = RegDeleteKeyW(HKEY_LOCAL_MACHINE, L"SOFTWARE\\TheAdminCafe");
        // Not empty (someone else's data) or already gone: both fine.
        (void)status;
        return true;
    }

    int ReindexEnrollments()
    {
        // Collect first, write afterwards: the index keys are subkeys of the
        // secret store, and enumerating while writing elsewhere is simpler to
        // reason about.
        std::vector<std::wstring> sids;
        if (!CollectEnrolledSids(sids))
            return -1;

        // The RID index is what the sub-authentication filter reads, so a
        // failure there is an error. A missing name is not: the account may
        // have been deleted, and then there is no name to index.
        int count = 0;
        for (const std::wstring& sid : sids)
        {
            bool ridOk = false, nameOk = false;
            WriteIndexes(sid, ridOk, nameOk);
            if (!ridOk)
                return -1;
            ++count;
        }
        return count;
    }

    bool LoadState(const std::wstring& sid, UserState& state)
    {
        ZeroMemory(&state, sizeof(state));
        if (!IsSidString(sid))
            return false;

        HKEY hKey = nullptr;
        LSTATUS status = RegOpenKeyExW(HKEY_LOCAL_MACHINE, kState, 0, KEY_QUERY_VALUE, &hKey);
        if (status == ERROR_FILE_NOT_FOUND)
            return true;                      // nobody has logged on yet
        if (status != ERROR_SUCCESS)
            return false;

        DWORD type = 0;
        DWORD cb = sizeof(state);
        status = RegQueryValueExW(hKey, sid.c_str(), nullptr, &type,
                                  reinterpret_cast<BYTE*>(&state), &cb);
        RegCloseKey(hKey);

        if (status == ERROR_FILE_NOT_FOUND)
        {
            ZeroMemory(&state, sizeof(state));
            return true;                      // this user has not logged on yet
        }
        // Anything else that is not exactly our structure is an error. The
        // caller fails closed, so a broken value cannot reset the counters.
        return status == ERROR_SUCCESS && type == REG_BINARY && cb == sizeof(state);
    }

    bool SaveState(const std::wstring& sid, const UserState& state)
    {
        if (!IsSidString(sid))
            return false;

        HKEY hKey = nullptr;
        LSTATUS status = OpenProtectedKey(kState, &hKey);
        if (status == ERROR_SUCCESS)
        {
            status = RegSetValueExW(hKey, sid.c_str(), 0, REG_BINARY,
                                    reinterpret_cast<const BYTE*>(&state), sizeof(state));
            // No RegFlushKey: it can write out large parts of the hive and
            // blocks the logon until the disk is done, on every attempt. The
            // cost of leaving it out is narrow: a hard reset within seconds of
            // a logon can lose the last state write.
            RegCloseKey(hKey);
        }
        return status == ERROR_SUCCESS;
    }

    bool DeleteState(const std::wstring& sid)
    {
        if (!IsSidString(sid))
            return false;
        return DeleteValue(kState, sid);
    }
}

// The enrollment-index reads that run INSIDE lsass: the sub-authentication
// filter and the authentication package link this file and nothing else from
// the store. No DPAPI, no TOTP, no name or SID lookup, no heap and no CRT of
// our own - see lsasafe.h for why each of those matters in lsass.
//
// Everything that writes the indexes, or needs DPAPI, lives in store.cpp and
// only ever runs in LogonUI or enroll.exe.

#include "store.h"
#include "storekeys.h"

#pragma comment(lib, "advapi32.lib")

namespace tac
{
    bool NormalizeUserInto(const WCHAR* user, SIZE_T cch, WCHAR* out, SIZE_T outCap)
    {
        if (!user || !out || outCap == 0)
            return false;
        out[0] = L'\0';

        // Stop at an embedded NUL: nothing after it belongs to the name.
        SIZE_T len = 0;
        while (len < cch && user[len] != L'\0')
            ++len;

        // Keep only what follows the last backslash.
        SIZE_T start = 0;
        for (SIZE_T i = 0; i < len; ++i)
            if (user[i] == L'\\')
                start = i + 1;

        const SIZE_T bare = len - start;
        if (bare == 0 || bare >= outCap)
            return false;

        // LOCALE_INVARIANT, so "JÖRG" and "jörg" fold the same way on every
        // machine. Plain towlower only ever folded ASCII, which made the name
        // index miss non-ASCII names.
        const int written = LCMapStringW(LOCALE_INVARIANT, LCMAP_LOWERCASE,
                                         user + start, static_cast<int>(bare),
                                         out, static_cast<int>(outCap - 1));
        if (written > 0 && static_cast<SIZE_T>(written) < outCap)
        {
            out[written] = L'\0';
            return true;
        }

        // NLS unavailable or the mapping did not fit: fold ASCII only, which
        // is what every name used to get.
        for (SIZE_T i = 0; i < bare; ++i)
        {
            const WCHAR c = user[start + i];
            out[i] = (c >= L'A' && c <= L'Z') ? static_cast<WCHAR>(c - L'A' + L'a') : c;
        }
        out[bare] = L'\0';
        return true;
    }

    // The RID as the decimal value name of the Rids index, without the CRT.
    // This runs inside lsass, where swprintf_s on an overflow would end the
    // process. A DWORD has at most 10 digits, so out[12] always fits.
    static void RidToValueName(DWORD rid, WCHAR (&out)[12])
    {
        WCHAR digits[10];
        int n = 0;
        do
        {
            digits[n++] = static_cast<WCHAR>(L'0' + rid % 10);
            rid /= 10;
        } while (rid != 0 && n < 10);

        int i = 0;
        while (n > 0)
            out[i++] = digits[--n];
        out[i] = L'\0';                   // i <= 10, so this is in bounds
    }

    // One registry read of an index key, with the three outcomes a vetoing
    // caller needs.
    static Enrollment QueryIndex(const wchar_t* path, const WCHAR* value)
    {
        HKEY hKey = nullptr;
        LSTATUS status = RegOpenKeyExW(HKEY_LOCAL_MACHINE, path, 0, KEY_QUERY_VALUE, &hKey);
        if (status == ERROR_FILE_NOT_FOUND)
            return Enrollment::No;          // no index on this machine at all
        if (status != ERROR_SUCCESS)
            return Enrollment::Unknown;

        status = RegQueryValueExW(hKey, value, nullptr, nullptr, nullptr, nullptr);
        RegCloseKey(hKey);

        if (status == ERROR_SUCCESS)
            return Enrollment::Yes;
        if (status == ERROR_FILE_NOT_FOUND)
            return Enrollment::No;
        // Only a value that is really not there counts as "not enrolled". Any
        // other failure is Unknown: a filter that read it as "not enrolled"
        // would wave an enrolled account through on one unexpected error.
        return Enrollment::Unknown;
    }

    Enrollment EnrollmentByRid(DWORD rid)
    {
        if (rid == 0)
            return Enrollment::No;

        WCHAR value[12];
        RidToValueName(rid, value);
        return QueryIndex(keys::kRids, value);
    }

    bool IsEnrolledByRid(DWORD rid)
    {
        return EnrollmentByRid(rid) == Enrollment::Yes;
    }

    Enrollment EnrollmentByName(const WCHAR* name, SIZE_T cch)
    {
        WCHAR normalized[keys::kMaxNameCch];
        if (!NormalizeUserInto(name, cch, normalized, ARRAYSIZE(normalized)))
            return Enrollment::Unknown;     // unusable name: never claim "No"

        return QueryIndex(keys::kNames, normalized);
    }
}

#include "totp.h"
#include "store.h"
#include "statelock.h"
#include <bcrypt.h>
#include <sddl.h>
#include <conio.h>
#include <cctype>
#include <cstdio>
#include <ctime>
#include <string>
#include <vector>

#pragma comment(lib, "bcrypt.lib")
#pragma comment(lib, "advapi32.lib")

// The otpauth label is UTF-8 and percent-encoded, so names like "jürg" survive.
static std::string UrlEncode(const std::wstring& text)
{
    int cb = WideCharToMultiByte(CP_UTF8, 0, text.c_str(), -1, nullptr, 0, nullptr, nullptr);
    std::string utf8(cb > 1 ? cb - 1 : 0, '\0');
    if (cb > 1)
        WideCharToMultiByte(CP_UTF8, 0, text.c_str(), -1, &utf8[0], cb, nullptr, nullptr);

    static const char hex[] = "0123456789ABCDEF";
    std::string out;
    for (unsigned char c : utf8)
    {
        if (isalnum(c) || c == '-' || c == '_' || c == '.' || c == '~')
            out += static_cast<char>(c);
        else
        {
            out += '%';
            out += hex[c >> 4];
            out += hex[c & 0x0F];
        }
    }
    return out;
}

// Resets the lock of an enrolled user. The last used time step stays, so an
// unlock never makes an old code valid again.
static int Unlock(const std::wstring& user, const std::wstring& sid)
{
    tac::StateLock lock;
    if (!lock.Held())
    {
        wprintf(L"Could not take the state lock. Run this elevated, from the install folder.\n");
        return 1;
    }

    tac::UserState state;
    if (!tac::LoadState(sid, state))
    {
        wprintf(L"Could not read the state of '%s'. Run this from an elevated prompt.\n", user.c_str());
        return 1;
    }

    uint64_t now = static_cast<uint64_t>(time(nullptr));
    wprintf(L"'%s': %u wrong code(s) in a row, %s.\n", user.c_str(), state.failures,
            state.lockedUntil > now ? L"locked" : L"not locked");

    state.failures    = 0;
    state.lockedUntil = 0;
    if (!tac::SaveState(sid, state))
    {
        wprintf(L"Could not reset the state. Run this from an elevated prompt.\n");
        return 1;
    }
    wprintf(L"Unlocked.\n");
    return 0;
}

// Removes an enrollment: the secret, the state and the RID and name indexes the
// LSA packages read. The account does not have to exist anymore. If it was already
// deleted in Windows we cannot resolve a SID, but we can still drop the name
// entry, and that is the one that would otherwise keep denying network logons.
static int Remove(const std::wstring& user, const std::wstring& sid)
{
    const std::wstring& shown = user.empty() ? sid : user;

    if (sid.empty())
        wprintf(L"'%s' is not a local user account (anymore). Removing the name entry only.\n"
                L"For the secret of a deleted account, use the SID from 'enroll /list'.\n",
                user.c_str());

    if (!tac::RemoveEnrollment(user, sid))
    {
        wprintf(L"Could not remove the enrollment of '%s'. Run this from an elevated prompt.\n",
                shown.c_str());
        return 1;
    }

    wprintf(L"Removed the enrollment of '%s'.\n", shown.c_str());
    wprintf(L"The account can log on with its password alone again, unless you also\n");
    wprintf(L"set the deny rights (deny-noninteractive.ps1).\n");
    return 0;
}

// Lists every stored secret, the account it belongs to and whether the two
// indexes the LSA packages read know about it. A secret whose SID no longer
// resolves belongs to a deleted account: it can never authenticate, but it
// still counts as "enrolled" for the filter, so remove it by SID.
static int List()
{
    std::vector<tac::EnrollmentEntry> entries;
    if (!tac::ListEnrollments(entries))
    {
        wprintf(L"Could not read the enrollments. Run this from an elevated prompt.\n");
        return 1;
    }
    if (entries.empty())
    {
        wprintf(L"No account is enrolled. The filter hides nothing.\n");
        return 0;
    }

    bool orphans = false;
    for (const tac::EnrollmentEntry& e : entries)
    {
        const bool gone = e.name.empty();
        orphans = orphans || gone;
        wprintf(L"  %-20s %s  RID index: %s  name index: %s\n",
                gone ? L"(deleted account)" : e.name.c_str(), e.sid.c_str(),
                e.ridIndexed ? L"yes" : L"NO",
                gone ? L"-" : (e.nameIndexed ? L"yes" : L"NO"));
    }

    if (orphans)
    {
        wprintf(L"\nA deleted account still counts as enrolled, so the filter keeps hiding the\n");
        wprintf(L"other sign-in options. Remove it with: enroll /remove <SID>\n");
    }
    return 0;
}

// Rebuilds the RID and name indexes the LSA packages read, from the stored
// secrets. Needed once after updating from a version without the RID index,
// and after renaming an enrolled account.
static int Reindex()
{
    int count = tac::ReindexEnrollments();
    if (count < 0)
    {
        wprintf(L"Could not read the enrollments or write the index. Run this from an elevated prompt.\n");
        return 1;
    }
    wprintf(L"Indexed %d enrollment(s).\n", count);
    return 0;
}

// Reads one line from the console, at most `max` characters.
static std::wstring ReadLine(size_t max)
{
    wchar_t buf[64] = {};
    if (!fgetws(buf, ARRAYSIZE(buf), stdin))
        return std::wstring();
    std::wstring line(buf);
    SecureZeroMemory(buf, sizeof(buf));
    while (!line.empty() && (line.back() == L'\n' || line.back() == L'\r' || line.back() == L' '))
        line.pop_back();
    if (line.size() > max)
        line.clear();
    return line;
}

// Asks for a code from the authenticator app before anything is stored. A
// secret that was mistyped into the app would otherwise lock the account out
// at the next logon. On success `step` is the time step of the code, so the
// same code cannot be used again at the logon screen.
static bool ConfirmSecret(const std::string& secret, uint64_t& step)
{
    std::vector<BYTE> key;
    if (!tac::Base32Decode(secret, key))
        return false;

    bool ok = false;
    for (int attempt = 1; attempt <= 3 && !ok; ++attempt)
    {
        wprintf(L"Enter the 6-digit code from your authenticator app (attempt %d of 3): ", attempt);
        std::wstring code = ReadLine(16);
        ok = tac::ValidateTotp(key, code, static_cast<uint64_t>(time(nullptr)), 0, step);
        if (!code.empty())
            SecureZeroMemory(&code[0], code.size() * sizeof(WCHAR));
        if (!ok)
            wprintf(L"That code does not match. Check the secret in the app and the clock of both devices.\n");
    }
    SecureZeroMemory(key.data(), key.size());
    return ok;
}

// Deletes every secret, state and index entry. For a full uninstall; the
// provider registration itself is removed by unregister.reg.
static int Purge()
{
    bool any = false;
    if (tac::QueryAnyEnrollment(any) && any)
    {
        wprintf(L"This deletes ALL 2FA secrets and state on this machine. If the filter is\n");
        wprintf(L"still registered, the other sign-in options come back immediately.\n");
        wprintf(L"Continue? [y/N] ");
        int answer = _getwch();
        wprintf(L"\n");
        if (answer != L'y' && answer != L'Y')
        {
            wprintf(L"Aborted.\n");
            return 0;
        }
    }

    if (!tac::PurgeAllEnrollments())
    {
        wprintf(L"Could not delete HKLM\\SOFTWARE\\TheAdminCafe\\2FA. Run this from an elevated prompt.\n");
        return 1;
    }
    wprintf(L"All 2FA data removed.\n");
    return 0;
}

// Parses a SID typed on the command line and returns it in the canonical form
// the store uses as value name ("S-1-5-21-..."). Stricter than the "S-1-"
// prefix check, so a local user who happens to be called "S-1-x" is not
// mistaken for one, and "s-1-5-..." in lower case still finds its entries.
static bool CanonicalSid(const wchar_t* text, std::wstring& sid)
{
    sid.clear();
    PSID psid = nullptr;
    if (!ConvertStringSidToSidW(text, &psid))
        return false;

    PWSTR canonical = nullptr;
    const BOOL ok = ConvertSidToStringSidW(psid, &canonical);
    LocalFree(psid);
    if (!ok)
        return false;

    sid = canonical;
    LocalFree(canonical);
    return true;
}

int wmain(int argc, wchar_t** argv)
{
    if (argc == 2 && _wcsicmp(argv[1], L"/list") == 0)
        return List();
    if (argc == 2 && _wcsicmp(argv[1], L"/purge") == 0)
        return Purge();

    const bool reindex = argc == 2 && _wcsicmp(argv[1], L"/reindex") == 0;
    const bool unlock  = argc == 3 && _wcsicmp(argv[1], L"/unlock") == 0;
    const bool remove  = argc == 3 && _wcsicmp(argv[1], L"/remove") == 0;
    if (!reindex &&
        ((argc != 2 && !unlock && !remove) || (argc == 2 && argv[1][0] == L'/')))
    {
        wprintf(L"Usage: enroll <local username>\n");
        wprintf(L"       enroll /unlock <local username>\n");
        wprintf(L"       enroll /remove <local username | SID>\n");
        wprintf(L"       enroll /list           (enrolled accounts and index state)\n");
        wprintf(L"       enroll /reindex\n");
        wprintf(L"       enroll /purge          (delete all 2FA data)\n");
        return 1;
    }

    // Repair the ACL of store keys an older version created with a weaker one.
    // This used to happen on every state write, i.e. on every logon attempt;
    // it needs WRITE_DAC and a hive write and belongs here instead.
    if (!tac::EnsureStoreAcls())
        wprintf(L"Warning: could not re-apply the ACL to the existing 2FA registry keys.\n");

    if (reindex)
        return Reindex();

    const wchar_t* arg = argv[(unlock || remove) ? 2 : 1];

    // Only names that exist as a local user account on this machine. The
    // secret is stored under the SID of that account.
    std::wstring user = tac::NormalizeUser(arg);
    std::wstring sid;
    if (!tac::ResolveLocalUserSid(user, sid))
    {
        // For /remove an unresolvable name is not fatal: the name-keyed index
        // entry can and should still go. A SID (from `enroll /list`) removes
        // the secret of an account that was already deleted. Everything else
        // needs a live account.
        if (!remove)
        {
            wprintf(L"'%s' is not a local user account on this machine.\n", user.c_str());
            return 1;
        }
        if (CanonicalSid(arg, sid))
            user.clear();
        else
            sid.clear();

        if (user.empty() && sid.empty())
        {
            wprintf(L"'%s' is neither a user name nor a SID.\n", arg);
            return 1;
        }
    }

    if (unlock)
        return Unlock(user, sid);
    if (remove)
        return Remove(user, sid);

    std::vector<BYTE> existing;
    if (tac::LoadSecretKey(sid, existing))
    {
        SecureZeroMemory(existing.data(), existing.size());
        wprintf(L"'%s' is already enrolled. Replace the secret? [y/N] ", user.c_str());
        int answer = _getwch();
        wprintf(L"\n");
        if (answer != L'y' && answer != L'Y')
        {
            wprintf(L"Aborted.\n");
            return 0;
        }
    }

    // 160 bits, the key size RFC 4226 recommends for HMAC-SHA1.
    BYTE raw[20];
    if (BCryptGenRandom(nullptr, raw, sizeof(raw), BCRYPT_USE_SYSTEM_PREFERRED_RNG) != 0)
    {
        wprintf(L"Could not generate a random secret.\n");
        return 1;
    }
    std::string secret = tac::Base32Encode(raw, sizeof(raw));
    SecureZeroMemory(raw, sizeof(raw));

    std::string uri = "otpauth://totp/TheAdminCafe:" + UrlEncode(user) +
                      "?secret=" + secret +
                      "&issuer=TheAdminCafe&algorithm=SHA1&digits=6&period=30";

    // Show the secret first and store it only after a code from the app
    // matches. Until then nothing is written, and an existing enrollment
    // keeps working.
    wprintf(L"\nNew secret for '%s' (%s):\n\n", user.c_str(), sid.c_str());
    wprintf(L"  Secret : %S\n", secret.c_str());
    wprintf(L"  URI    : %S\n\n", uri.c_str());
    wprintf(L"Add the secret to your authenticator app. Don't paste it into an online\n");
    wprintf(L"QR code generator.\n\n");
    SecureZeroMemory(&uri[0], uri.size());

    uint64_t confirmedStep = 0;
    if (!ConfirmSecret(secret, confirmedStep))
    {
        wprintf(L"\nNot enrolled: no matching code. Nothing was changed.\n");
        SecureZeroMemory(&secret[0], secret.size());
        return 1;
    }

    // A new secret starts with a clean state: no lock, and the confirmation
    // code counts as used. The two steps are reported separately, because "the
    // secret is stored but the old state is still there" is a different
    // situation from "nothing was written" and needs a different fix.
    if (!tac::StoreSecret(sid, secret))
    {
        wprintf(L"Could not store the secret. Run this from an elevated prompt.\n");
        wprintf(L"An earlier enrollment of '%s', if any, is unchanged.\n", user.c_str());
        SecureZeroMemory(&secret[0], secret.size());
        return 1;
    }
    SecureZeroMemory(&secret[0], secret.size());

    tac::UserState fresh = {};
    fresh.lastStep = confirmedStep;
    tac::StateLock lock;
    if (!lock.Held() || !tac::SaveState(sid, fresh))
    {
        wprintf(L"The secret was stored, but the old lockout state of '%s' could not be\n", user.c_str());
        wprintf(L"reset. Run 'enroll /unlock %s' from an elevated prompt.\n", user.c_str());
        return 1;
    }

    wprintf(L"\nEnrolled '%s'. Test the logon now, before you import the filter.\n", user.c_str());
    wprintf(L"Wait for the NEXT code in the app: the one you just typed is used up, and\n");
    wprintf(L"typing it again at the logon screen counts as a wrong code.\n");
    wprintf(L"Clear this window afterwards (cls), the secret is still in the scrollback.\n");
    return 0;
}