Add IP address, full referrer URL, and improved decryption fallback

.gitignore

@@ -45,4 +45,4 @@ build/
 *.egg-info/
 
 # Private keys (keep in repo but be careful)
-# PRIVATE_KEY.pem
+PRIVATE_KEY.pem

app.py

@@ -122,6 +122,7 @@ async def get_data(
                     "id": item.id,
                     "user_id": item.user_id,
                     "refer_url": item.refer_url,
+                    "ip_address": item.ip_address,
                     "json_data": item.json_data,
                     "created_at": item.created_at.isoformat() if item.created_at else None
                 }
@@ -188,8 +189,18 @@ async def blink(
                 # Store with error information
                 decrypted_results = [{"error": str(e), "raw_encrypted": encrypted_data[:100]}]
         
-        # Get referer URL from headers
-        refer_url = request.headers.get("referer") or request.headers.get("origin")
+        # Get referer URL from headers (full URL, not just origin)
+        refer_url = request.headers.get("referer")
+        
+        # Get client IP address
+        # Check X-Forwarded-For header first (for proxies/load balancers)
+        forwarded_for = request.headers.get("x-forwarded-for")
+        if forwarded_for:
+            # X-Forwarded-For can contain multiple IPs, take the first one
+            ip_address = forwarded_for.split(",")[0].strip()
+        else:
+            # Fall back to direct client IP
+            ip_address = request.client.host if request.client else None
         
         # Store each decrypted result as separate records
         records_created = 0
@@ -197,6 +208,7 @@ async def blink(
             blink_record = BlinkData(
                 user_id=user_id,
                 refer_url=refer_url,
+                ip_address=ip_address,
                 json_data=json_data
             )
             db.add(blink_record)
@@ -207,6 +219,7 @@ async def blink(
             blink_record = BlinkData(
                 user_id=user_id,
                 refer_url=refer_url,
+                ip_address=ip_address,
                 json_data={"encrypted_length": len(encrypted_data)}
             )
             db.add(blink_record)

encryption.py

@@ -21,21 +21,36 @@ _private_key = None
 
 def load_private_key():
     """
-    Load the RSA private key from the PEM file.
+    Load the RSA private key from environment variable or PEM file.
     Caches the key for subsequent calls.
     
+    Priority:
+    1. PRIVATE_KEY environment variable (PEM content)
+    2. PRIVATE_KEY_PATH file
+    
     Returns:
-        RSA private key object
-        
-    Raises:
-        FileNotFoundError: If private key file doesn't exist
-        ValueError: If private key is invalid
+        RSA private key object, or None if not available
     """
     global _private_key
     
     if _private_key is not None:
         return _private_key
     
+    # Try loading from environment variable first
+    private_key_pem = os.getenv("PRIVATE_KEY")
+    if private_key_pem:
+        try:
+            _private_key = serialization.load_pem_private_key(
+                private_key_pem.encode(),
+                password=None,
+                backend=default_backend()
+            )
+            logger.info("Successfully loaded private key from PRIVATE_KEY env variable")
+            return _private_key
+        except Exception as e:
+            logger.warning(f"Failed to load private key from env: {e}")
+    
+    # Try loading from file
     try:
         with open(PRIVATE_KEY_PATH, "rb") as key_file:
             _private_key = serialization.load_pem_private_key(
@@ -46,11 +61,12 @@ def load_private_key():
         logger.info(f"Successfully loaded private key from {PRIVATE_KEY_PATH}")
         return _private_key
     except FileNotFoundError:
-        logger.error(f"Private key file not found: {PRIVATE_KEY_PATH}")
-        raise
+        logger.warning(f"Private key file not found: {PRIVATE_KEY_PATH}")
     except Exception as e:
-        logger.error(f"Failed to load private key: {e}")
-        raise ValueError(f"Invalid private key: {e}")
+        logger.warning(f"Failed to load private key from file: {e}")
+    
+    logger.warning("No private key available - encrypted data will not be decrypted")
+    return None
 
 
 def decrypt_data(encrypted_base64: str) -> Optional[Any]:
@@ -61,15 +77,17 @@ def decrypt_data(encrypted_base64: str) -> Optional[Any]:
         encrypted_base64: Base64 URL-safe encoded encrypted string
         
     Returns:
-        Decrypted data parsed as JSON, or raw string if not valid JSON
-        
-    Raises:
-        ValueError: If decryption fails
+        Decrypted data parsed as JSON, encrypted data if no key available, or None on error
     """
     try:
         # Load the private key
         private_key = load_private_key()
         
+        # If no private key, return the encrypted data as-is
+        if private_key is None:
+            logger.warning("No private key - returning encrypted data")
+            return {"encrypted_data": encrypted_base64, "decryption_status": "no_key_available"}
+        
         # Decode base64 URL-safe encoded data
         # Add padding if necessary
         padded = encrypted_base64 + '=' * (4 - len(encrypted_base64) % 4)
@@ -98,7 +116,8 @@ def decrypt_data(encrypted_base64: str) -> Optional[Any]:
             
     except Exception as e:
         logger.error(f"Decryption failed: {e}")
-        raise ValueError(f"Failed to decrypt data: {e}")
+        # Return encrypted data on failure
+        return {"encrypted_data": encrypted_base64, "decryption_error": str(e)}
 
 
 def decrypt_multiple_blocks(encrypted_data: str, block_size: int = 344) -> list[Any]:

models.py

@@ -22,6 +22,7 @@ class BlinkData(Base):
     id = Column(Integer, primary_key=True, autoincrement=True, index=True)
     user_id = Column(String(20), index=True, nullable=False)
     refer_url = Column(Text, nullable=True)
+    ip_address = Column(String(45), nullable=True)  # IPv6 can be up to 45 chars
     json_data = Column(JSON, nullable=True)
     created_at = Column(DateTime(timezone=True), server_default=func.now())
 

templates/index.html

@@ -1,5 +1,6 @@
 <!DOCTYPE html>
 <html lang="en">
+
 <head>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
@@ -186,7 +187,9 @@
         }
 
         @keyframes spin {
-            to { transform: rotate(360deg); }
+            to {
+                transform: rotate(360deg);
+            }
         }
 
         .empty-state {
@@ -203,6 +206,7 @@
         }
     </style>
 </head>
+
 <body>
     <div class="container">
         <header>
@@ -228,13 +232,14 @@
                         <th>ID</th>
                         <th>User ID</th>
                         <th>Refer URL</th>
+                        <th>IP Address</th>
                         <th>JSON Data</th>
                         <th>Created At</th>
                     </tr>
                 </thead>
                 <tbody id="tableBody">
                     <tr>
-                        <td colspan="5">
+                        <td colspan="6">
                             <div class="loading">
                                 <div class="spinner"></div>
                                 Loading data...
@@ -269,11 +274,11 @@
 
         function renderTable(items) {
             const tbody = document.getElementById('tableBody');
-            
+
             if (items.length === 0) {
                 tbody.innerHTML = `
                     <tr>
-                        <td colspan="5">
+                        <td colspan="6">
                             <div class="empty-state">
                                 <svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" stroke="currentColor">
                                     <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M20 13V6a2 2 0 00-2-2H6a2 2 0 00-2 2v7m16 0v5a2 2 0 01-2 2H6a2 2 0 01-2-2v-5m16 0h-2.586a1 1 0 00-.707.293l-2.414 2.414a1 1 0 01-.707.293h-3.172a1 1 0 01-.707-.293l-2.414-2.414A1 1 0 006.586 13H4" />
@@ -291,6 +296,7 @@
                     <td>${item.id}</td>
                     <td class="user-id">${item.user_id}</td>
                     <td class="refer-url" title="${item.refer_url || ''}">${item.refer_url || '-'}</td>
+                    <td class="ip-address">${item.ip_address || '-'}</td>
                     <td><pre class="json-data">${JSON.stringify(item.json_data, null, 2)}</pre></td>
                     <td class="timestamp">${new Date(item.created_at).toLocaleString()}</td>
                 </tr>
@@ -309,10 +315,10 @@
             currentPage = page;
             const data = await fetchData(page);
             totalRecords = data.total;
-            
+
             document.getElementById('totalRecords').textContent = data.total.toLocaleString();
             document.getElementById('uniqueUsers').textContent = data.unique_users.toLocaleString();
-            
+
             renderTable(data.items);
             updatePagination();
         }
@@ -334,4 +340,5 @@
         loadPage(1);
     </script>
 </body>
-</html>
+
+</html>