#!/usr/bin/env python3
"""
vulnerable_handler.py — Simulates AshAuthentication's vulnerable OAuth2/OIDC
user resolution logic (CVE-2026-49757).

This module reproduces the email-based upsert that AshAuthentication used
before the fix in versions 4.14.0 and 5.0.0-rc.10.

The vulnerable code path:
  1. IdentityChange.change/3 receives OAuth user_info from provider
  2. Upsert with upsert_identity: :unique_email matches on email
  3. SignInPreparation.prepare/3 creates session for matched user

Usage:
    from vulnerable_handler import VulnerableAuthHandler
    handler = VulnerableAuthHandler(db_path=":memory:")
"""

import sqlite3
import json
import hashlib
import secrets
from datetime import datetime, timezone


class VulnerableAuthHandler:
    """
    Simulates AshAuthentication's vulnerable OAuth2/OIDC callback handler.

    The vulnerability: user resolution is performed by email address
    (upsert_identity: :unique_email) instead of the (strategy, sub) pair.
    """

    def __init__(self, db_path=":memory:"):
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self._init_db()

    def _init_db(self):
        self.conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                email       TEXT UNIQUE NOT NULL,
                username    TEXT,
                role        TEXT DEFAULT 'user',
                created_at  TEXT DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS user_identities (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                strategy        TEXT NOT NULL,
                uid             TEXT NOT NULL,
                user_id         INTEGER NOT NULL REFERENCES users(id),
                user_info       TEXT,
                access_token    TEXT,
                refresh_token   TEXT,
                created_at      TEXT DEFAULT (datetime('now')),
                updated_at      TEXT DEFAULT (datetime('now'))
            );

            -- This is the vulnerable unique constraint: on email, not on (strategy, uid)
            CREATE UNIQUE INDEX IF NOT EXISTS idx_identities_email
                ON user_identities(
                    json_extract(user_info, '$.email')
                );
        """)
        self.conn.commit()

    def register_local_user(self, email, username=None, role="user"):
        cur = self.conn.execute(
            "INSERT INTO users (email, username, role) VALUES (?, ?, ?)",
            (email, username or email.split("@")[0], role),
        )
        self.conn.commit()
        return cur.lastrowid

    def oauth_callback(self, strategy_name, user_info, access_token=None,
                       refresh_token=None):
        """
        Simulates IdentityChange.change/3 + SignInPreparation.prepare/3.

        This is the VULNERABLE path — matches user by email.
        """
        email = user_info.get("email")
        sub = user_info.get("sub")
        email_verified = user_info.get("email_verified", False)

        print(f"  [handler] OAuth callback received:")
        print(f"  [handler]   strategy:       {strategy_name}")
        print(f"  [handler]   email:          {email}")
        print(f"  [handler]   sub:            {sub}")
        print(f"  [handler]   email_verified: {email_verified}")

        # Step 1: Upsert on :unique_email (VULNERABLE)
        existing = self.conn.execute(
            """SELECT ui.*, u.email as local_email, u.role
               FROM user_identities ui
               JOIN users u ON ui.user_id = u.id
               WHERE json_extract(ui.user_info, '$.email') = ?""",
            (email,),
        ).fetchone()

        if existing:
            user_id = existing["user_id"]
            print(f"  [handler] MATCH FOUND via email upsert:")
            print(f"  [handler]   user_id:    {user_id}")
            print(f"  [handler]   local_email: {existing['local_email']}")
            print(f"  [handler]   role:       {existing['role']}")

            self.conn.execute(
                """UPDATE user_identities
                   SET uid = ?, user_info = ?, access_token = ?,
                       refresh_token = ?, updated_at = datetime('now')
                   WHERE id = ?""",
                (sub, json.dumps(user_info), access_token,
                 refresh_token, existing["id"]),
            )
            self.conn.commit()
        else:
            # Also check if a local user already has this email
            local_user = self.conn.execute(
                "SELECT id FROM users WHERE email = ?", (email,),
            ).fetchone()

            if local_user:
                user_id = local_user["id"]
                print(f"  [handler] No identity match, but found local user_id={user_id}")
                print(f"  [handler] Linking OAuth identity to existing user")
            else:
                cur = self.conn.execute(
                    """INSERT INTO users (email, username)
                       VALUES (?, ?)""",
                    (email, email.split("@")[0]),
                )
                user_id = cur.lastrowid
                print(f"  [handler] No existing match — created new user_id={user_id}")

            self.conn.execute(
                """INSERT INTO user_identities
                   (strategy, uid, user_id, user_info, access_token, refresh_token)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (strategy_name, sub, user_id, json.dumps(user_info),
                 access_token, refresh_token),
            )
            self.conn.commit()

        return self._create_session(user_id)

    def _create_session(self, user_id):
        user = self.conn.execute(
            "SELECT * FROM users WHERE id = ?", (user_id,),
        ).fetchone()
        token = secrets.token_hex(16)
        print(f"  [handler] SignInPreparation: session created for user_id={user_id}")
        return {
            "session_token": token,
            "user_id": user_id,
            "email": user["email"],
            "username": user["username"],
            "role": user["role"],
        }

    def get_user(self, user_id):
        return self.conn.execute(
            "SELECT * FROM users WHERE id = ?", (user_id,),
        ).fetchone()

    def list_identities(self, user_id):
        return self.conn.execute(
            "SELECT * FROM user_identities WHERE user_id = ?", (user_id,),
        ).fetchall()

    def close(self):
        self.conn.close()


class FixedAuthHandler:
    """
    Simulates the PATCHED AshAuthentication handler (v4.14.0 / v5.0.0-rc.10).

    Uses (strategy, sub) as the primary lookup key.
    Applies on_untrusted_email_match policy for unknown subs.
    """

    def __init__(self, db_path=":memory:", on_untrusted_email_match="reject",
                 trust_email_verified=False):
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self.on_untrusted_email_match = on_untrusted_email_match
        self.trust_email_verified = trust_email_verified
        self._init_db()

    def _init_db(self):
        self.conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                email       TEXT UNIQUE NOT NULL,
                username    TEXT,
                role        TEXT DEFAULT 'user',
                created_at  TEXT DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS user_identities (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                strategy        TEXT NOT NULL,
                uid             TEXT NOT NULL,
                user_id         INTEGER NOT NULL REFERENCES users(id),
                user_info       TEXT,
                access_token    TEXT,
                refresh_token   TEXT,
                created_at      TEXT DEFAULT (datetime('now')),
                updated_at      TEXT DEFAULT (datetime('now'))
            );

            -- Fixed: unique constraint is on (strategy, uid), NOT email
            CREATE UNIQUE INDEX IF NOT EXISTS idx_identities_strategy_uid
                ON user_identities(strategy, uid);
        """)
        self.conn.commit()

    def register_local_user(self, email, username=None, role="user"):
        cur = self.conn.execute(
            "INSERT INTO users (email, username, role) VALUES (?, ?, ?)",
            (email, username or email.split("@")[0], role),
        )
        self.conn.commit()
        return cur.lastrowid

    def link_identity(self, strategy_name, sub, user_id, user_info,
                      access_token=None, refresh_token=None):
        self.conn.execute(
            """INSERT INTO user_identities
               (strategy, uid, user_id, user_info, access_token, refresh_token)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (strategy_name, sub, user_id, json.dumps(user_info),
             access_token, refresh_token),
        )
        self.conn.commit()

    def oauth_callback(self, strategy_name, user_info, access_token=None,
                       refresh_token=None):
        """
        Simulates the FIXED UserResolver.resolve/3.

        1. Look up identity by (strategy, sub)
        2. If found → return linked user
        3. If not found → apply on_untrusted_email_match policy
        """
        email = user_info.get("email")
        sub = user_info.get("sub")
        email_verified = user_info.get("email_verified", False)

        print(f"  [handler] OAuth callback received:")
        print(f"  [handler]   strategy:       {strategy_name}")
        print(f"  [handler]   email:          {email}")
        print(f"  [handler]   sub:            {sub}")
        print(f"  [handler]   email_verified: {email_verified}")

        # Step 1: Resolve by (strategy, sub) — THE FIX
        identity = self.conn.execute(
            """SELECT ui.*, u.email as local_email, u.role
               FROM user_identities ui
               JOIN users u ON ui.user_id = u.id
               WHERE ui.strategy = ? AND ui.uid = ?""",
            (strategy_name, sub),
        ).fetchone()

        if identity:
            user_id = identity["user_id"]
            print(f"  [handler] Identity found via (strategy, sub):")
            print(f"  [handler]   user_id: {user_id}")
            print(f"  [handler]   email:   {identity['local_email']}")

            self.conn.execute(
                """UPDATE user_identities
                   SET access_token = ?, refresh_token = ?, updated_at = datetime('now')
                   WHERE id = ?""",
                (access_token, refresh_token, identity["id"]),
            )
            self.conn.commit()
            return self._create_session(user_id)

        # Step 2: Unknown sub — apply on_untrusted_email_match policy
        print(f"  [handler] No identity found for (strategy={strategy_name}, sub={sub})")

        if self.trust_email_verified and email_verified:
            print(f"  [handler] trust_email_verified=true and provider verified email")
            print(f"  [handler] Auto-linking by verified email...")
            return self._link_and_sign_in(strategy_name, sub, email, user_info,
                                          access_token, refresh_token)

        if self.on_untrusted_email_match == "reject":
            print(f"  [handler] on_untrusted_email_match = :reject")
            print(f"  [handler] LOGIN DENIED — no auto-linking by untrusted email")
            return {"error": "Authentication failed: untrusted email match"}

        elif self.on_untrusted_email_match == "confirm":
            print(f"  [handler] on_untrusted_email_match = :confirm")
            print(f"  [handler] Sending confirmation token to {email}...")
            token = secrets.token_hex(8)
            print(f"  [handler] Confirmation token: {token}")
            print(f"  [handler] Identity will be linked ONLY after email confirmation")
            return {
                "error": "Confirmation required",
                "confirmation_token": token,
                "email": email,
            }

        elif self.on_untrusted_email_match == "warn":
            print(f"  [handler] on_untrusted_email_match = :warn")
            print(f"  [handler] WARNING: Untrusted email match denied")
            return {"error": "Authentication failed: untrusted email match (warning logged)"}

        return {"error": "Unknown policy"}

    def _link_and_sign_in(self, strategy_name, sub, email, user_info,
                          access_token, refresh_token):
        existing_user = self.conn.execute(
            "SELECT * FROM users WHERE email = ?", (email,),
        ).fetchone()

        if existing_user:
            user_id = existing_user["id"]
        else:
            cur = self.conn.execute(
                "INSERT INTO users (email, username) VALUES (?, ?)",
                (email, email.split("@")[0]),
            )
            user_id = cur.lastrowid

        self.link_identity(strategy_name, sub, user_id, user_info,
                          access_token, refresh_token)
        return self._create_session(user_id)

    def _create_session(self, user_id):
        user = self.conn.execute(
            "SELECT * FROM users WHERE id = ?", (user_id,),
        ).fetchone()
        token = secrets.token_hex(16)
        print(f"  [handler] Session created for user_id={user_id}")
        return {
            "session_token": token,
            "user_id": user_id,
            "email": user["email"],
            "username": user["username"],
            "role": user["role"],
        }

    def get_user(self, user_id):
        return self.conn.execute(
            "SELECT * FROM users WHERE id = ?", (user_id,),
        ).fetchone()

    def close(self):
        self.conn.close()


if __name__ == "__main__":
    print("=" * 70)
    print("CVE-2026-49757 — Vulnerable Handler Module")
    print("This module simulates AshAuthentication's email-based user matching.")
    print("Use exploit.py to demonstrate the attack.")
    print("=" * 70)