Update helper/subscriptions.py

helper/subscriptions.py

@@ -1,11 +1,14 @@
-import os
-import psycopg
-from psycopg.rows import dict_row
-from supabase import create_client, Client
 import asyncio
+import hashlib
 import json
+import os
+from datetime import datetime, timezone
 from pathlib import Path
 
+import psycopg
+from psycopg.rows import dict_row
+from supabase import Client, create_client
+
 ADMINS_PATH = Path("assets/admins.json")
 
 try:
@@ -111,6 +114,39 @@ async def get_conn():
     return conn
 
 
+async def execute_query(query: str, params=(), *, fetchone: bool = False, commit: bool = False):
+    async with conn_lock:
+        connection = await get_conn()
+
+        try:
+            with connection.cursor() as cur:
+                cur.execute(query, params)
+                if cur.description is None:
+                    result = None if fetchone else []
+                else:
+                    result = cur.fetchone() if fetchone else cur.fetchall()
+                if commit:
+                    connection.commit()
+                return result
+        except psycopg.OperationalError:
+            connection = psycopg.connect(
+                POSTGRE_SECRET,
+                row_factory=dict_row,
+                sslmode="verify-full",
+                sslrootcert="prod-ca-2021.crt",
+            )
+            globals()["conn"] = connection
+            with connection.cursor() as cur:
+                cur.execute(query, params)
+                if cur.description is None:
+                    result = None if fetchone else []
+                else:
+                    result = cur.fetchone() if fetchone else cur.fetchall()
+                if commit:
+                    connection.commit()
+                return result
+
+
 def normalize_plan_key(plan_name: str | None) -> str:
     if not plan_name:
         return "free"
@@ -127,24 +163,17 @@ def normalize_plan_key(plan_name: str | None) -> str:
 
 
 async def fetch_subscription(jwt: str):
-    auth_res = supabase.auth.get_user(jwt)
-    if auth_res.user is None:
-        return {"error": "Invalid or expired session"}
-
-    user = auth_res.user
-    email = user.email
-
-    auth_res = supabase.auth.get_user(jwt)
-    if auth_res.user is None:
-        return {"error": "Invalid or expired session"}
-    
-    user = auth_res.user
-    email = user.email.lower()
-    
+    identity = await resolve_token_identity(jwt)
+    if "error" in identity:
+        return identity
+
+    email = identity["email"]
+    signed_up = identity["signed_up"]
+
     if email in ADMIN_EMAIL_MAP:
         admin = ADMIN_EMAIL_MAP[email]
         print(f"[ADMIN OVERRIDE] {email} → forcing plan '{admin['plan_key']}'")
-    
+
         subscription_obj = {
             "subscription_id": admin.get("subscription_id", "admin-override"),
             "status": admin.get("status", "active"),
@@ -154,92 +183,56 @@ async def fetch_subscription(jwt: str):
             "nickname": admin.get("nickname"),
             "plan_key": admin.get("plan_key"),
         }
-    
+
         return {
             "email": email,
-            "signed_up": user.created_at.isoformat(),
+            "signed_up": signed_up,
             "subscription": [subscription_obj],
             "plan_key": admin["plan_key"],
+            "auth_type": identity["auth_type"],
         }
 
-    async with conn_lock:
-        connection = await get_conn()
-
-        try:
-            with connection.cursor() as cur:
-                cur.execute("""
-                    with cust as (
-                        select id 
-                        from stripe.customers 
-                        where email = %s
-                    ),
-                    subs as (
-                        select
-                            s.id as subscription_id,
-                            s.status,
-                            s.current_period_end,
-                            s.items->'data'->0->'price'->>'id' as price_id
-                        from stripe.subscriptions s
-                        join cust on s.customer = cust.id
-                        where s.status in ('active', 'trialing', 'past_due')
-                    )
-                    select
-                        subs.subscription_id,
-                        subs.status,
-                        subs.current_period_end,
-                        subs.price_id,
-                        prices.nickname,
-                        prices.product as product_id,
-                        products.name as product_name
-                    from subs
-                    left join stripe.prices prices
-                      on prices.id = subs.price_id
-                    left join stripe.products products
-                      on prices.product = products.id;
-                """, (email,))
-                rows = cur.fetchall()
-
-        except psycopg.OperationalError:
-            connection = psycopg.connect(POSTGRE_SECRET, row_factory=dict_row)
-            with connection.cursor() as cur:
-                cur.execute("""
-                    with cust as (
-                        select id 
-                        from stripe.customers 
-                        where email = %s
-                    ),
-                    subs as (
-                        select
-                            s.id as subscription_id,
-                            s.status,
-                            s.current_period_end,
-                            s.items->'data'->0->'price'->>'id' as price_id
-                        from stripe.subscriptions s
-                        join cust on s.customer = cust.id
-                        where s.status in ('active', 'trialing', 'past_due')
-                    )
-                    select
-                        subs.subscription_id,
-                        subs.status,
-                        subs.current_period_end,
-                        subs.price_id,
-                        prices.nickname,
-                        prices.product as product_id,
-                        products.name as product_name
-                    from subs
-                    left join stripe.prices prices
-                      on prices.id = subs.price_id
-                    left join stripe.products products
-                      on prices.product = products.id;
-                """, (email,))
-                rows = cur.fetchall()
+    rows = await execute_query(
+        """
+            with cust as (
+                select id
+                from stripe.customers
+                where email = %s
+            ),
+            subs as (
+                select
+                    s.id as subscription_id,
+                    s.status,
+                    s.current_period_end,
+                    s.items->'data'->0->'price'->>'id' as price_id
+                from stripe.subscriptions s
+                join cust on s.customer = cust.id
+                where s.status in ('active', 'trialing', 'past_due')
+            )
+            select
+                subs.subscription_id,
+                subs.status,
+                subs.current_period_end,
+                subs.price_id,
+                prices.nickname,
+                prices.product as product_id,
+                products.name as product_name
+            from subs
+            left join stripe.prices prices
+              on prices.id = subs.price_id
+            left join stripe.products products
+              on prices.product = products.id;
+        """,
+        (email,),
+    )
 
     if not rows:
         return {
             "email": email,
-            "signed_up": user.created_at.isoformat(),
+            "signed_up": signed_up,
             "subscription": None,
             "plan_key": "free",
+            "auth_type": identity["auth_type"],
         }
 
     subscriptions = []
@@ -263,7 +256,95 @@ async def fetch_subscription(jwt: str):
 
     return {
         "email": email,
-        "signed_up": user.created_at.isoformat(),
+        "signed_up": signed_up,
         "subscription": subscriptions,
         "plan_key": preferred_plan_key,
+        "auth_type": identity["auth_type"],
+    }
+
+
+async def resolve_token_identity(token: str):
+    jwt_identity = resolve_jwt_identity(token)
+    if jwt_identity is not None:
+        return jwt_identity
+
+    api_key_identity = await resolve_api_key_identity(token)
+    if api_key_identity is not None:
+        return api_key_identity
+
+    return {"error": "Invalid or expired credentials"}
+
+
+def resolve_jwt_identity(token: str):
+    try:
+        auth_res = supabase.auth.get_user(token)
+    except Exception:
+        return None
+
+    user = getattr(auth_res, "user", None)
+    email = getattr(user, "email", None)
+    created_at = getattr(user, "created_at", None)
+    user_id = getattr(user, "id", None)
+    if user is None or not isinstance(email, str) or not email.strip():
+        return None
+
+    return {
+        "auth_type": "jwt",
+        "user_id": user_id,
+        "email": email.strip().lower(),
+        "signed_up": created_at.isoformat() if hasattr(created_at, "isoformat") else None,
+    }
+
+
+async def resolve_api_key_identity(token: str):
+    token_hash = hashlib.sha256(token.encode("utf-8")).hexdigest()
+    row = await execute_query(
+        """
+            select
+                k.id,
+                k.user_id,
+                k.name,
+                k.expires_at,
+                k.revoked_at,
+                u.email,
+                u.created_at as user_created_at
+            from public.lightning_api_keys k
+            join auth.users u on u.id = k.user_id
+            where k.key_hash = %s
+            limit 1;
+        """,
+        (token_hash,),
+        fetchone=True,
+    )
+    if not row:
+        return None
+
+    if row.get("revoked_at") is not None:
+        return {"error": "API key has been revoked"}
+
+    expires_at = row.get("expires_at")
+    if expires_at is not None and expires_at <= datetime.now(timezone.utc):
+        return {"error": "API key has expired"}
+
+    email = row.get("email")
+    created_at = row.get("user_created_at")
+    user_id = row.get("user_id")
+    if not isinstance(email, str) or not email.strip():
+        return {"error": "API key is not linked to a valid user"}
+
+    await execute_query(
+        """
+            update public.lightning_api_keys
+            set last_used_at = timezone('utc', now())
+            where id = %s;
+        """,
+        (row["id"],),
+        commit=True,
+    )
+
+    return {
+        "auth_type": "api_key",
+        "user_id": str(user_id) if user_id is not None else None,
+        "email": email.strip().lower(),
+        "signed_up": created_at.isoformat() if hasattr(created_at, "isoformat") else None,
     }