"""Hardening tests: the findings from the self-review.

Each test here corresponds to a defect that was found and fixed. If any of
these regress, the corresponding vulnerability is back.
"""

import base64
import os
import tempfile
import threading
import time
import unittest

from sigil import errors, identity, keys, sessions, viewcodes
from sigil.store import Store

import helpers


def enrolled(store=None, agent="bigpickle"):
    store = store or Store(":memory:")
    private, _ = helpers.enrol(store, agent)
    return store, private, None


sign = helpers.sign


def issue(store, private, agent="bigpickle", **kw):
    return helpers.issue(store, private, agent, **kw)


class TestSingleUseUnderRace(unittest.TestCase):
    """Finding: check-then-update let one code mint several sessions.

    Measured before the fix: a single-use code succeeded 753 times across 200
    trials. These tests hold the invariant under real concurrency.
    """

    TRIALS = 40
    THREADS = 8

    def _race(self, setup, threads=THREADS):
        """Set up one reusable credential per trial, then race for it.

        setup returns (store, shared, private). `shared` is the single-use
        credential the threads all compete for, created once outside the
        threads, so this genuinely tests the same-secret-twice path.
        """
        for _ in range(self.TRIALS):
            store, shared, private = setup()
            barrier = threading.Barrier(threads)
            outcomes = []

            def go():
                barrier.wait()
                try:
                    work(store, shared, private)
                    outcomes.append(1)
                except Exception:
                    outcomes.append(0)

            pool = [threading.Thread(target=go) for _ in range(threads)]
            for t in pool:
                t.start()
            for t in pool:
                t.join()
            yield sum(outcomes)

    def test_one_code_yields_at_most_one_session(self):
        def setup():
            store, private, _ = enrolled()
            return store, helpers.issue(store, private)["code"], private

        def work(store, code, private):
            viewcodes.redeem_code(store, code)

        for successes in self._race(setup):
            self.assertLessEqual(successes, 1, "one code minted multiple sessions")

    def test_one_bootstrap_secret_enrols_at_most_once(self):
        def setup():
            store, private, _ = enrolled()
            return store, identity.seed_bootstrap(store, "racer"), private

        def work(store, secret, private):
            identity.enrol(
                store, "racer", keys.export_public_b64(keys.generate()[1]), secret
            )

        for successes in self._race(setup, threads=6):
            self.assertLessEqual(successes, 1, "one secret enrolled multiple keys")

    def test_one_nonce_verifies_at_most_once(self):
        def setup():
            store, private, _ = enrolled()
            nonce = identity.challenge(store, "bigpickle")["nonce"]
            return store, nonce, private

        def work(store, nonce, private):
            identity.verify(store, nonce, sign(private, nonce))

        for successes in self._race(setup, threads=6):
            self.assertLessEqual(successes, 1, "one nonce produced multiple sessions")


class TestTokensNotStoredInClear(unittest.TestCase):
    """Finding: session tokens were stored verbatim, so a database leak
    handed an attacker live sessions."""

    def setUp(self):
        self.store, self.private, self.public = enrolled()

    def _agent_token(self):
        chal = identity.challenge(self.store, "bigpickle")
        return identity.verify(self.store, chal["nonce"], sign(self.private, chal["nonce"]))["token"]

    def test_agent_token_appears_nowhere_in_the_database_file(self):
        import sqlite3

        path = tempfile.mktemp(suffix=".db")
        store = Store(path)
        private, public = keys.generate()
        secret = identity.seed_bootstrap(store, "bigpickle")
        identity.enrol(store, "bigpickle", keys.export_public_b64(public), secret)
        chal = identity.challenge(store, "bigpickle")
        token = identity.verify(
            store, chal["nonce"], sign(private, chal["nonce"])
        )["token"]
        store.close()
        with open(path, "rb") as handle:
            blob = handle.read()
        os.unlink(path)
        self.assertNotIn(token.encode(), blob, "session token stored in the clear")

    def test_human_token_appears_nowhere_in_the_database_file(self):
        import sqlite3

        path = tempfile.mktemp(suffix=".db")
        store = Store(path)
        private, public = keys.generate()
        secret = identity.seed_bootstrap(store, "bigpickle")
        identity.enrol(store, "bigpickle", keys.export_public_b64(public), secret)
        code = helpers.issue(store, private)["code"]
        token = viewcodes.redeem_code(store, code)["token"]
        store.close()
        with open(path, "rb") as handle:
            blob = handle.read()
        os.unlink(path)
        self.assertNotIn(token.encode(), blob, "view session token stored in the clear")

    def test_hashed_token_still_resolves(self):
        token = self._agent_token()
        self.assertEqual(
            sessions.resolve(self.store, token)["on_behalf_of"], "bigpickle"
        )


class TestAuditChain(unittest.TestCase):
    """The log is an attestation substrate, so it must be tamper-evident."""

    def test_empty_log_verifies(self):
        self.assertTrue(Store(":memory:").verify_chain()["ok"])

    def test_intact_log_verifies(self):
        store, private, public = enrolled()
        for _ in range(5):
            identity.challenge(store, "bigpickle")
        result = store.verify_chain()
        self.assertTrue(result["ok"])
        self.assertEqual(result["events"], len(store.events(limit=50)))

    def test_edited_detail_is_detected(self):
        store, private, public = enrolled()
        identity.challenge(store, "bigpickle")
        store.conn.execute(
            "UPDATE events SET detail = '{\"tampered\":true}' WHERE id = 1"
        )
        store.conn.commit()
        result = store.verify_chain()
        self.assertFalse(result["ok"])
        self.assertEqual(result["broken_at"], 1)

    def test_deleted_event_is_detected(self):
        store, private, public = enrolled()
        identity.challenge(store, "bigpickle")
        store.conn.execute("DELETE FROM events WHERE id = 1")
        store.conn.commit()
        self.assertFalse(store.verify_chain()["ok"])

    def test_reordered_events_are_detected(self):
        store, private, public = enrolled()
        for _ in range(3):
            identity.challenge(store, "bigpickle")
        store.conn.execute("UPDATE events SET kind = 'tampered' WHERE id = 2")
        store.conn.commit()
        self.assertFalse(store.verify_chain()["ok"])


class TestKeyRotationKillsSessions(unittest.TestCase):
    """Finding: rekey replaced the key but left existing sessions alive, so
    rotating after a compromise bought nothing."""

    def setUp(self):
        self.store, self.private, self.public = enrolled()

    def _token(self):
        chal = identity.challenge(self.store, "bigpickle")
        return identity.verify(
            self.store, chal["nonce"], sign(self.private, chal["nonce"])
        )["token"]

    def test_rekey_revokes_existing_sessions(self):
        token = self._token()
        sessions.resolve(self.store, token)  # alive before
        secret = identity.seed_bootstrap(self.store, "bigpickle")
        identity.enrol_rekey(
            self.store, "bigpickle",
            keys.export_public_b64(keys.generate()[1]), secret,
        )
        with self.assertRaises(errors.SessionExpired):
            sessions.resolve(self.store, token)

    def test_old_key_cannot_authenticate_after_rekey(self):
        secret = identity.seed_bootstrap(self.store, "bigpickle")
        identity.enrol_rekey(
            self.store, "bigpickle",
            keys.export_public_b64(keys.generate()[1]), secret,
        )
        chal = identity.challenge(self.store, "bigpickle")
        with self.assertRaises(errors.BadSignature):
            identity.verify(self.store, chal["nonce"], sign(self.private, chal["nonce"]))

    def test_rotation_is_recorded(self):
        secret = identity.seed_bootstrap(self.store, "bigpickle")
        identity.enrol_rekey(
            self.store, "bigpickle",
            keys.export_public_b64(keys.generate()[1]), secret,
        )
        kinds = [e["kind"] for e in self.store.events(limit=50)]
        self.assertIn("identity.rekeyed", kinds)


class TestRateLimits(unittest.TestCase):
    """Finding: challenge() was unmetered, an open signing oracle."""

    def setUp(self):
        self.store, self.private, self.public = enrolled()

    def test_challenge_is_rate_limited(self):
        raised = False
        for _ in range(200):
            try:
                identity.challenge(self.store, "bigpickle")
            except errors.RateLimited:
                raised = True
                break
        self.assertTrue(raised, "challenge() was never rate limited")

    def test_code_issue_is_rate_limited(self):
        raised = False
        for _ in range(200):
            try:
                helpers.issue(self.store, self.private)
            except errors.RateLimited:
                raised = True
                break
        self.assertTrue(raised, "issue_code() was never rate limited")

    def test_redemption_guessing_is_rate_limited(self):
        refused = 0
        for i in range(500):
            try:
                viewcodes.redeem_code(self.store, f"ZZZZ-{i:04d}")
            except errors.RateLimited:
                break
            except errors.SigilError:
                refused += 1
        else:
            self.fail("code guessing was not rate limited")

    def test_limits_are_per_agent(self):
        for _ in range(40):
            try:
                identity.challenge(self.store, "bigpickle")
            except errors.RateLimited:
                break
        # a different agent is unaffected
        self.store2, _, _ = enrolled(self.store, agent="second")
        identity.challenge(self.store2, "second")


class TestKeyEncryptionAtRest(unittest.TestCase):
    """Finding: keys were written with NoEncryption by default."""

    def test_encrypted_key_needs_passphrase(self):
        private, _ = keys.generate()
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "k.pem")
            keys.save_private(private, path, "a-long-enough-passphrase")
            self.assertTrue(keys.is_encrypted(path))
            with self.assertRaises(ValueError):
                keys.load_private(path)

    def test_wrong_passphrase_refuses(self):
        private, _ = keys.generate()
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "k.pem")
            keys.save_private(private, path, "a-long-enough-passphrase")
            with self.assertRaises(Exception):
                keys.load_private(path, "the-wrong-passphrase")

    def test_correct_passphrase_roundtrips(self):
        private, _ = keys.generate()
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "k.pem")
            keys.save_private(private, path, "a-long-enough-passphrase")
            loaded = keys.load_private(path, "a-long-enough-passphrase")
            self.assertEqual(loaded.sign(b"m"), private.sign(b"m"))

    def test_short_passphrase_refused(self):
        private, _ = keys.generate()
        with tempfile.TemporaryDirectory() as tmp:
            with self.assertRaises(ValueError):
                keys.save_private(private, os.path.join(tmp, "k.pem"), "short")

    def test_plaintext_key_still_loads(self):
        private, _ = keys.generate()
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "k.pem")
            keys.save_private(private, path)
            self.assertFalse(keys.is_encrypted(path))
            keys.load_private(path)


class TestAgentIdValidation(unittest.TestCase):
    """Finding: agent ids were unvalidated, so an id carrying a newline
    could forge audit lines."""

    def test_bad_ids_refused(self):
        store = Store(":memory:")
        for bad in ["", "Big Pickle", "a" * 65, "with\nnewline", "with space",
                    "-leading", "with/slash", "with;drop"]:
            with self.assertRaises(errors.SigilError, msg=f"accepted {bad!r}"):
                identity.seed_bootstrap(store, bad)

    def test_good_ids_accepted(self):
        store = Store(":memory:")
        for good in ["bigpickle", "agent-2", "a.b_c", "x" * 64]:
            private, public = keys.generate()
            secret = identity.seed_bootstrap(store, good)
            identity.enrol(store, good, keys.export_public_b64(public), secret)
            self.assertIn(good, store.agents())

    def test_newline_id_cannot_reach_the_audit_log(self):
        store = Store(":memory:")
        with self.assertRaises(errors.SigilError):
            identity.seed_bootstrap(store, "evil\nfake-event")


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


class TestClaimCodeExpiryCAS(unittest.TestCase):
    """R6: claim_code CAS must include redeem_expires_at predicate."""

    def test_claim_code_after_expiry_returns_false(self):
        store, private = helpers.fresh()
        issued = helpers.issue(store, private, code_ttl=1, view_ttl=60)
        digest = viewcodes._hash_code(issued["code"])
        # Force expiry in the row.
        store.conn.execute(
            "UPDATE codes SET redeem_expires_at = ? WHERE code_hash = ?",
            (int(time.time()) - 1, digest),
        )
        store.conn.commit()
        self.assertFalse(store.claim_code(digest))
        row = store.get_code(digest)
        self.assertIsNone(row["used_at"])


class TestEnrolRekeyRateLimit(unittest.TestCase):
    """N3: enrol_rekey uses a dedicated rekey: bucket; junk cannot lock ops."""

    def test_exhausting_enrol_bucket_does_not_block_rekey(self):
        store, _ = helpers.fresh("limited")
        for _ in range(20):
            store.allow("enrol:limited")
        secret = identity.seed_bootstrap(store, "limited")
        # Valid rekey must succeed even if enrol: is exhausted.
        identity.enrol_rekey(
            store, "limited", keys.export_public_b64(keys.generate()[1]), secret
        )

    def test_junk_rekey_fills_rekey_bucket_not_enrol(self):
        store, _ = helpers.fresh("victim")
        for _ in range(15):
            try:
                identity.enrol_rekey(
                    store, "victim", keys.export_public_b64(keys.generate()[1]), "junk"
                )
            except (errors.EnrolmentFailed, errors.RateLimited):
                pass
        # enrol: bucket still usable
        self.assertTrue(store.allow("enrol:victim"))
        # Valid ops rekey still works: failures charge rekey: only after bad secret,
        # and a valid secret bypasses the limiter.
        secret = identity.seed_bootstrap(store, "victim")
        identity.enrol_rekey(
            store, "victim", keys.export_public_b64(keys.generate()[1]), secret
        )

    def test_failed_rekey_eventually_rate_limited(self):
        store, _ = helpers.fresh("hammer")
        hit = False
        for _ in range(30):
            try:
                identity.enrol_rekey(
                    store, "hammer", keys.export_public_b64(keys.generate()[1]), "nope"
                )
            except errors.RateLimited:
                hit = True
                break
            except errors.EnrolmentFailed:
                pass
        self.assertTrue(hit, "junk rekey attempts should eventually RateLimited")



class TestBootstrapRowDeletedOnSpend(unittest.TestCase):
    """R8: spent bootstrap secrets are deleted, not retained as hashes."""

    def test_spend_deletes_the_row(self):
        store = Store(":memory:")
        secret = identity.seed_bootstrap(store, "dave")
        digest = __import__("hashlib").sha256(secret.encode()).hexdigest()
        self.assertIsNotNone(store.conn.execute(
            "SELECT 1 FROM bootstrap WHERE secret_hash = ?", (digest,)
        ).fetchone())
        owner = store.spend_bootstrap(digest)
        self.assertEqual(owner, "dave")
        self.assertIsNone(store.conn.execute(
            "SELECT 1 FROM bootstrap WHERE secret_hash = ?", (digest,)
        ).fetchone())


class TestResolveRechecksAgentState(unittest.TestCase):
    """N1: resolve must refuse sessions that outran revoke/rekey."""

    def _plant_agent_session(self, store, agent, fingerprint):
        """Insert a live session as if put_session won a race against revoke."""
        import time as _t
        token = "race-survivor-token"
        store.put_session(
            sessions.hash_token(token),
            sessions.AGENT,
            agent,
            agent,
            list(sessions.AGENT_SCOPES),
            int(_t.time()) + 900,
            fingerprint=fingerprint,
        )
        return token

    def test_resolve_refuses_race_survivor_after_permanent_revoke(self):
        store, private = helpers.fresh()
        fp = store.get_agent("bigpickle")["fingerprint"]
        identity.revoke(store, "bigpickle", reason="compromise")
        # Session inserted *after* revoke (race); not marked revoked_at.
        token = self._plant_agent_session(store, "bigpickle", fp)
        with self.assertRaises(errors.SessionExpired):
            sessions.resolve(store, token)

    def test_resolve_refuses_stale_fingerprint_after_rekey(self):
        store, private = helpers.fresh()
        old_fp = store.get_agent("bigpickle")["fingerprint"]
        new_pub = keys.generate()[1]
        secret = identity.seed_bootstrap(store, "bigpickle")
        identity.enrol_rekey(
            store, "bigpickle", keys.export_public_b64(new_pub), secret
        )
        token = self._plant_agent_session(store, "bigpickle", old_fp)
        with self.assertRaises(errors.SessionExpired):
            sessions.resolve(store, token)

    def test_redeem_refuses_code_after_rekey_even_if_not_invalidated(self):
        """Fingerprint binding: plant a code with a stale fingerprint."""
        store, private = helpers.fresh()
        agent_row = store.get_agent("bigpickle")
        old_fp = agent_row["fingerprint"]
        # Issue normally, then surgically restore used_at and keep old fp after rekey
        issued = helpers.issue(store, private)
        digest = viewcodes._hash_code(issued["code"])
        new_pub = keys.generate()[1]
        secret = identity.seed_bootstrap(store, "bigpickle")
        identity.enrol_rekey(
            store, "bigpickle", keys.export_public_b64(new_pub), secret
        )
        # Rekey invalidated the code; un-burn it to simulate put_code after invalidate.
        store.conn.execute(
            "UPDATE codes SET used_at = NULL, fingerprint = ? WHERE code_hash = ?",
            (old_fp, digest),
        )
        store.conn.commit()
        with self.assertRaises(errors.CodeInvalid):
            viewcodes.redeem_code(store, issued["code"])

    def test_redeem_refuses_without_live_key(self):
        store, private = helpers.fresh()
        issued = helpers.issue(store, private)
        digest = viewcodes._hash_code(issued["code"])
        identity.revoke(store, "bigpickle", reason="compromise")
        identity.reinstate(store, "bigpickle")  # ban lifted, key still wiped
        store.conn.execute(
            "UPDATE codes SET used_at = NULL WHERE code_hash = ?", (digest,)
        )
        store.conn.commit()
        with self.assertRaises(errors.CodeInvalid):
            viewcodes.redeem_code(store, issued["code"])

    def test_verify_then_revoke_session_unusable(self):
        """Session minted just before revoke must not resolve."""
        store, private = helpers.fresh()
        chal = identity.challenge(store, "bigpickle")
        sig = helpers.sign(private, chal["nonce"])
        sess = identity.verify(store, chal["nonce"], sig)
        identity.revoke(store, "bigpickle", reason="compromise")
        with self.assertRaises(errors.SessionExpired):
            sessions.resolve(store, sess["token"])


class TestHumanSessionFingerprintEpoch(unittest.TestCase):
    """N7: rekey must kill race-surviving human_via_agent sessions too."""

    def test_resolve_refuses_stale_human_fingerprint_after_rekey(self):
        store, private = helpers.fresh()
        old_fp = store.get_agent("bigpickle")["fingerprint"]
        # Plant a human session as if redeem's put_session won a race against rekey.
        import time as _t
        token = "human-race-survivor"
        store.put_session(
            sessions.hash_token(token),
            sessions.HUMAN_VIA_AGENT,
            "bigpickle's human",
            "bigpickle",
            list(sessions.HUMAN_VIEW_SCOPES),
            int(_t.time()) + 3600,
            granted_by="bigpickle",
            fingerprint=old_fp,
        )
        new_pub = keys.generate()[1]
        secret = identity.seed_bootstrap(store, "bigpickle")
        identity.enrol_rekey(
            store, "bigpickle", keys.export_public_b64(new_pub), secret
        )
        # Clear revoked_at that rekey set, to isolate the fingerprint check.
        store.conn.execute(
            "UPDATE sessions SET revoked_at = NULL, revoked_reason = NULL"
            " WHERE token_hash = ?",
            (sessions.hash_token(token),),
        )
        store.conn.commit()
        with self.assertRaises(errors.SessionExpired) as ctx:
            sessions.resolve(store, token)
        self.assertIn("key epoch", str(ctx.exception))

    def test_issue_code_return_has_single_view_ttl(self):
        """N8: issue_code return dict must not duplicate view_ttl."""
        import inspect
        from sigil import viewcodes as vc
        src = inspect.getsource(vc.issue_code)
        # Count literal "view_ttl" keys in the return dict construction region.
        self.assertEqual(src.count('"view_ttl":'), 1, src[src.rfind("return"):])