1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2025 Sam Blenny
#
# See NOTES.md for documentation links and pinout info.
#
import time

from controller import (Controller, NEOKEY0, NEOKEY1, NEOKEY2, NEOKEY3,
    NP_AMBER, NP_RED)


# The controller manages IO with I2C and SPI devices, and it keeps track of the
# EEPROM-backed database of TOTP accounts
ctrl = Controller()
ctrl.load_totp_accounts()


# These variables help with edge detection of changing values in the main loop
t = ctrl.datetime()
prev_t = t
prev_nk = ctrl.get_neokey_bits()
need_refresh = True

# Main loop
while True:

    # Is it time to update the TOTP code? (period over or refresh requested)
    if t.tm_sec % 30 == 0 or need_refresh:
        # Generate new totp code at multiples of 30 seconds
        unix_time = time.mktime(ctrl.datetime())
        slot, label, totp_code = ctrl.get_selected_totp(unix_time)
        code_changed = True

    # Redraw the display only when backlight is on
    if ctrl.bl_enable:
        seconds = t[5] % 30
        ctrl.set_text('%02d\n%s\n%s' % (seconds, label, totp_code))

    # Make sure the right NeoKey selected slot LED is lit
    for i in range(4):
        if i == ctrl.get_selected_slot():
            if ctrl.is_selected_slot_empty():
                ctrl.set_neokey(i, NP_RED)    # red for empty slots
            else:
                ctrl.set_neokey(i, NP_AMBER)  # amber for configured slots
        else:
            ctrl.set_neokey_off(i)

    # Clear the manual update requested flag
    need_refresh = False

    # Wait until second rolls over
    while t.tm_sec == prev_t.tm_sec and not need_refresh:

        # 1. Spend about 100ms polling for input events
        t_10Hz = time.monotonic() + 0.1
        while time.monotonic() < t_10Hz and not need_refresh:

            # Sleep to rate limit I2C, debounce buttons, do VM background tasks
            time.sleep(0.01)

            # Check the currently selected TOTP account slot
            slot = ctrl.get_selected_slot()

            # Monitor the 4 NeoKeys for state transitions
            nk = ctrl.get_neokey_bits()  # get bitfield of current key states
            diff = nk ^ prev_nk          # use XOR to find which bits changed
            prev_nk = nk                 # update previous values for next loop

            # Did any key get pressed while backlight was off?
            waking_up = False
            if (diff & nk) and not ctrl.bl_enable:
                # Rising edge of key press -> wake up the screen
                ctrl.backlight_on()
                need_refresh = True
                waking_up = True  # don't go back to sleep when key released

            # Check for key press to select TOTP account slots. When the
            # currently selected key is pressed, turn off the display. When a
            # different key is pressed, select the slot for that key.
            for i, bitmask in enumerate((NEOKEY0, NEOKEY1, NEOKEY2, NEOKEY3)):
                if diff & bitmask and ctrl.bl_enable:
                    # Check if other keys were already down to avoid potential
                    # weird glitchy behavior if you fat-finger multiple keys at
                    # the same time
                    other_keys_pressed = bool(nk & (~bitmask))
                    if nk & bitmask and not other_keys_pressed:
                        # Key was just pressed
                        if slot == i:
                            # Slot was already selected (toggle backlight?)
                            if not waking_up:
                                ctrl.backlight_off()
                                ctrl.set_neokey_off(i)
                        else:
                            # Slot wasn't selected -> select it
                            if ctrl.select_account(i):
                                # Selected a non-empty slot
                                ctrl.set_neokey(i, NP_AMBER)
                                need_refresh = True
                            else:
                                # Empty slot
                                ctrl.set_neokey(i, NP_RED)
                                need_refresh = True
                    else:
                        # Key released or other keys were already pressed
                        pass

            # -- END of input polling loop --


        # 2. Poll RTC at about 10 Hz to detect when seconds have changed
        if ctrl.bl_enable:
            t = ctrl.datetime()

    # After the seconds have changed, update the previous time
    prev_t = t

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2025 Sam Blenny
#
# See NOTES.md for documentation links and pinout info.
#
import atexit
import board
import busio
import collections
import digitalio
import displayio
from fourwire import FourWire
import gc
from micropython import const
from pwmio import PWMOut
import os
import terminalio


from adafruit_24lc32 import EEPROM_I2C
from adafruit_display_text import label
from adafruit_ds3231 import DS3231
from adafruit_neokey.neokey1x4 import NeoKey1x4
from adafruit_st7789 import ST7789

from eeprom_db import check_eeprom_format, is_slot_in_use, load_totp_account
from sb_totp import base32_encode, totp_sha1


# Begin TFT backlight dimming (100% brightness is PWM duty_cycle=0xffff)
BACKLIGHT_ON = const(26214)  # 40% of 0xffff
BACKLIGHT_OFF = const(0)

# NeoKey keypress bitfield constants
NEOKEY0 = const(1)
NEOKEY1 = const(2)
NEOKEY2 = const(4)
NEOKEY3 = const(8)

# NeoKey NeoPixel color constants
NP_AMBER = const(0x202000)
NP_RED   = const(0x400000)
NP_OFF   = const(0)


class Controller:

    def __init__(self):
        # I2C RTC, EEPROM, and NeoKey
        i2c = busio.I2C(board.SCL, board.SDA, frequency=250_000)
        self.i2c = i2c
        self.rtc = DS3231(i2c)
        self.neokey = NeoKey1x4(i2c, addr=0x30)
        self.eeprom = EEPROM_I2C(i2c)
        self.accounts = []                            # TOTP account slot data
        self.selected_acct = None

        # Display backlight PWM dimming
        self.backlight = PWMOut(board.TFT_BACKLIGHT, frequency=500,
            duty_cycle=BACKLIGHT_ON)
        self.bl_enable = True

        # CLUE 240x240 ST7789 TFT Display
        displayio.release_displays()
        gc.collect()
        spi = busio.SPI(board.TFT_SCK, MOSI=board.TFT_MOSI)
        display_bus = FourWire(spi, command=board.TFT_DC,
            chip_select=board.TFT_CS, reset=board.TFT_RESET)
        display = ST7789(display_bus, width=240, height=240, rowstart=80,
            bgr=True, rotation=0, auto_refresh=False)
        group = displayio.Group()
        display.root_group = group
        display.refresh()
        textbox = label.Label(font=terminalio.FONT, scale=4, color=0xefef00)
        textbox.anchor_point = (0, 0)
        textbox.anchored_position = (16, 8)
        textbox.line_spacing = 1.25  # default is 1.25
        group.append(textbox)
        self.spi = spi
        self.display_bus = display_bus
        self.display = display
        self.group = group
        self.textbox = textbox

        # Set an atexit handler to release the display once code.py ends. This
        # is an aesthetic filter to prevent CircuitPython's supervisor from
        # hijacking the display to show noisy stuff I don't want to see.
        def atexit_shutdown_display():
            try:
                self.set_text('OFFLINE')
                displayio.release_displays()
                self.spi.deinit()
                for i in range(4):
                    self.neokey.pixels[i] = 0
            except AttributeError:
                pass
        atexit.register(atexit_shutdown_display)


    def backlight_off(self):
        self.backlight.duty_cycle = BACKLIGHT_OFF
        self.bl_enable = False
        self.set_text('')
        # Also make sure NeoKey NeoPixels are off
        for i in range(4):
            self.set_neokey_off(i)

    def backlight_on(self):
        self.backlight.duty_cycle = BACKLIGHT_ON
        self.bl_enable = True

    def set_text(self, txt):
        self.textbox.text = txt
        self.display.refresh()

    def load_totp_accounts(self):
        # Load TOTP account slot data from EEPROM database
        TOTPAccount = collections.namedtuple("TOTPAccount",
            ["slot", "label", "secret_b32"])
        try:
            eeprom = self.eeprom
            check_eeprom_format(eeprom)
            first_used_index = None
            for slot in range(1, 5):
                if is_slot_in_use(eeprom, slot):
                    label, secret_bytes = load_totp_account(eeprom, slot)
                    secret_b32 = base32_encode(secret_bytes)
                    self.accounts.append(TOTPAccount(slot, label, secret_b32))
                    if first_used_index is None:
                        first_used_index = slot-1
                else:
                    self.accounts.append(None)
            self.selected_acct = first_used_index
        except ValueError as e:
            print(e)

        # Print summary from loading the EEPROM account database
        print('Loaded data for %d TOTP account slots from EEPROM' % (
            len(self.accounts)))
        for i, a in enumerate(self.accounts):
            if a is None:
                print(" slot %d: -- empty --")
            else:
                print(" slot %d: '%s'" % (a.slot, a.label))
        if self.selected_acct is not None:
            print("Selected Slot:", self.accounts[self.selected_acct].slot)

    def select_account(self, slot):
        # Try to select requested slot (return False for empty slot)
        if (0 <= slot < len(self.accounts)):
            self.selected_acct = slot
        else:
            raise ValueError(f"slot number out of range: {slot}")

        if self.accounts[self.selected_acct] is None:
            # This slot is empty
            return False
        # Happy path: slot is configured with an account
        return True

    def get_selected_slot(self):
        return self.selected_acct

    def is_selected_slot_empty(self):
        # Return True if there's no account configured for the selected slot
        return self.accounts[self.selected_acct] is None

    def get_selected_totp(self, unix_time):
        # Get TOTP slot, label, and code for the selected account
        if self.selected_acct is None:
            return ('', '', '')
        acct = self.accounts[self.selected_acct]
        if acct is None:
            return ('', 'empty', '')
        # CAUTION! For boards that use adafruit_hashlib.hashlib for SHA1, this
        # call may take on the order of 2 seconds to finish
        code = totp_sha1(acct.secret_b32, unix_time, digits=6, period=30)
        return (acct.slot, acct.label, code)

    def datetime(self):
        return self.rtc.datetime

    def get_neokey_bits(self):
        # Convert the list of booleans from NeoKey1x4 to an integer bitfield so
        # it's easier to detect state changes using XOR and AND
        k = self.neokey.get_keys()
        return (NEOKEY0 * int(k[0]) + NEOKEY1 * int(k[1]) + NEOKEY2 * int(k[2])
            + NEOKEY3 * int(k[3]))

    def set_neokey(self, key, color):
        self.neokey.pixels[key] = color

    def set_neokey_off(self, key):
        self.neokey.pixels[key] = NP_OFF

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2025 Sam Blenny
#
# Helper routines for managing a 4KB 24LC32 I2C EEPROM as a database file for
# TOTP account info. The first 32 bytes (1 EEPROM page) are a header.
#
# Database file format:
# - The first 4 bytes (eeprom[0:4]) of a formatted EEPROM should be b'TOTP'
# - Next 4 bytes (eeprom[4:8]) of the header are markers to indicate which
#   account slots are in use
# - Remainder of header (eeprom[8:32]) is reserved
# - Account slots are 64 bytes. Slot 1 is at eeprom[32:96], slot 2 is at
#   eeprom[96:160], etc.
# - Account slot format is: 8 bytes null padded label, 24 bytes reserved, 32
#   bytes null filled TOTP secret (stored as raw bytes, *not* as base32 string)
#


def check_eeprom_format(eeprom):
    # Checks if the EEPROM is formatted with 'TOTP' magic bytes
    if eeprom[0:4] != b'TOTP':
        raise ValueError("EEPROM not formatted. Try util.format_eeprom().")


def is_slot_in_use(eeprom, slot):
    # Checks if the slot is in use by checking the in-use marker
    in_use_marker = 4 + (slot - 1)
    return eeprom[in_use_marker] == b'\xFF'


def load_totp_account(eeprom, slot):
    # Ensure slot number is valid and slot is in use
    if not (1 <= slot <= 4):
        raise ValueError("Invalid slot number. Must be between 1 and 4.")
    if not is_slot_in_use(eeprom, slot):
        raise ValueError(f"Slot {slot} is not in use.")

    # Define the base address for the slot
    base = 32 + (slot - 1) * 64

    # Load the label (trim nulls) and secret bytes
    label = eeprom[base:base + 8].decode('utf-8').rstrip('\x00')
    secret_bytes = eeprom[base + 32:base + 64]

    return label, secret_bytes

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2024 Sam Blenny
#
# Related Docs:
# - https://datatracker.ietf.org/doc/html/rfc2104  (HMAC RFC)
# - https://nvlpubs.nist.gov/nistpubs/fips/nist.fips.198-1.pdf (FIPS PUB 198-1)
#
try:
    import hashlib
except ImportError:
    import adafruit_hashlib as hashlib


def hmac(K, text, hash_fn, B, L):
    # Return HMAC digest for key and text using the named hash algorithm.
    #
    # The non-pythonic parameter names correspond to the notation used in
    # RFC 2104 and FIPS 198-1. The point of doing it this way is to make
    # it easier to compare this implementation against the specs.
    # Params:
    #   K: secret key (UTF-8 string or bytes)
    #   text: hash function input data (UTF-8 string or bytes)
    #   hash_fn: name of hash function for use with hashlib.new()
    #   B: block size in bytes of hash function input
    #   L: byte-length of hash function output
    # FIPS 198-1 Crypto Math Notation:
    #   ⨁: plus sign in a circle means bitwise XOR
    #   ||: two vertical bars means concatenate bitstreams

    # Convert text from string to bytes if needed
    if isinstance(text, str):
        text = text.encode('utf-8')
    # Convert key from string to bytes if needed
    if isinstance(K, str):
        K = K.encode('utf-8')
    # Prepare key as block sized byte buffer (hash or pad as needed)
    if len(K) == B:
        K0 = K
    elif len(K) > B:
        H = hashlib.new('sha1')
        H.update(K)
        K0 = H.digest() + (b'\x00' * (B - L))
    elif len(K) < B:
        K0 = K + (b'\x00' * (B - len(K)))
    # Calculate inner hash: H((K₀ ⨁ ipad)||text)
    K0_xored = bytearray([b ^ 0x36 for b in K0])
    H = hashlib.new('sha1')
    H.update(K0_xored)
    H.update(text)
    inner_hash = H.digest()
    # Calculate outer hash: H((K₀ ⨁ opad)||inner_digest)
    K0_xored = bytearray([b ^ 0x5c for b in K0])
    H = hashlib.new('sha1')
    H.update(K0_xored)
    H.update(inner_hash)
    return H.digest()


def hmac_sha1(key, data):
    # HMAC-SHA1 digest for key and data
    return hmac(K=key, text=data, hash_fn='sha1', B=64, L=20)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2024 Sam Blenny
#
# Related RFCs:
# - [RFC 4226: HOTP](https://datatracker.ietf.org/doc/html/rfc4226)
# - [RFC 6238: TOTP](https://datatracker.ietf.org/doc/html/rfc6238)
#
import struct

from sb_hmac import hmac_sha1


def normalize_base32(s):
    # Normalize a Base32 secret string from a TOTP URI.
    # This is meant to handle known TOTP "secret=..." query parameter quirks.
    # - Converts URL-encoded padding (%3D / %3d) to '='
    # - Strips whitespace
    # - Uppercases all letters
    #
    return s.strip().replace("%3D", "=").replace("%3d", "=").upper()


def parse_uri(uri):
    # Parse a TOTP QR code URI and extract the Base32 secret, validating that
    # algorithm=SHA1, digits=6, period=30
    # - Returns: Normalized Base32 secret string
    # - Raises: ValueError if the URI is malformed or parameters do not match
    #           expected values

    # 1. Validate scheme
    if not uri.startswith("otpauth://totp/"):
        raise ValueError("Invalid scheme or method (expected otpauth://totp/)")

    # 2. Split label and query
    try:
        _, query = uri.split("?", 1)
    except ValueError:
        raise ValueError("URI missing query string")

    # 3. Parse query manually
    params = {}
    for pair in query.split("&"):
        if "=" in pair:
            k, v = pair.split("=", 1)
            params[k] = v

    # 4. Extract secret and normalize it
    secret = params.get("secret")
    if not secret:
        raise ValueError("Missing secret")
    secret = normalize_base32(secret)  # handles padding, uppercasing, etc

    # 5. Extract and validate other fields
    algorithm = params.get("algorithm", "SHA1")
    digits = params.get("digits", "6")
    period = params.get("period", "30")

    if algorithm.upper() != "SHA1":
        raise ValueError("Unsupported algorithm (expected SHA1)")
    if digits != "6":
        raise ValueError("Unsupported digits (expected 6)")
    if period != "30":
        raise ValueError("Unsupported period (expected 30)")

    return secret


def base32_encode(data):
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"

    # Ensure the input is bytes
    if isinstance(data, str):
        data = data.encode('utf-8')

    # Add padding to make the data length a multiple of 5 bytes
    unpadded_len = len(data)
    padding = (5 - unpadded_len % 5) % 5
    data_padded = data + b'\x00' * padding

    # Create the encoded string by processing every 5 bytes
    encoded = []
    def append(bits):
        encoded.append(alphabet[bits])

    for i in range(0, len(data), 5):
        (a, b, c, d, e) = data_padded[i:i+5]
        last_iteration = bool(i+5 > unpadded_len)

        # For each chunk of 5 bytes, extract 5-bit groups
        append(  a >> 3)
        append((((a << 5) & 0xFF) >> 3) | (b >> 6))
        if last_iteration and padding == 4:
            encoded.append("======")
            break
        append((((b << 2) & 0xFF) >> 3))
        append((((b << 7) & 0xFF) >> 3) | (c >> 4))
        if last_iteration and padding == 3:
            encoded.append("====")
            break
        append((((c << 4) & 0xFF) >> 3) | (d >> 7))
        if last_iteration and padding == 2:
            encoded.append("===")
            break
        append( ((d << 1) & 0xFF) >> 3)
        append((((d << 6) & 0xFF) >> 3) | (e >> 5))
        if last_iteration and padding == 1:
            encoded.append("=")
            break
        append( ((e << 3) & 0xFF) >> 3)

    # Join the characters and remove the padding
    return ''.join(encoded)


def base32_decode(s):
    # Decode a Base32 string to bytes (ignores '=' and is case-insensitive)
    # CircuitPython doesn't appear to have this either built in or in a library
    alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
    s = s.strip().replace("=", "").upper()
    bits = 0
    bit_buffer = 0
    result = bytearray()

    for c in s:
        if c not in alphabet:
            raise ValueError("Invalid Base32 character: " + c)
        val = alphabet.index(c)
        bit_buffer = (bit_buffer << 5) | val
        bits += 5
        while bits >= 8:
            bits -= 8
            byte = (bit_buffer >> bits) & 0xff
            result.append(byte)

    return bytes(result)


def totp_sha1(secret_b32, timestamp, digits=6, period=30):
    # Return a string with the TOTP code for the given Unix timestamp. The
    # configurable timestamp is meant to support tests that check for the text
    # vectors from RFC6238 Appendix B.

    key =  base32_decode(secret_b32)

    counter = timestamp // period
    msg = struct.pack(">Q", counter)  # hash timestamp as big-endian int64

    hash_ = hmac_sha1(key, msg)

    # Do the dynamic truncation thing from RFC 6238 Appendix A
    offset = hash_[len(hash_)-1] & 0x0F
    # Extract the 31-bit integer following sample code in Appendix A
    binary = ((hash_[offset] & 0x7f) << 24 |
            (hash_[offset+1] & 0xff) << 16 |
            (hash_[offset+2] & 0xff) << 8 |
            (hash_[offset+3] & 0xff))

    # Do modulo (10^digits) so the final code is the right number of digits
    binary = str(binary % (10 ** digits))
    return ("0" * max(0, digits - len(binary))) + binary  # zero fill on left


# ------------------------------------------------------
# Tests
#
# >>> import sb_totp
# >>> sb_totp.test_totp_edge_cases()
# Running TOTP edge-case tests...
#
# RFC 6238 SHA1 tests:
# t=59 expected=94287082 got=94287082 OK
# t=1111111109 expected=07081804 got=07081804 OK
# t=1111111111 expected=14050471 got=14050471 OK
# t=1234567890 expected=89005924 got=89005924 OK
# t=2000000000 expected=69279037 got=69279037 OK
# t=20000000000 expected=65353130 got=65353130 OK
#
# RFC 2202 HMAC-SHA1 tests:
# Test 0: OK
# Test 1: OK
#
# Weird-but-valid Base32 secrets:
# AA -> OK (1 bytes)
# AE -> OK (1 bytes)
# AI -> OK (1 bytes)
# AM -> OK (1 bytes)
# AQ -> OK (1 bytes)
# AU -> OK (1 bytes)
# AY -> OK (1 bytes)
# A4====== -> OK (1 bytes)
# ABCDEFGHAA -> OK (6 bytes)
# ABAA==== -> OK (2 bytes)
# ACAA==== -> OK (2 bytes)
# ABCD%3D -> OK (2 bytes)
# EFGH%3d -> OK (2 bytes)
# Lowercase secret -> OK
#
# RFC 4648 Base32 test vectors (encode direction):
# t= expected='' got='' OK
# t=f expected='MY======' got='MY======' OK
# t=fo expected='MZXQ====' got='MZXQ====' OK
# t=foo expected='MZXW6===' got='MZXW6===' OK
# t=foob expected='MZXW6YQ=' got='MZXW6YQ=' OK
# t=fooba expected='MZXW6YTB' got='MZXW6YTB' OK
# t=foobar expected='MZXW6YTBOI======' got='MZXW6YTBOI======' OK
#
# RFC 4648 Base32 test vectors (decode direction):
# t= expected='' got='b''' OK
# t=f expected='f' got='b'f'' OK
# t=fo expected='fo' got='b'fo'' OK
# t=foo expected='foo' got='b'foo'' OK
# t=foob expected='foob' got='b'foob'' OK
# t=fooba expected='fooba' got='b'fooba'' OK
# t=foobar expected='foobar' got='b'foobar'' OK
#
# Round-Trip Base32 encode/decode tests:
# Test 0 - Round-trip: OK (input: b'hello', encoded: NBSWY3DP, decoded: b'hello')
# Test 1 - Round-trip: OK (input: b'1234567890', encoded: GEZDGNBVGY3TQOJQ, decoded: b'1234567890')
# Test 2 - Round-trip: OK (input: b'\x00\x01\x02\x03\x04\x05', encoded: AAAQEAYEAU======, decoded: b'\x00\x01\x02\x03\x04\x05')
# Test 3 - Round-trip: OK (input: b'ABCDEF1234567890', encoded: IFBEGRCFIYYTEMZUGU3DOOBZGA======, decoded: b'ABCDEF1234567890')
# Test 4 - Round-trip: OK (input: b'abcdefghijklmnopqrstuvwxyz', encoded: MFRGGZDFMZTWQ2LKNNWG23TPOBYXE43UOV3HO6DZPI======, decoded: b'abcdefghijklmnopqrstuvwxyz')
# >>>
# ------------------------------------------------------

def test_totp_edge_cases():
    print("Running TOTP edge-case tests...")

    # ------------------------------------------------------
    # 1. RFC 6238 SHA1 TOTP vectors
    # ------------------------------------------------------

    # RFC 6238 secret: "12345678901234567890"
    secret = "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"
    test_times = [
        59,
        1111111109,
        1111111111,
        1234567890,
        2000000000,
        20000000000,
    ]
    expected_totps = [
        "94287082",
        "07081804",
        "14050471",
        "89005924",
        "69279037",
        "65353130",
    ]
    print("\nRFC 6238 SHA1 tests:")
    for t, exp in zip(test_times, expected_totps):
        got = totp_sha1(secret, t, digits=8)
        print(f"t={t} expected={exp} got={got}", "OK" if got == exp else "FAIL")

    # ------------------------------------------------------
    # 2. HMAC-SHA1 tests (RFC 2202)
    # ------------------------------------------------------
    print("\nRFC 2202 HMAC-SHA1 tests:")
    test_cases = [
        (bytes([0x0b]*20), b"Hi There",
         bytes([0xb6,0x17,0x31,0x86,0x55,0x05,0x72,0x64,
                0xe2,0x8b,0xc0,0xb6,0xfb,0x37,0x8c,0x8e,
                0xf1,0x46,0xbe,0x00])),
        (b"Jefe", b"what do ya want for nothing?",
         bytes([0xef,0xfc,0xdf,0x6a,0xe5,0xeb,0x2f,0xa2,
                0xd2,0x74,0x16,0xd5,0xf1,0x84,0xdf,0x9c,
                0x25,0x9a,0x7c,0x79])),
    ]
    for i, (key, data, expected) in enumerate(test_cases):
        got = hmac_sha1(key, data)
        print(f"Test {i}:", "OK" if got == expected else "FAIL")

    # ------------------------------------------------------
    # 3. Weird-but-valid Base32 secrets
    # ------------------------------------------------------
    print("\nWeird-but-valid Base32 secrets:")
    secrets = [
        "AA", "AE", "AI", "AM", "AQ", "AU", "AY",
        "A4======", "ABCDEFGHAA", "ABAA====", "ACAA====",
        "ABCD%3D", "EFGH%3d"  # URL-encoded padding
    ]
    for s in secrets:
        try:
            decoded = base32_decode(normalize_base32(s))
            print(f"{s} -> OK ({len(decoded)} bytes)")
        except Exception as e:
            print(f"{s} -> FAIL ({e})")

    # Lowercase secret test
    try:
        decoded = base32_decode("jbswy3dpehpk3pxp")
        print("Lowercase secret -> OK")
    except Exception as e:
        print("Lowercase secret -> FAIL", e)


    # ------------------------------------------------------
    # 4. Base32 encode with RFC 4648 test vectors
    # ------------------------------------------------------
    test_data = [
        "",
        "f",
        "fo",
        "foo",
        "foob",
        "fooba",
        "foobar",
    ]
    base32_string = [
        "",
        "MY======",
        "MZXQ====",
        "MZXW6===",
        "MZXW6YQ=",
        "MZXW6YTB",
        "MZXW6YTBOI======"
    ]
    print("\nRFC 4648 Base32 test vectors (encode direction):")
    for td, b32 in zip(test_data, base32_string):
        got = base32_encode(td)
        print(f"t={td} expected='{b32}' got='{got}'", "OK" if got == b32 else "FAIL")

    print("\nRFC 4648 Base32 test vectors (decode direction):")
    for b32, td in zip(base32_string, test_data):
        got = base32_decode(b32)
        print(f"t={td} expected='{td}' got='{got}'",
            "OK" if got == td.encode('utf-8') else "FAIL")


    # ------------------------------------------------------
    # 5. Round-Trip Base32 Encode/Decode Tests
    # ------------------------------------------------------
    print("\nRound-Trip Base32 encode/decode tests:")
    test_data = [
        b"hello",  # Simple string
        b"1234567890",  # Numeric string
        b"\x00\x01\x02\x03\x04\x05",  # Binary data
        b"ABCDEF1234567890",  # Alphanumeric string
        b"abcdefghijklmnopqrstuvwxyz",  # Lowercase letters
    ]
    for i, data in enumerate(test_data):
        try:
            # Encode and then decode the data
            encoded = base32_encode(data)
            decoded = base32_decode(encoded)
            print(f"Test {i} - Round-trip: OK (input: {data}, encoded: "
                f"{encoded}, decoded: {decoded}) ")
            # Check that the round-trip decode matches the original data
            assert decoded == data, (f"Decoded data doesn't match original "
                f"(input: {data}, decoded: {decoded})")
        except Exception as e:
            print(f"Test {i} - Round-trip: FAIL ({e})")
            raise e

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270

# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright 2024 Sam Blenny
#
# This is meant to be imported from the CircuitPython REPL to set the DS3231
# clock after installing a new battery.
#
# See NOTES.md for documentation links.
#
# Usage example (CircuitPython REPL in USB serial console):
#
#     >>> import util
#     >>> util.menu()
#     Available functions:
#
#      1. add_totp_account()    - Add a new TOTP account.
#      2. copy_totp_account()   - Copy a TOTP account to another slot.
#      3. erase_totp_account()  - Erase a TOTP account from EEPROM.
#      4. format_eeprom()       - Format the EEPROM for TOTP storage.
#      5. get_time()            - Get the current DS3231 RTC time.
#      6. list_totp_accounts()  - List all stored TOTP accounts.
#      7. set_time()            - Set the DS3231 RTC time.
#
#     Choose a function by number (or Enter to cancel):
#     Operation canceled.
#     >>>
#
import board
import busio
import time
from time import mktime, sleep, struct_time

from adafruit_24lc32 import EEPROM_I2C
from adafruit_datetime import datetime
from adafruit_ds3231 import DS3231

from eeprom_db import check_eeprom_format, is_slot_in_use, load_totp_account
from sb_totp import base32_decode, parse_uri


i2c = busio.I2C(board.SCL, board.SDA, frequency=250_000)
rtc = DS3231(i2c)
eeprom = EEPROM_I2C(i2c)


def set_time():
    print("Set DS3231 RTC time...")
    try:
        y    = int(input("   year: "))
        mon  = int(input("  month: "))
        d    = int(input("    day: "))
        h    = int(input("   hour: "))
        min_ = int(input(" minute: "))
        s    = int(input("seconds: "))
        t = struct_time((y, mon, d, h, min_, s, 0, -1, -1))
        rtc.datetime = t
        print("new RTC time: ", now())
    except ValueError as e:
        print("ERROR Bad value:", e)


def now():
    # Return RTC time formatted as a string
    return "%04d-%02d-%02d %02d:%02d:%02d" % ((rtc.datetime)[0:6])


def get_time():
    print(now())


def get_slot_input(tag=None):
    # Prompts the user for a valid slot number
    prompt = "Enter slot number (1-4): "
    if tag is not None:
        prompt = f"{tag}: {prompt}"
    slot = int(input(prompt))
    if not (1 <= slot <= 4):
        raise ValueError("Invalid slot number. Must be between 1 and 4.")
    return slot


def confirm_overwrite(slot):
    # Prompts the user to confirm overwriting a slot if it is in use
    if is_slot_in_use(eeprom, slot):  # Slot is in use
        overwrite = input("Slot is in use. Overwrite? (y/n): ").strip().lower()
        if overwrite != 'y':
            print("Operation canceled.")
            return False
    return True


def format_eeprom():
    # 1. Confirm if user really wants to format the EEPROM
    yn = input("Are you sure? This will erase all data. (y/n): ")
    if yn.strip().lower() != 'y':
        print("Operation canceled.")
        return

    # 2. Check EEPROM length
    if (len_ := len(eeprom)) != 4096:
        raise ValueError("Expected EEPROM length 4096 bytes, got:", len_)

    # 3. Write header with 'TOTP' magic bytes and reserved fields
    print("Writing header page...")
    eeprom[0:32] = b'TOTP' + (b'\x00' * 28) # magic + nulls for reserved bytes

    # 4. Zero all record slots (each record is 64 bytes or 2 pages)
    blank_page = b'\x00' * 32
    print("Zeroing 127 pages:")
    for i in range(1, 128):  # 127 pages, 32 bytes each, skip page 0 (header)
        print(".", end=('\n' if i % 64 == 0 else ''))
        start = i * 32
        eeprom[start:start+32] = blank_page

    print("\nEEPROM formatted successfully.")


def write_slot_data(slot, label_padded, secret_bytes):
    # Writes label and secret data to the specified slot
    base = 32 + (slot - 1) * 64
    eeprom[base+0:base+8] = label_padded      # 8 null padded label bytes
    eeprom[base+8:base+32] = b'\x00' * 24     # 24 reserved bytes
    eeprom[base+32:base+64] = secret_bytes    # 32 null padded secret bytes


def add_totp_account():
    # 1. Check for 'TOTP' magic bytes in the EEPROM
    check_eeprom_format(eeprom)

    # 2. Prompt for slot number
    slot = get_slot_input()

    # 3. Check if the slot is already in use
    if is_slot_in_use(eeprom, slot):
        if not confirm_overwrite(slot):
            return

    # 4. Prompt for label (max 8 chars)
    label_bytes = input("Enter label (max 8 chars): ").strip().encode('utf-8')
    if len(label_bytes) > 8:
        raise ValueError("Label too long. Max 8 utf-8 bytes.")
    label_padded = label_bytes + b'\x00' * (8 - len(label_bytes))  # null pad

    # 5. Prompt for TOTP URI
    uri = input("Enter TOTP URI: ").strip()

    # 6. Parse URI, extract secret and decode
    secret_b32 = parse_uri(uri)
    secret_bytes = base32_decode(secret_b32)  # Decode the Base32 secret

    # Check if the secret is too long
    if len(secret_bytes) > 32:
        raise ValueError("Too many secret bytes. Must be 32 bytes or fewer.")

    # Pad the secret to 32 bytes if shorter (with 0x00 padding). This is okay
    # to do now because the HMAC hash will pad short keys later anyway.
    secret_bytes += b'\x00' * (32 - len(secret_bytes))

    # 7. Write label and secret to the EEPROM slot
    write_slot_data(slot, label_padded, secret_bytes)

    # 8. Mark slot in use in the header
    in_use_marker = 4 + (slot - 1)
    eeprom[in_use_marker] = 0xFF

    print(f"Record added to slot {slot}.")


def erase_totp_account():
    # 1. Check for 'TOTP' magic bytes in the EEPROM
    check_eeprom_format(eeprom)

    # 2. Prompt for slot number
    slot = get_slot_input()

    # 3. Erase the slot: set all its bytes to 0x00
    base = 32 + (slot - 1) * 64
    eeprom[base:base + 64] = b'\x00' * 64  # Clear the slot's record (64 bytes)

    # 4. Mark the slot as free in the header
    in_use_marker = 4 + (slot - 1)
    eeprom[in_use_marker] = 0x00  # Mark slot as not in use

    print(f"Slot {slot} has been erased.")


def list_totp_accounts():
    # Check for 'TOTP' magic bytes in the EEPROM
    check_eeprom_format(eeprom)

    # Iterate through each of the 4 possible slots
    for slot in range(1, 5):
        try:
            # Load label and secret from EEPROM for the given slot
            label, _ = load_totp_account(eeprom, slot)
            print(f"Slot {slot}: '{label}'")
        except ValueError:
            # If the slot isn't in use, we get a ValueError
            print(f"Slot {slot}: -- empty --")


def copy_totp_account():
    # Prompt for the source slot number
    src_slot = get_slot_input("Source")

    # Check if the source slot is in use
    if not is_slot_in_use(eeprom, src_slot):
        raise ValueError(f"Source slot {src_slot} is not in use.")

    # Prompt for the destination slot number
    dest_slot = get_slot_input("Destination")

    # Check if the destination slot is in use
    if is_slot_in_use(eeprom, dest_slot):
        # Prompt to confirm overwriting the destination slot
        if not confirm_overwrite(dest_slot):
            return

    # Define the base addresses for the source and destination slots
    src_base = 32 + (src_slot - 1) * 64
    dest_base = 32 + (dest_slot - 1) * 64

    # Copy TOTP account data (label + secret) from the source to destination
    eeprom[dest_base:dest_base + 64] = eeprom[src_base:src_base + 64]

    # Update the in-use marker for the destination slot
    eeprom[4 + (dest_slot - 1)] = 0xFF

    print(f"TOTP account copied from slot {src_slot} to slot {dest_slot}.")


def menu():
    # List the available functions with brief descriptions
    print("Available functions:")
    print()
    print(" 1. add_totp_account()    - Add a new TOTP account.")
    print(" 2. copy_totp_account()   - Copy a TOTP account to another slot.")
    print(" 3. erase_totp_account()  - Erase a TOTP account from EEPROM.")
    print(" 4. format_eeprom()       - Format the EEPROM for TOTP storage.")
    print(" 5. get_time()            - Get the current DS3231 RTC time.")
    print(" 6. list_totp_accounts()  - List all stored TOTP accounts.")
    print(" 7. set_time()            - Set the DS3231 RTC time.")
    print()

    # Prompt for which function to invoke
    choice = input(
        "Choose a function by number (or Enter to cancel): ").strip()

    # If the user entered an empty string, do nothing
    if choice == "":
        print("Operation canceled.")
        return

    # Convert input to an integer and invoke the corresponding function
    choice = int(choice)
    if choice == 1:
        add_totp_account()
    elif choice == 2:
        copy_totp_account()
    elif choice == 3:
        erase_totp_account()
    elif choice == 4:
        format_eeprom()
    elif choice == 5:
        get_time()
    elif choice == 6:
        list_totp_accounts()
    elif choice == 7:
        set_time()
    else:
        raise ValueError(f"Invalid choice: {choice}")