"""Hardware-free unit tests for apex_otp_import (OATH protocol layer + input)."""
import os
import tempfile
import unittest

from openpyxl import Workbook

from apex_otp_import import (
    tlv, parse_tlvs, decode_secret, pad_key, format_cred_id,
    build_put_data, build_validate_data, build_set_code_data,
    derive_key, compute_response,
    Credential, validate_input, load_credentials,
    CardChannel, OathSession, OathError, parse_memory,
    TOTP_TYPE, HOTP_TYPE, HMAC_SHA1, HMAC_SHA256, HMAC_SHA512,
)


class FakeConn:
    """Fake pyscard connection: replays a scripted list of (data, sw1, sw2)."""
    def __init__(self, script):
        self.script = list(script)
        self.sent = []

    def transmit(self, apdu):
        self.sent.append(apdu)
        data, sw1, sw2 = self.script.pop(0)
        return list(data), sw1, sw2


def cred(issuer="Iss", account="acct", secret="GEZDGNBVGY3TQOJQ", **kw):
    return Credential(issuer=issuer, account=account, secret=secret, **kw)


class TestTlv(unittest.TestCase):
    def test_short_form_length(self):
        self.assertEqual(tlv(0x71, b"AB"), bytes([0x71, 0x02, 0x41, 0x42]))

    def test_boundary_127_is_short_form(self):
        self.assertEqual(tlv(0x73, b"\x00" * 127)[:2], bytes([0x73, 0x7F]))

    def test_extended_form_128(self):
        # length 128 uses 0x81 <len>
        self.assertEqual(tlv(0x73, b"\x00" * 128)[:3], bytes([0x73, 0x81, 0x80]))


class TestParseTlvs(unittest.TestCase):
    def test_parses_ordered_pairs(self):
        data = bytes([0x79, 0x03, 0x00, 0x01, 0x02, 0x71, 0x02, 0xAA, 0xBB])
        self.assertEqual(
            parse_tlvs(data),
            [(0x79, b"\x00\x01\x02"), (0x71, b"\xAA\xBB")],
        )

    def test_preserves_repeated_tags(self):
        # LIST returns repeated 0x72 entries
        data = tlv(0x72, b"\x21A") + tlv(0x72, b"\x21B")
        self.assertEqual(parse_tlvs(data), [(0x72, b"\x21A"), (0x72, b"\x21B")])

    def test_round_trips_extended_length(self):
        val = bytes(range(256)) * 2  # 512 bytes -> 0x82 form
        self.assertEqual(parse_tlvs(tlv(0x73, val)), [(0x73, val)])


class TestSecret(unittest.TestCase):
    def test_rfc6238_vector_decodes_to_known_bytes(self):
        # RFC 6238 seed "12345678901234567890" == this Base32
        self.assertEqual(
            decode_secret("GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"),
            b"12345678901234567890",
        )

    def test_lowercase_and_spaces_and_missing_padding(self):
        self.assertEqual(decode_secret("ge zd gn bv"), decode_secret("GEZDGNBV"))

    def test_pad_short_key_to_14_zero_filled(self):
        raw = decode_secret("JBSWY3DPEHPK3PXP")  # 16 chars -> 10 bytes
        self.assertEqual(len(raw), 10)
        padded = pad_key(raw)
        self.assertEqual(len(padded), 14)
        self.assertEqual(padded[:10], raw)
        self.assertEqual(padded[10:], b"\x00\x00\x00\x00")

    def test_pad_leaves_long_enough_key_unchanged(self):
        raw = b"\x01" * 20
        self.assertEqual(pad_key(raw), raw)


class TestCredId(unittest.TestCase):
    def test_default_period_is_issuer_colon_account(self):
        self.assertEqual(format_cred_id("Google", "amal@amal.net", 30),
                         "Google:amal@amal.net")

    def test_non_default_period_prefixed(self):
        self.assertEqual(format_cred_id("Steam", "user", 15), "15/Steam:user")

    def test_empty_issuer_is_just_account(self):
        self.assertEqual(format_cred_id("", "amal", 30), "amal")


class TestPutData(unittest.TestCase):
    def test_byte_exact_put_payload(self):
        key = b"\x01" * 14
        data = build_put_data("AB", TOTP_TYPE | HMAC_SHA1, 6, key)
        expected = (
            bytes([0x71, 0x02, 0x41, 0x42])           # NAME "AB"
            + bytes([0x73, 0x10, 0x21, 0x06]) + key    # KEY: len=2+14, typealgo, digits
        )
        self.assertEqual(data, expected)

    def test_totp_sha1_type_byte_is_0x21(self):
        self.assertEqual(TOTP_TYPE | HMAC_SHA1, 0x21)


class TestValidateData(unittest.TestCase):
    def test_byte_exact_validate_payload(self):
        resp = b"\xAA" * 20
        chal = b"\xBB" * 8
        data = build_validate_data(resp, chal)
        self.assertEqual(data, bytes([0x75, 0x14]) + resp + bytes([0x74, 0x08]) + chal)


class TestSetCodeData(unittest.TestCase):
    def test_byte_exact_set_code_payload(self):
        key = b"\x11" * 16
        chal = b"\x22" * 8
        data = build_set_code_data(key, chal)
        resp = compute_response(key, chal, HMAC_SHA1)
        expected = (tlv(0x73, bytes([0x21]) + key)     # KEY: TOTP|SHA1 + 16-byte key
                    + tlv(0x74, chal)                   # CHALLENGE
                    + tlv(0x75, resp))                  # RESPONSE = HMAC(key, challenge)
        self.assertEqual(data, expected)


class TestPasswordAuth(unittest.TestCase):
    SALT = bytes([0x19, 0x83, 0x57, 0x1E, 0xC3, 0xE5, 0xAB, 0xD7])

    def test_derive_key_pbkdf2_1000_rounds_16_bytes(self):
        key = derive_key("test1234", self.SALT)
        self.assertEqual(key.hex(), "aa69a3f71eec1660849d7bb16c313e32")
        self.assertEqual(len(key), 16)

    def test_compute_response_matches_rfc2202_hmac_sha1(self):
        # RFC 2202 HMAC-SHA1 test case 1
        resp = compute_response(b"\x0b" * 20, b"Hi There", HMAC_SHA1)
        self.assertEqual(resp.hex(), "b617318655057264e28bc0b6fb378c8ef146be00")


class TestValidateInput(unittest.TestCase):
    def test_clean_input_has_no_problems(self):
        creds = [cred(account="a@x.com"), cred(account="b@x.com")]
        self.assertEqual(validate_input(creds), [])

    def test_duplicate_cred_id_is_reported(self):
        # same issuer+account, different secret -> the SAW/michelle case
        creds = [
            cred(issuer="SAW", account="m@d.com", secret="GEZDGNBVGY3TQOJQ", source="a.png"),
            cred(issuer="SAW", account="m@d.com", secret="MZXW6YTBOI======", source="b.png"),
        ]
        problems = validate_input(creds)
        self.assertEqual(len(problems), 1)
        self.assertIn("SAW:m@d.com", problems[0])

    def test_invalid_base32_is_reported(self):
        problems = validate_input([cred(secret="not-base32!")])
        self.assertTrue(any("base32" in p.lower() for p in problems))

    def test_oversize_name_is_reported(self):
        problems = validate_input([cred(issuer="I" * 40, account="a" * 40)])
        self.assertTrue(any("64" in p for p in problems))

    def test_sha512_algorithm_is_reported(self):
        problems = validate_input([cred(algorithm=HMAC_SHA512)])
        self.assertTrue(any("sha512" in p.lower() or "algorithm" in p.lower()
                            for p in problems))


class TestLoadCredentials(unittest.TestCase):
    def _make_xlsx(self, rows):
        wb = Workbook()
        ws = wb.active
        ws.title = "OTP Keys"
        ws.append(["Issuer", "Account", "Secret (Base32)", "Type", "Algorithm",
                   "Digits", "Period (s)", "QR Label", "otpauth:// URI", "Source File"])
        for r in rows:
            ws.append(r)
        fd, path = tempfile.mkstemp(suffix=".xlsx")
        os.close(fd)
        wb.save(path)
        return path

    def test_maps_columns_and_string_enums(self):
        path = self._make_xlsx([
            ["Google", "a@x.com", "GEZDGNBVGY3TQOJQ", "TOTP", "SHA1", 6, 30, "lbl", "uri", "g.png"],
            ["Steam", "user", "MZXW6YTBOI======", "TOTP", "SHA256", 8, 15, "lbl", "uri", "s.png"],
        ])
        try:
            creds = load_credentials(path)
        finally:
            os.remove(path)
        self.assertEqual(len(creds), 2)
        g = creds[0]
        self.assertEqual((g.issuer, g.account, g.secret), ("Google", "a@x.com", "GEZDGNBVGY3TQOJQ"))
        self.assertEqual((g.oath_type, g.algorithm, g.digits, g.period), (TOTP_TYPE, HMAC_SHA1, 6, 30))
        self.assertEqual(g.source, "g.png")
        self.assertEqual((creds[1].algorithm, creds[1].digits, creds[1].period), (HMAC_SHA256, 8, 15))


class TestCardChannel(unittest.TestCase):
    def test_reassembles_61xx_via_send_remaining(self):
        conn = FakeConn([
            (b"\x72\x02\x21A", 0x61, 0x04),   # first chunk + "4 more"
            (b"\x72\x02\x21B", 0x90, 0x00),   # remainder + OK
        ])
        ch = CardChannel(conn)
        data, sw = ch.send(0x00, 0xA1, 0x00, 0x00, le=0x00)  # LIST
        self.assertEqual(sw, 0x9000)
        self.assertEqual(data, b"\x72\x02\x21A\x72\x02\x21B")
        # second APDU must be SEND REMAINING (INS 0xA5)
        self.assertEqual(conn.sent[1][:2], [0x00, 0xA5])

    def test_single_response_returns_sw(self):
        conn = FakeConn([(b"", 0x90, 0x00)])
        data, sw = CardChannel(conn).send(0x00, 0x01, 0x00, 0x00, data=b"\x71\x01A")
        self.assertEqual(sw, 0x9000)
        self.assertEqual(data, b"")


DEV_ID = bytes(range(8))            # 00..07
VERSION_TLV = tlv(0x79, b"\x00\x01\x02")


def select_unlocked():
    return VERSION_TLV + tlv(0x71, DEV_ID)


def select_locked(challenge):
    # real applet reports the full type|algo byte (0x21 = TOTP|SHA1), not 0x01
    return (VERSION_TLV + tlv(0x71, DEV_ID)
            + tlv(0x74, challenge) + tlv(0x7B, bytes([TOTP_TYPE | HMAC_SHA1])))


class TestOathSessionSelect(unittest.TestCase):
    def test_parses_unlocked_select(self):
        s = OathSession(CardChannel(FakeConn([(select_unlocked(), 0x90, 0x00)])))
        s.select()
        self.assertFalse(s.locked)
        self.assertEqual(s.device_id, DEV_ID)
        self.assertEqual(s.version, (0, 1, 2))

    def test_parses_locked_select(self):
        chal = b"\xC1" * 8
        s = OathSession(CardChannel(FakeConn([(select_locked(chal), 0x90, 0x00)])))
        s.select()
        self.assertTrue(s.locked)
        self.assertEqual(s.challenge, chal)
        self.assertEqual(s.algorithm, HMAC_SHA1)

    def test_falls_back_to_next_aid_on_failure(self):
        conn = FakeConn([(b"", 0x6A, 0x82),                 # first AID: not found
                         (select_unlocked(), 0x90, 0x00)])  # second AID: ok
        s = OathSession(CardChannel(conn))
        s.select()
        self.assertFalse(s.locked)
        self.assertEqual(len(conn.sent), 2)  # tried two AIDs


class TestOathSessionList(unittest.TestCase):
    def test_list_names_parses_entries(self):
        listing = (tlv(0x72, bytes([0x21]) + b"Google:a@x.com")
                   + tlv(0x72, bytes([0x21]) + b"AWS:b"))
        conn = FakeConn([(select_unlocked(), 0x90, 0x00), (listing, 0x90, 0x00)])
        s = OathSession(CardChannel(conn)); s.select()
        self.assertEqual(s.list_names(), ["Google:a@x.com", "AWS:b"])


class TestOathSessionPut(unittest.TestCase):
    def _cred(self):
        return Credential("Iss", "acct", "GEZDGNBVGY3TQOJQ")

    def test_put_success_sends_put_apdu(self):
        conn = FakeConn([(select_unlocked(), 0x90, 0x00), (b"", 0x90, 0x00)])
        s = OathSession(CardChannel(conn)); s.select()
        s.put(self._cred())
        self.assertEqual(conn.sent[1][1], 0x01)                  # INS PUT
        self.assertEqual(bytes(conn.sent[1][5:]), self._cred().put_data())

    def test_put_full_raises_with_sw(self):
        conn = FakeConn([(select_unlocked(), 0x90, 0x00), (b"", 0x6A, 0x84)])
        s = OathSession(CardChannel(conn)); s.select()
        with self.assertRaises(OathError) as ctx:
            s.put(self._cred())
        self.assertEqual(ctx.exception.sw, 0x6A84)


class TestParseMemory(unittest.TestCase):
    def test_parses_available_and_total(self):
        # avail=84000, total=84336 (0x00014970), tRes=0x0100, tDes=0x0080
        data = (84000).to_bytes(4, "big") + bytes.fromhex("00014970") \
            + bytes.fromhex("0100") + bytes.fromhex("0080")
        m = parse_memory(data)
        self.assertEqual(m["available_persistent"], 84000)
        self.assertEqual(m["total_persistent"], 84336)


class TestOathSessionValidate(unittest.TestCase):
    def test_validate_success_sends_correct_response(self):
        from apex_otp_import import derive_key, compute_response
        chal = b"\xC1" * 8
        client = b"\xD2" * 8
        key = derive_key("hunter2", DEV_ID)
        device_reply = tlv(0x75, compute_response(key, client))
        conn = FakeConn([(select_locked(chal), 0x90, 0x00),
                         (device_reply, 0x90, 0x00)])
        s = OathSession(CardChannel(conn), rng=lambda n: client)
        s.select()
        s.validate("hunter2")                                   # must not raise
        # our response to the device must be HMAC(key, device challenge)
        sent = bytes(conn.sent[1])
        self.assertIn(compute_response(key, chal), sent)

    def test_validate_wrong_password_raises(self):
        conn = FakeConn([(select_locked(b"\xC1" * 8), 0x90, 0x00),
                         (b"", 0x6A, 0x80)])
        s = OathSession(CardChannel(conn), rng=lambda n: b"\xD2" * 8)
        s.select()
        with self.assertRaises(OathError):
            s.validate("wrong")


class TestOathSessionAdmin(unittest.TestCase):
    def test_set_password_sends_set_code(self):
        from apex_otp_import import derive_key
        client = b"\x22" * 8
        conn = FakeConn([(select_unlocked(), 0x90, 0x00), (b"", 0x90, 0x00)])
        s = OathSession(CardChannel(conn), rng=lambda n: client)
        s.select()
        s.set_password("test1234")
        self.assertEqual(conn.sent[1][1], 0x03)  # INS SET CODE
        key = derive_key("test1234", DEV_ID)
        self.assertEqual(bytes(conn.sent[1][5:]), build_set_code_data(key, client))

    def test_reset_sends_reset_with_dead(self):
        conn = FakeConn([(select_unlocked(), 0x90, 0x00), (b"", 0x90, 0x00)])
        s = OathSession(CardChannel(conn)); s.select()
        s.reset()
        self.assertEqual(conn.sent[1][:4], [0x00, 0x04, 0xDE, 0xAD])


if __name__ == "__main__":
    unittest.main()