Add IP-based daily free search quota and Anti-DDoS rate limits

auth.py

@@ -8,18 +8,34 @@ from typing import Optional, Dict, Any, List, Tuple
 KEYS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "api_keys.json")
 ADMIN_MASTER_KEY = os.environ.get("ADMIN_MASTER_KEY", "secretz_master_admin_2026")
 INTERNAL_BOT_KEY = os.environ.get("INTERNAL_BOT_KEY", "sk_internal_telebot_secretz_9921")
-REQUIRE_AUTH = os.environ.get("REQUIRE_AUTH", "false").lower() in ("true", "1", "yes")
+DAILY_FREE_LIMIT = int(os.environ.get("DAILY_FREE_LIMIT", "5"))
 
 _lock = threading.Lock()
-# In-memory sliding window for rate limiting: {api_key: [timestamp1, timestamp2, ...]}
+
+# Rate limiting sliding windows
+# {api_key: [timestamp1, timestamp2, ...]}
 _rate_limits: Dict[str, List[float]] = {}
 
+# Public IP Anti-DDoS sliding windows
+# {ip: [timestamp1, timestamp2, ...]}
+_public_ip_burst: Dict[str, List[float]] = {}
+
+# Public IP Daily usage: {ip: {"date": "YYYY-MM-DD", "count": int}}
+_public_ip_daily: Dict[str, Dict[str, Any]] = {}
+
 
 def load_keys() -> Dict[str, Dict[str, Any]]:
     if os.path.exists(KEYS_FILE):
         try:
             with open(KEYS_FILE, "r", encoding="utf-8") as f:
-                return json.load(f)
+                data = json.load(f)
+                # Ensure all keys have both 'key' and 'api_key' fields
+                for k, v in data.items():
+                    if "key" not in v:
+                        v["key"] = k
+                    if "api_key" not in v:
+                        v["api_key"] = k
+                return data
         except Exception:
             return {}
     return {}
@@ -47,12 +63,13 @@ class KeyManager:
         expires_at = now + (days_valid * 86400) if (days_valid and days_valid > 0) else None
 
         key_info = {
+            "key": api_key,
             "api_key": api_key,
-            "name": name.strip() or "Unnamed Client",
+            "name": name,
             "created_at": now,
             "expires_at": expires_at,
-            "rate_limit_per_min": max(1, int(rate_limit_per_min)),
-            "total_credits": int(total_credits),  # -1 for unlimited
+            "rate_limit_per_min": rate_limit_per_min,
+            "total_credits": total_credits,
             "used_credits": 0,
             "status": "active"
         }
@@ -68,7 +85,13 @@ class KeyManager:
     def list_keys() -> List[Dict[str, Any]]:
         with _lock:
             keys = load_keys()
-            return list(keys.values())
+            result = []
+            for k, v in keys.items():
+                item = dict(v)
+                item["key"] = k
+                item["api_key"] = k
+                result.append(item)
+            return result
 
     @staticmethod
     def revoke_key(api_key: str) -> bool:
@@ -91,66 +114,121 @@ class KeyManager:
             return False
 
     @staticmethod
-    def validate_and_consume(api_key: Optional[str]) -> Tuple[bool, int, str]:
+    def validate_request(api_key: Optional[str], client_ip: str) -> Tuple[bool, int, str, Dict[str, Any]]:
         """
-        Validates API key, checks expiration, rate limit, and consumes 1 credit.
-        Returns: (is_valid: bool, error_code: int, error_message: str)
+        Validates an incoming request either by API Key or Public Free Tier (IP-based).
+        Returns: (is_valid: bool, error_code: int, error_message: str, meta: dict)
         """
-        if not REQUIRE_AUTH:
-            return True, 200, "OK"
+        now = time.time()
 
-        if not api_key:
-            return False, 401, "API Key required. Pass header 'X-API-Key' or query parameter 'api_key'."
+        # ========================================================
+        # Case 1: Master Admin or Internal Bot Key
+        # ========================================================
+        if api_key in (ADMIN_MASTER_KEY, INTERNAL_BOT_KEY):
+            return True, 200, "OK", {"tier": "admin", "remaining": -1}
+
+        # ========================================================
+        # Case 2: User provided a custom API Key
+        # ========================================================
+        if api_key:
+            with _lock:
+                keys = load_keys()
+                if api_key not in keys:
+                    return False, 401, "Invalid API Key. Please verify your credentials.", {}
+
+                info = keys[api_key]
+
+                # Status check
+                if info.get("status") != "active":
+                    return False, 403, "API Key has been revoked or deactivated.", {}
+
+                # Expiry check
+                expires_at = info.get("expires_at")
+                if expires_at and now > expires_at:
+                    info["status"] = "expired"
+                    save_keys(keys)
+                    return False, 403, "API Key has expired.", {}
+
+                # Credits check
+                total_credits = info.get("total_credits", -1)
+                used_credits = info.get("used_credits", 0)
+                if total_credits != -1 and used_credits >= total_credits:
+                    return False, 403, "API Key search credits exhausted.", {}
+
+                # Rate limit (sliding window 60s)
+                rate_limit = info.get("rate_limit_per_min", 60)
+                if api_key not in _rate_limits:
+                    _rate_limits[api_key] = []
+
+                cutoff = now - 60.0
+                _rate_limits[api_key] = [t for t in _rate_limits[api_key] if t > cutoff]
+
+                if len(_rate_limits[api_key]) >= rate_limit:
+                    retry_after = round(60.0 - (now - _rate_limits[api_key][0]), 1)
+                    return False, 429, f"Rate limit exceeded ({rate_limit} req/min). Retry after {retry_after}s.", {}
+
+                # Deduct credit
+                _rate_limits[api_key].append(now)
+                info["used_credits"] = used_credits + 1
+                save_keys(keys)
 
-        # Master admin key bypass
-        if api_key == ADMIN_MASTER_KEY or api_key == INTERNAL_BOT_KEY:
-            return True, 200, "OK"
+                remaining = (total_credits - info["used_credits"]) if total_credits != -1 else -1
+                return True, 200, "OK", {"tier": "custom", "remaining": remaining}
 
+        # ========================================================
+        # Case 3: Public / Anonymous User (Bina Key ke - Daily Free Limit & Anti-DDoS)
+        # ========================================================
         with _lock:
-            keys = load_keys()
-            if api_key not in keys:
-                return False, 401, "Invalid API Key."
-
-            info = keys[api_key]
+            # 1. Anti-DDoS / Rapid Burst Check (Max 1 request per 2 seconds, max 6 per minute)
+            if client_ip not in _public_ip_burst:
+                _public_ip_burst[client_ip] = []
 
-            # 1. Status check
-            if info.get("status") != "active":
-                return False, 403, "API Key has been revoked or deactivated."
+            cutoff = now - 60.0
+            _public_ip_burst[client_ip] = [t for t in _public_ip_burst[client_ip] if t > cutoff]
 
-            now = time.time()
+            if _public_ip_burst[client_ip] and (now - _public_ip_burst[client_ip][-1] < 1.5):
+                return False, 429, "Anti-DDoS: Rapid requests detected. Please wait 2 seconds between queries.", {}
 
-            # 2. Expiry check
-            expires_at = info.get("expires_at")
-            if expires_at and now > expires_at:
-                info["status"] = "expired"
-                save_keys(keys)
-                return False, 403, "API Key has expired."
+            if len(_public_ip_burst[client_ip]) >= 6:
+                return False, 429, "Public rate limit exceeded (max 6 req/min). Please slow down.", {}
 
-            # 3. Credit check
-            total_credits = info.get("total_credits", -1)
-            used_credits = info.get("used_credits", 0)
-            if total_credits != -1 and used_credits >= total_credits:
-                return False, 403, "API Key credit quota exhausted."
+            # 2. Daily Free Search Limit (Reset at 00:00 UTC)
+            today = time.strftime("%Y-%m-%d", time.gmtime())
+            if client_ip not in _public_ip_daily or _public_ip_daily[client_ip].get("date") != today:
+                _public_ip_daily[client_ip] = {"date": today, "count": 0}
 
-            # 4. Rate limit check (sliding window 60s)
-            rate_limit = info.get("rate_limit_per_min", 60)
-            if api_key not in _rate_limits:
-                _rate_limits[api_key] = []
+            used_today = _public_ip_daily[client_ip]["count"]
+            if used_today >= DAILY_FREE_LIMIT:
+                return (
+                    False,
+                    429,
+                    f"Daily free limit reached ({DAILY_FREE_LIMIT}/{DAILY_FREE_LIMIT} used today). Resets at midnight UTC, or enter an Access Key for unlimited searches.",
+                    {"tier": "free", "remaining": 0}
+                )
 
-            # Keep only timestamps in the last 60 seconds
-            cutoff = now - 60.0
-            _rate_limits[api_key] = [t for t in _rate_limits[api_key] if t > cutoff]
+            # Record usage
+            _public_ip_burst[client_ip].append(now)
+            _public_ip_daily[client_ip]["count"] = used_today + 1
+            remaining = max(0, DAILY_FREE_LIMIT - _public_ip_daily[client_ip]["count"])
 
-            if len(_rate_limits[api_key]) >= rate_limit:
-                retry_after = round(60.0 - (now - _rate_limits[api_key][0]), 1)
-                return False, 429, f"Rate limit exceeded ({rate_limit} req/min). Retry after {retry_after}s."
+            return True, 200, "OK", {"tier": "free", "remaining": remaining}
 
-            # Record request & consume credit
-            _rate_limits[api_key].append(now)
-            info["used_credits"] = used_credits + 1
-            save_keys(keys)
-
-            return True, 200, "OK"
+    @staticmethod
+    def get_public_quota(client_ip: str) -> Dict[str, Any]:
+        """Returns current daily quota status for an IP without consuming credit."""
+        today = time.strftime("%Y-%m-%d", time.gmtime())
+        with _lock:
+            record = _public_ip_daily.get(client_ip, {})
+            if record.get("date") != today:
+                used = 0
+            else:
+                used = record.get("count", 0)
+            return {
+                "tier": "free",
+                "daily_limit": DAILY_FREE_LIMIT,
+                "used_today": used,
+                "remaining_today": max(0, DAILY_FREE_LIMIT - used)
+            }
 
     @staticmethod
     def get_key_status(api_key: str) -> Optional[Dict[str, Any]]:
@@ -161,4 +239,10 @@ class KeyManager:
 
         with _lock:
             keys = load_keys()
-            return keys.get(api_key)
+            info = keys.get(api_key)
+            if info:
+                res = dict(info)
+                res["key"] = api_key
+                res["api_key"] = api_key
+                return res
+            return None

main.py

@@ -71,13 +71,27 @@ class CreateKeyRequest(BaseModel):
 # ---------------------------------------------------------
 # Authentication Dependency
 # ---------------------------------------------------------
+def get_client_ip(request: Request) -> str:
+    """Extracts client IP considering Cloudflare, proxies, and standard headers."""
+    cf_ip = request.headers.get("cf-connecting-ip")
+    if cf_ip:
+        return cf_ip.strip()
+    forwarded = request.headers.get("x-forwarded-for")
+    if forwarded:
+        return forwarded.split(",")[0].strip()
+    real_ip = request.headers.get("x-real-ip")
+    if real_ip:
+        return real_ip.strip()
+    return request.client.host if request.client else "127.0.0.1"
+
+
 def verify_api_key(
     request: Request,
     api_key_query: Optional[str] = Query(None, alias="api_key"),
     x_api_key: Optional[str] = Header(None, alias="X-API-Key"),
     authorization: Optional[str] = Header(None)
-) -> str:
-    """Extracts and verifies API key from Header or Query Param."""
+) -> Dict[str, Any]:
+    """Extracts and verifies API key or applies IP-based daily free search with anti-DDoS."""
     token = None
     if x_api_key:
         token = x_api_key
@@ -86,11 +100,12 @@ def verify_api_key(
     elif api_key_query:
         token = api_key_query
 
-    is_valid, code, msg = KeyManager.validate_and_consume(token)
+    client_ip = get_client_ip(request)
+    is_valid, code, msg, meta = KeyManager.validate_request(token, client_ip)
     if not is_valid:
         raise HTTPException(status_code=code, detail=msg)
 
-    return token or "unauthenticated"
+    return meta
 
 
 def verify_admin_key(
@@ -121,6 +136,13 @@ async def health_check():
     return {"status": "ok", "timestamp": time.time()}
 
 
+@app.get("/api/public/quota")
+async def check_public_quota(request: Request):
+    """Returns current daily free search quota for the caller IP."""
+    client_ip = get_client_ip(request)
+    return KeyManager.get_public_quota(client_ip)
+
+
 @app.get("/api/stats")
 async def get_system_stats():
     """Returns engine health, cache status, and database metrics."""
@@ -143,34 +165,59 @@ async def get_system_stats():
 
 
 # ---------------------------------------------------------
-# Protected Search Endpoints (Requires Valid API Key)
+# Search Endpoints (Public Free Tier or Key Authorized)
 # ---------------------------------------------------------
 @app.get("/api/lookup")
 async def unified_lookup(
     q: Optional[str] = Query(None, description="Search query (phone or aadhar)"),
     phone: Optional[str] = Query(None, description="Direct 10-digit phone number"),
     aadhar: Optional[str] = Query(None, description="Direct 12-digit aadhar number"),
-    _key: str = Depends(verify_api_key)
+    auth_meta: Dict[str, Any] = Depends(verify_api_key)
 ):
+    target = phone or aadhar or q
+    if not target:
+        raise HTTPException(status_code=400, detail="Missing parameter: provide 'q', 'phone', or 'aadhar'.")
+
+    # Anti-abuse sanitization
+    clean = "".join(c for c in str(target) if c.isdigit())
+    if len(clean) not in (10, 12):
+        raise HTTPException(status_code=400, detail="Invalid query. Please provide a 10-digit phone or 12-digit Aadhaar number.")
+
     eng = get_engine()
-    if phone:
-        return eng.query_by_phone(phone)
-    elif aadhar:
-        return eng.query_by_aadhar(aadhar)
-    elif q:
-        return eng.auto_lookup(q)
+    if len(clean) == 12:
+        res = eng.query_by_aadhar(clean)
     else:
-        raise HTTPException(status_code=400, detail="Missing parameter: provide 'q', 'phone', or 'aadhar'.")
+        res = eng.query_by_phone(clean)
+
+    if isinstance(res, dict):
+        res["tier"] = auth_meta.get("tier", "free")
+        res["remaining_credits"] = auth_meta.get("remaining", -1)
+
+    return res
 
 
 @app.get("/api/phone/{phone_number}")
-async def lookup_by_phone(phone_number: str, _key: str = Depends(verify_api_key)):
-    return get_engine().query_by_phone(phone_number)
+async def lookup_by_phone(phone_number: str, auth_meta: Dict[str, Any] = Depends(verify_api_key)):
+    clean = "".join(c for c in phone_number if c.isdigit())
+    if len(clean) != 10:
+        raise HTTPException(status_code=400, detail="Phone number must be exactly 10 digits.")
+    res = get_engine().query_by_phone(clean)
+    if isinstance(res, dict):
+        res["tier"] = auth_meta.get("tier", "free")
+        res["remaining_credits"] = auth_meta.get("remaining", -1)
+    return res
 
 
 @app.get("/api/aadhar/{aadhar_number}")
-async def lookup_by_aadhar(aadhar_number: str, _key: str = Depends(verify_api_key)):
-    return get_engine().query_by_aadhar(aadhar_number)
+async def lookup_by_aadhar(aadhar_number: str, auth_meta: Dict[str, Any] = Depends(verify_api_key)):
+    clean = "".join(c for c in aadhar_number if c.isdigit())
+    if len(clean) != 12:
+        raise HTTPException(status_code=400, detail="Aadhaar number must be exactly 12 digits.")
+    res = get_engine().query_by_aadhar(clean)
+    if isinstance(res, dict):
+        res["tier"] = auth_meta.get("tier", "free")
+        res["remaining_credits"] = auth_meta.get("remaining", -1)
+    return res
 
 
 @app.get("/api/key/status")